def _replace_submodules(gm: GraphModule, replacement: torch.nn.Module) -> None: gm.delete_all_unused_submodules() if isinstance(replacement, GraphModule): replacement.graph.lint() def try_get_submodule_or_attr(mod: torch.nn.Module, target: str) -> Optional[torch.nn.Module]: try: mod_match = mod.get_submodule(target) return mod_match except AttributeError: pass # supports getattr as well try: attr = getattr(mod, target) return attr except AttributeError: return None for node in gm.graph.nodes: if node.op == "call_module" or node.op == "get_attr": gm_submod = try_get_submodule_or_attr(gm, node.target) replacement_submod = try_get_submodule_or_attr( replacement, node.target) # CASE 1: This target already exists as a submodule in our # result GraphModule. Whether or not it exists in # `replacement`, the existing submodule takes precedence. if gm_submod is not None: continue # CASE 2: The target exists as a submodule in `replacement` # only, so we need to copy it over. elif replacement_submod is not None: new_submod = copy.deepcopy(getattr(replacement, node.target)) if isinstance(new_submod, torch.nn.Module): gm.add_submodule(node.target, new_submod) else: setattr(gm, node.target, new_submod) # CASE 3: The target doesn't exist as a submodule in `gm` # or `replacement` else: continue raise RuntimeError( "Attempted to create a \"", node.op, "\" node during subgraph rewriting " f"with target {node.target}, but " "the referenced submodule does not " "exist in either the original " "GraphModule `gm` or the replacement" " GraphModule `replacement`") gm.graph.lint()
def get_size_of_node(fx_module: GraphModule, node: Node) -> size_bytes: """Given a node with node.dtype and node.shape, return its total size and its output size. total_size = weights + bias + output_size """ # Total num of elements total_num_of_elems = 0 # For a module, conside all parameters if node.op == "call_module": submodule_dict = dict(fx_module.named_modules()) submodule = submodule_dict[node.target] parameters = submodule.named_parameters() # Parameters are named tuples for name, p in parameters: total_num_of_elems += p.numel() # Don't forget the output size # node.shape is the shape of this node's output tensor_meta = get_tensor_meta(node) output_elem = tensor_meta.shape.numel() total_num_of_elems += output_elem # Assume for now if it's quantized then it's qint8 or quint8 if tensor_meta.is_quantized: size_per_elem_bytes = torch._empty_affine_quantized( [], dtype=tensor_meta.dtype).element_size() else: size_per_elem_bytes = torch.tensor( [], dtype=tensor_meta.dtype).element_size() total_size = size_per_elem_bytes * total_num_of_elems output_size = size_per_elem_bytes * output_elem return size_bytes(output_size, total_size)
def get_size_of_node(fx_module: GraphModule, node: Node) -> size_bytes: """Given a node with node.dtype and node.shape, return its total size and its output size. total_size = weights + bias + output_size """ # Total num of elements total_num_of_elems = 0 # For a module, conside all parameters if node.op == "call_module": submodule_dict = dict(fx_module.named_modules()) submodule = submodule_dict[node.target] parameters = submodule.named_parameters() # Parameters are named tuples for name, p in parameters: total_num_of_elems += p.numel() # Don't forget the output size # node.shape is the shape of this node's output shape = getattr(node, "shape", None) if shape: output_elem = shape.numel() else: raise RuntimeError("Node has no shape attr") total_num_of_elems += output_elem size_per_elem_bytes = 0 dtype = getattr(node, "dtype", None) if dtype: size_per_elem_bytes = torch.tensor([], dtype=dtype).element_size() else: raise RuntimeError("Node has no dtype attr") total_size = size_per_elem_bytes * total_num_of_elems output_size = size_per_elem_bytes * output_elem return size_bytes(output_size, total_size)
def legalize_graph(gm: GraphModule): """ Replace the graph of the given GraphModule with one that contains the same nodes as the original, but in topologically sorted order. This is used by the merge_matmul transformation below, which disturbs the topologically sorted order of its input GraphModule, so that this order is restored before further transformation. Arguments: gm: The graph module to topologically sort. It is modified in-place. """ # Build an adjacency list representation of node dependencies in the graph. This also # serves as a list of nodes that still need to be inserted into the new, topologically # sorted graph. dependencies = { node: node.all_input_nodes.copy() for node in gm.graph.nodes } # Construct a new graph that will contain all nodes in topologically sorted order. new_graph = Graph() value_remap: Dict[Node, Node] = {} # Copy over all nodes with no dependencies. for node, deps in dependencies.items(): if not deps: value_remap[node] = new_graph.node_copy(node, lambda n: value_remap[n]) # Remove the copied over nodes from the adjacency list. for copied_node in value_remap.keys(): del dependencies[copied_node] # While there are still nodes to insert into the new graph: while dependencies: copied_this_round = [] # Copy over all nodes whose dependencies already exist in the new graph. for node, deps in dependencies.items(): all_deps_copied = True for dep in deps: if dep not in value_remap: all_deps_copied = False if all_deps_copied: value_remap[node] = new_graph.node_copy( node, lambda n: value_remap[n]) copied_this_round.append(node) # Delete all nodes copied over in this iteration from dependencies. for copied_node in copied_this_round: del dependencies[copied_node] # Replace the old graph with the new, topologically sorted one. gm.graph = new_graph
def lift_lowering_attrs_to_nodes(fx_module: GraphModule) -> None: """Recursively traverse all `fx_module` nodes and fetch the module's attributes if the node is a leaf module. """ submodules = dict(fx_module.named_modules()) for node in fx_module.graph.nodes: if node.op == "call_module": if isinstance(submodules[node.target], GraphModule): lift_lowering_attrs_to_nodes(submodules[node.target]) else: node.attrs_for_lowering = extract_attrs_for_lowering( submodules[node.target])
def insert_subgm(gm: GraphModule, sub_gm: GraphModule, orig_inputs: Tuple[Node, ...], orig_outputs: Tuple[Node, ...]): # add sub_gm into gm submodule_name = sub_gm.__class__.__name__ gm.add_submodule(submodule_name, sub_gm) # Create a call_module node in main graph. module_node = gm.graph.call_module(submodule_name, args=orig_inputs, kwargs=None) if len(orig_outputs) == 1: # main_remapping[comp.orig_outputs[0]] = module_node orig_outputs[0].replace_all_uses_with(module_node) else: for i, orig_output in enumerate(orig_outputs): # Use Proxy to record getitem access. proxy_out = torch.fx.Proxy( module_node)[i].node # type: ignore[index] orig_output.replace_all_uses_with(proxy_out) return gm
def lift_subgraph_as_module(gm: GraphModule, subgraph: Graph, class_name: str = 'GraphModule') -> GraphModule: """ Create a GraphModule for subgraph, which copies the necessory attributes from the original parent graph_module. Args: gm (GraphModule): parent graph module subgraph (Graph): a valid subgraph that contains copied nodes from the parent graph class_name (str): name for the submodule """ # Loop through all module calls (call_module) and param fetches (get_attr) # in this component, creating HolderModules as necessary to match the path. # e.g. if in the original module there's a get_attr node fetches "conv.weight". # We create a HolderModule as root -> add a HolderModule named "conv" -> # make "weight" a attribute of "conv" HolderModule and point to conv.weight in # the original module. submodule = HolderModule({}) for n in subgraph.nodes: if n.op not in ("call_module", "get_attr"): continue target = n.target assert isinstance(target, str) target_name_parts = target.split(".") curr = submodule orig_gm = gm for name in target_name_parts[:-1]: if not hasattr(curr, name): curr.add_module(name, HolderModule({})) curr = getattr(curr, name) orig_gm = getattr(orig_gm, name) leaf_node_name = target_name_parts[-1] leaf_node = getattr(orig_gm, leaf_node_name) # Relies on custom __setattr__ magic. setattr(curr, leaf_node_name, leaf_node) return GraphModule(submodule, subgraph, class_name)
def replace_target_nodes_with( fx_module: GraphModule, old_op: str, old_target: Target, new_op: str, new_target: Target, ): """Modifies all nodes in fx_module.graph.nodes which match the specified op code and target, and updates them to match the new op code and target""" new_graph = Graph() val_map : Dict[Node, Node] = {} for node in fx_module.graph.nodes: if node.op == old_op and node.target == old_target: args = map_arg(node.args, lambda n: val_map[n]) kwargs = map_arg(node.kwargs, lambda n: val_map[n]) assert isinstance(args, tuple) assert isinstance(kwargs, dict) val_map[node] = new_graph.create_node(new_op, new_target, args, kwargs, node.name) else: val_map[node] = new_graph.node_copy(node, lambda n : val_map[n]) fx_module.graph = new_graph
def symbolic_trace_with_rewrite( root: Union[torch.nn.Module, Callable]) -> GraphModule: return GraphModule( root if isinstance(root, torch.nn.Module) else torch.nn.Module(), RewritingTracer().trace(root), )
def serialize_module(fx_module: GraphModule, weights: Dict, name_prefix="") -> Dict: """Recursively Serializes a graph module (fx_module) to a dictionary which is later exported to JSON. It also adds all weights the provided weights dictionary by qualified_name. Dictionary Schema: MODULE { modules: {module_name: MODULE], nodes: [NODE], weights {qualified_name: WEIGHT}, } NODE { shape: [], dtype: dtype, target: target, op_code: op_code, name: name, args: [], kwargs: {} } WEIGHT { dtype: dtype, is_quantized: bool, shape: [], quantization_info: QUANTIZATION } QUANTIZATION { qscheme: qscheme, q_scale: float, q_zero_point: float, q_per_channel_scales, [], q_per_channel_zero_points: [], q_per_channel_axis, int } """ serialized_dict: Dict[str, Any] = {} serialized_dict["modules"] = {} serialized_dict["weights"] = {} serialized_dict["nodes"] = [] parameters = fx_module.named_parameters() prefix = f"{name_prefix}." if name_prefix else "" submodules = dict(fx_module.named_modules()) for name, p in parameters: if isinstance(p, torch.Tensor): weight = serialize_weight(p) serialized_dict["weights"][prefix + name] = weight weights[prefix + name] = p # Note: lift_lowering_attrs_to_nodes is only used to support leaf modules # that cannot currently be symbolically traced into, e.g. batch norm. lift_lowering_attrs_to_nodes(fx_module) for node in fx_module.graph.nodes: node_rep: Dict[str, Any] = {} # Get shape/type info, currently not needed for call_module. if node.op != "call_module" or not isinstance(submodules[node.target], GraphModule): shape, dtype = get_shape_and_dtype(node) node_rep["shape"] = serialize_shape(shape) node_rep["dtype"] = str(dtype) # Recurse down into any submodules we are calling. if node.op == "call_module": if isinstance(submodules[node.target], GraphModule): serialized_module = serialize_module( getattr(fx_module, node.target), weights, node.target) serialized_dict["modules"][node.target] = serialized_module else: node_rep["parameters"] = serialize_leaf_module( node, serialized_dict["weights"], weights, prefix + node.target, ) if node.op == "call_function": node_rep["target"] = _get_qualified_name(node.target) else: node_rep["target"] = str(node.target) # Make sure we capture all constants. if node.op == "get_attr": # If we are targeting a parent constant we update the target. if node.target.startswith("parent."): node.name = node.name[len("parent."):] node_rep["target"] = str(node.target[len("parent."):]) weight = serialize_weight( weights[node.target[len("parent."):]]) serialized_dict["weights"][ node.target[len("parent."):]] = weight else: # Iterate through the module hierarchy to find the attr. target = fx_module split = node.target.split(".") assert len(split) while len(split): target = getattr(target, split.pop(0)) qualname = prefix + node.target # Check that the target is a tensor, and that we haven't added it already from a leaf module. if isinstance(target, torch.Tensor) and qualname not in weights: weight = serialize_weight(target) serialized_dict["weights"][prefix + node.target] = weight weights[prefix + node.target] = target node_rep["op_code"] = node.op node_rep["name"] = node.name if node.op == "output": def get_output_info(arg: Node) -> Argument: shape, dtype = get_shape_and_dtype(arg) return { "is_node": True, "name": str(arg), "shape": serialize_shape(shape), "dtype": str(dtype), } node_rep["args"] = map_arg( node.args, get_output_info, ) # If there're multiple outputs then node_rep["args"][0] will be a tuple. # In this case we want to unpack the tuple. if isinstance(node_rep["args"][0], tuple): node_rep["args"] = node_rep["args"][0] else: node_rep["args"] = map_arg( node.args, lambda arg: { "is_node": True, "name": str(arg) }) node_rep["kwargs"] = map_arg( node.kwargs, lambda arg: { "is_node": True, "name": str(arg) }) serialized_dict["nodes"] += [node_rep] return serialized_dict
def serialize_module(fx_module: GraphModule, weights: Dict, name_prefix="") -> Dict: """Recursively Serializes a graph module (fx_module) to a dictionary which is later exported to JSON. It also adds all weights the provided weights dictionary by qualified_name. Dictionary Schema: MODULE { modules: {module_name: MODULE], nodes: [NODE], weights {qualified_name: WEIGHT}, } NODE { shape: [], dtype: dtype, is_quantized: bool, target: target, op_code: op_code, name: name, args: [], kwargs: {} } WEIGHT { dtype: dtype, is_quantized: bool, shape: [], QUANTIZATION, } QUANTIZATION { qscheme: qscheme, q_scale: float, q_zero_point: float, q_per_channel_scales, [], q_per_channel_zero_points: [], q_per_channel_axis, int } """ serialized_dict: Dict[str, Any] = {} serialized_dict["modules"] = {} serialized_dict["weights"] = {} serialized_dict["nodes"] = [] submodules = dict(fx_module.named_modules()) prefix = f"{name_prefix}." if name_prefix else "" def add_weight_tensors(named_tensors): for name, p in named_tensors: if name.startswith("parent.") or not isinstance(p, torch.Tensor): continue weight = serialize_weight(p) serialized_dict["weights"][prefix + name] = weight weights[prefix + name] = p add_weight_tensors(fx_module.named_parameters()) add_weight_tensors(fx_module.named_buffers()) def get_node_info(node): shape, dtype = get_shape_and_dtype(node) tensor_meta = node.meta.get('tensor_meta') if not tensor_meta: raise RuntimeError(f'Node {node} has no tensor metadata! Ensure shape ' f'propagation has been run!') node_rep = { "shape": serialize_shape(shape), "dtype": str(dtype), "is_quantized": tensor_meta.is_quantized, } if tensor_meta.is_quantized: node_rep["qscheme"] = str(tensor_meta.qscheme) if tensor_meta.qscheme in {torch.per_tensor_affine, torch.per_tensor_symmetric}: node_rep["q_scale"] = tensor_meta.q_scale node_rep["q_zero_point"] = tensor_meta.q_zero_point return node_rep # Note: lift_lowering_attrs_to_nodes is only used to support leaf modules # that cannot currently be symbolically traced into, e.g. batch norm. lift_lowering_attrs_to_nodes(fx_module) for node in fx_module.graph.nodes: node_rep: Dict[str, Any] = {} # Get shape/type info, currently not needed for call_module node # whose target is a GraphModule and output node. if ( not ( node.op == "call_module" and isinstance(submodules[node.target], GraphModule) ) and node.op != "output" ): node_rep.update(get_node_info(node)) # Recurse down into any submodules we are calling. if node.op == "call_module": if isinstance(submodules[node.target], GraphModule): serialized_module = serialize_module( getattr(fx_module, node.target), weights, node.target ) serialized_dict["modules"][node.target] = serialized_module else: node_rep["parameters"] = serialize_leaf_module( node, serialized_dict["weights"], weights, prefix + node.target, ) if node.op == "call_function": node_rep["target"] = _get_qualified_name(node.target) else: node_rep["target"] = str(node.target) # Make sure we capture all constants. if node.op == "get_attr": # If we are targeting a parent constant we update the target. if node.target.startswith("parent."): stripped_name = node.target[len("parent.") :] node.name = stripped_name node_rep["target"] = stripped_name weight = serialize_weight(weights[stripped_name]) serialized_dict["weights"][stripped_name] = weight else: # Find the actual target parameter/buffer from the fx_module. submod_path, _, target_name = node.target.rpartition(".") submod: Optional[torch.nn.Module] = ( fx_module.get_submodule(submod_path) if submod_path else fx_module ) assert submod is not None, f"submod {submod_path} not found" target = getattr(submod, target_name, None) assert target is not None, f"{target_name} not an attr of {submod_path}" qualname = prefix + node.target # Check that the target is a tensor, and that we haven't added it already from a leaf module. if isinstance(target, torch.Tensor) and qualname not in weights: weight = serialize_weight(target) serialized_dict["weights"][qualname] = weight weights[qualname] = target node_rep["op_code"] = node.op node_rep["name"] = node.name def get_arg_info(arg: Argument) -> Any: if isinstance(arg, torch.fx.Node): return {"is_node": True, "name": str(arg)} elif isinstance(arg, torch.dtype): return str(arg) else: return arg def get_output_arg_info(arg: Node) -> Dict[str, Any]: node_rep: Dict[str, Any] = get_arg_info(arg) node_rep.update(get_node_info(arg)) return node_rep if node.op == "output": node_rep["args"] = map_arg( node.args, get_output_arg_info, ) # If there're multiple outputs then node_rep["args"][0] will be a tuple. # In this case we want to unpack the tuple. if isinstance(node_rep["args"][0], tuple): node_rep["args"] = node_rep["args"][0] else: node_rep["args"] = map_aggregate( node.args, get_arg_info ) node_rep["kwargs"] = map_aggregate( node.kwargs, get_arg_info ) serialized_dict["nodes"] += [node_rep] return serialized_dict
def serialize_module(fx_module: GraphModule, weights: Dict, name_prefix="") -> Dict: """Recursively Serializes a graph module (fx_module) to a dictionary which is later exported to JSON. It also adds all weights the provided weights dictionary by qualified_name. Dictionary Schema: MODULE { modules: {module_name: MODULE], nodes: [NODE], weights {qualified_name: WEIGHT}, } NODE { shape: [], dtype: dtype, target: target, op_code: op_code, name: name, args: [], kwargs: {} } WEIGHT { dtype: dtype, is_quantized: bool, shape: [], quantization_info: QUANTIZATION } QUANTIZATION { qscheme: qscheme, q_scale: float, q_zero_point: float, q_per_channel_scales, [], q_per_channel_zero_points: [], q_per_channel_axis, int } """ serialized_dict: Dict[str, Any] = {} serialized_dict["modules"] = {} serialized_dict["weights"] = {} serialized_dict["nodes"] = [] parameters = fx_module.named_parameters() prefix = f"{name_prefix}." if name_prefix else "" submodules = dict(fx_module.named_modules()) for name, p in parameters: if isinstance(p, torch.Tensor): weight = serialize_weight(p) serialized_dict["weights"][prefix + name] = weight weights[prefix + name] = p for node in fx_module.graph.nodes: node_rep: Dict[str, Any] = {} # Get shape/type info, currently not needed for call_module. if node.op != "call_module" or not isinstance(submodules[node.target], GraphModule): shape = getattr(node, "shape", None) if shape: node_rep["shape"] = serialize_shape(shape) else: raise RuntimeError( "Node has no shape attr, this is likely because shape propagation has not been run on this Graph." ) dtype = getattr(node, "dtype", None) if dtype: node_rep["dtype"] = str(dtype) else: raise RuntimeError( "Node has no dtype attr, this is likely because shape propagation has not been run on this Graph." ) # Recurse down into any submodules we are calling. if node.op == "call_module": if isinstance(submodules[node.target], GraphModule): serialized_module = serialize_module( getattr(fx_module, node.target), weights, node.target) serialized_dict["modules"][node.target] = serialized_module else: node_rep["parameters"] = serialize_leaf_module( submodules[node.target], serialized_dict["weights"], weights, prefix + node.target, ) if node.op == "call_function": node_rep["target"] = get_qualified_name(node.target) else: node_rep["target"] = str(node.target) # Make sure we capture all constants. if node.op == "get_attr": target = getattr(fx_module, node.target) qualname = prefix + node.target if isinstance(target, torch.Tensor) and qualname not in weights: weight = serialize_weight(target) serialized_dict["weights"][prefix + node.target] = weight weights[prefix + node.target] = target node_rep["op_code"] = node.op node_rep["name"] = node.name node_rep["args"] = map_arg( node.args, lambda arg: { "is_node": True, "name": str(arg) }) node_rep["kwargs"] = map_arg( node.kwargs, lambda arg: { "is_node": True, "name": str(arg) }) serialized_dict["nodes"] += [node_rep] return serialized_dict
def serialize_module(fx_module: GraphModule, weights: Dict, name_prefix="") -> Dict: """Recursively Serializes a graph module (fx_module) to a dictionary which is later exported to JSON. It also adds all weights the provided weights dictionary by qualified_name. Dictionary Schema: MODULE { modules: {module_name: MODULE], nodes: [NODE], weights {qualified_name: WEIGHT}, } NODE { shape: [], stride: [], dtype: dtype, is_quantized: bool, target: target, op_code: op_code, name: name, args: [], kwargs: {} } WEIGHT { dtype: dtype, is_quantized: bool, shape: [], QUANTIZATION, } QUANTIZATION { qscheme: qscheme, q_scale: float, q_zero_point: float, q_per_channel_scales, [], q_per_channel_zero_points: [], q_per_channel_axis, int } """ serialized_dict: Dict[str, Any] = {} serialized_dict["modules"] = {} serialized_dict["weights"] = {} serialized_dict["nodes"] = [] submodules = dict(fx_module.named_modules()) prefix = f"{name_prefix}." if name_prefix else "" def add_weight_tensors(named_tensors): for name, p in named_tensors: if name.startswith("parent.") or not isinstance(p, torch.Tensor): continue weight_dict = serialize_weight(p, weights, prefix + name) serialized_dict["weights"].update(weight_dict) weights[prefix + name] = p add_weight_tensors(fx_module.named_parameters()) add_weight_tensors(fx_module.named_buffers()) def get_node_info(node): tensor_meta = get_tensor_meta(node) node_rep = { "shape": serialize_shape(tensor_meta.shape), "dtype": str(tensor_meta.dtype), "requires_grad": str(tensor_meta.requires_grad), "stride": serialize_stride(tensor_meta.stride), "is_quantized": tensor_meta.is_quantized, } if tensor_meta.is_quantized: node_rep["qscheme"] = str(tensor_meta.qscheme) if tensor_meta.qscheme in { torch.per_tensor_affine, torch.per_tensor_symmetric, }: node_rep["q_scale"] = tensor_meta.q_scale node_rep["q_zero_point"] = tensor_meta.q_zero_point return node_rep # Note: lift_lowering_attrs_to_nodes is only used to support leaf modules # that cannot currently be symbolically traced into, e.g. batch norm. lift_lowering_attrs_to_nodes(fx_module) for node in fx_module.graph.nodes: node_rep: Dict[str, Any] = {} # Get shape/type info, currently not needed for call_module node # whose target is a GraphModule and output node. if ( not ( node.op == "call_module" and isinstance(submodules[node.target], GraphModule) ) and node.op != "output" ): node_rep.update(get_node_info(node)) # Recurse down into any submodules we are calling. if node.op == "call_module": if isinstance(submodules[node.target], GraphModule): serialized_module = serialize_module( getattr(fx_module, node.target), weights, node.target ) serialized_dict["modules"][node.target] = serialized_module else: node_rep["parameters"] = serialize_leaf_module( node, serialized_dict["weights"], weights, prefix + node.target, ) if node.op == "call_function": node_rep["target"] = _get_qualified_name(node.target) else: node_rep["target"] = str(node.target) # Make sure we capture all constants. if node.op == "get_attr": # If we are targeting a parent constant we update the target. if node.target.startswith("parent."): stripped_name = node.target[len("parent.") :] node.name = stripped_name node_rep["target"] = stripped_name weight = serialize_weight( weights[stripped_name], weights, node.target[len("parent.") :] ) # For quantized embedding tables we need to update the shape/type, # so we check if the users of this get_attr is a quantized EB and this is the weight for the EB. user_targets = { _get_qualified_name( n.target ).replace("torch.fx.experimental.fx_acc.", "").replace("glow.fb.fx.", ""): n for n in node.users.keys() } if ( "acc_ops.embedding_bag_byte_rowwise_offsets" in user_targets and str( user_targets[ "acc_ops.embedding_bag_byte_rowwise_offsets" ].kwargs["weight"] ) == stripped_name ): weight[stripped_name]["dtype"] = "acc.uint8fused" # Same as above, but for the 4 bit version. if ( "acc_ops.embedding_bag_4bit_rowwise_offsets" in user_targets and str( user_targets[ "acc_ops.embedding_bag_4bit_rowwise_offsets" ].kwargs["weight"] ) == stripped_name ): weight[stripped_name]["dtype"] = "acc.uint4fused" serialized_dict["weights"].update(weight) else: # Find the actual target parameter/buffer from the fx_module. submod_path, _, target_name = node.target.rpartition(".") submod: Optional[torch.nn.Module] = ( fx_module.get_submodule(submod_path) if submod_path else fx_module ) assert submod is not None, f"submod {submod_path} not found" target = getattr(submod, target_name, None) assert target is not None, f"{target_name} not an attr of {submod_path}" qualname = prefix + node.target # Check that the target is a tensor, and that we haven't added it already from a leaf module. if isinstance(target, torch.Tensor) and qualname not in weights: weight = serialize_weight(target, weights, qualname) serialized_dict["weights"].update(weight) weights[qualname] = target node_rep["op_code"] = node.op node_rep["name"] = node.name def get_user_info(user_node: Argument) -> Any: return {"is_node": True, "name": str(user_node)} def get_arg_info(arg: Argument) -> Any: if isinstance(arg, torch.fx.Node): return {"is_node": True, "name": str(arg)} elif isinstance(arg, (torch.dtype, torch.memory_format, torch.qscheme)): return str(arg) else: return arg def get_output_arg_info(arg: Node) -> Dict[str, Any]: node_rep: Dict[str, Any] = get_arg_info(arg) node_rep.update(get_node_info(arg)) return node_rep if node.op == "output": node_rep["args"] = map_arg( node.args, get_output_arg_info, ) # If there're multiple outputs then node_rep["args"][0] will be a tuple. # In this case we want to unpack the tuple. if isinstance(node_rep["args"][0], tuple): node_rep["args"] = node_rep["args"][0] else: node_rep["args"] = map_aggregate(node.args, get_arg_info) node_rep["kwargs"] = map_aggregate(node.kwargs, get_arg_info) node_rep["users"] = map_aggregate(list(node.users.keys()), get_user_info) serialized_dict["nodes"] += [node_rep] return serialized_dict
def _match_nodes(self, pn: Node, gn: Node, original_module: GraphModule, pattern_module: GraphModule) -> bool: if isinstance(pn, (tuple, list)): if not isinstance(gn, type(pn)): return False return all( self._match_nodes(a1, a2, original_module, pattern_module) for a1, a2 in zip(pn, gn)) # type: ignore[call-overload] elif isinstance(pn, dict): if not isinstance(gn, dict): return False return pn.keys() == gn.keys() and \ all(self._match_nodes(v1, v2, original_module, pattern_module) for v1, v2 in zip(pn.values(), gn.values())) # Check if we've already matched these nodes in the current # traversal if pn in self.nodes_map: return self.nodes_map[pn] == gn PRIM_TYPES = (int, float, torch.dtype) # if both pattern and graph are not Node, we check for equality of these values if not isinstance(pn, Node) and not isinstance(gn, Node): return pn == gn # trying to match the input in pattern graph with a primitive type values if isinstance(gn, PRIM_TYPES): if isinstance(pn, Node) and pn.op == "placeholder": self.nodes_map[pn] = gn return True else: return False original_modules = dict(original_module.named_modules()) pattern_modules = dict(pattern_module.named_modules()) def attributes_are_equal(pn: Node, gn: Node) -> bool: # Use placeholder and output nodes as wildcards. The # only exception is that an output node can't match # a placeholder if (pn.op == "placeholder" or (pn.op == "output" and gn.op != "placeholder")): return True elif pn.op == "get_attr" and gn.op == "get_attr": # assuming get_attr nodes are the same return True elif pn.op == "call_module" and gn.op == "call_module": original_m = original_modules[gn.target] pattern_m = pattern_modules[pn.target] return type(original_m) == type(pattern_m) return pn.op == gn.op and pn.target == gn.target # Terminate early if the node attributes are not equal if not attributes_are_equal(pn, gn): return False # Optimistically mark `pn` as a match for `gn` self.nodes_map[pn] = gn # Traverse the use-def relationships to ensure that `pn` is a true # match for `gn` if pn.op == "placeholder": return True if (pn.op != "output" and len(pn.args) != len(gn.args)): return False if pn.op == "output": match_found = any( self._match_nodes(pn.all_input_nodes[0], gn_, original_module, pattern_module) for gn_ in gn.all_input_nodes) else: # using args here to make sure we can match Node and non-Node # arguments # also allows us to match a Node with a primitive type value match_found = (len(pn.args) == len(gn.args) and all(self._match_nodes(pn_, gn_, original_module, pattern_module) # type: ignore[arg-type] for pn_, gn_ \ in zip(pn.args, gn.args))) if not match_found: self.nodes_map.pop(pn) return False return True
def replace_pattern( gm: GraphModule, pattern: Callable, replacement: Callable, is_match_filters: Optional[List[Callable]] = None) -> List[Match]: """ Matches all possible non-overlapping sets of operators and their data dependencies (``pattern``) in the Graph of a GraphModule (``gm``), then replaces each of these matched subgraphs with another subgraph (``replacement``). Args: ``gm``: The GraphModule that wraps the Graph to operate on ``pattern``: The subgraph to match in ``gm`` for replacement ``replacement``: The subgraph to replace ``pattern`` with Returns: List[Match]: A list of ``Match`` objects representing the places in the original graph that ``pattern`` was matched to. The list is empty if there are no matches. ``Match`` is defined as: .. code-block:: python class Match(NamedTuple): # Node from which the match was found anchor: Node # Maps nodes in the pattern subgraph to nodes in the larger graph nodes_map: Dict[Node, Node] Examples: .. code-block:: python import torch from torch.fx import symbolic_trace, subgraph_rewriter class M(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, w1, w2): m1 = torch.cat([w1, w2]).sum() m2 = torch.cat([w1, w2]).sum() return x + torch.max(m1) + torch.max(m2) def pattern(w1, w2): return torch.cat([w1, w2]).sum() def replacement(w1, w2): return torch.stack([w1, w2]) traced_module = symbolic_trace(M()) subgraph_rewriter.replace_pattern(traced_module, pattern, replacement) The above code will first match ``pattern`` in the ``forward`` method of ``traced_module``. Pattern-matching is done based on use-def relationships, not node names. For example, if you had ``p = torch.cat([a, b])`` in ``pattern``, you could match ``m = torch.cat([a, b])`` in the original ``forward`` function, despite the variable names being different (``p`` vs ``m``). The ``return`` statement in ``pattern`` is matched based on its value only; it may or may not match to the ``return`` statement in the larger graph. In other words, the pattern doesn't have to extend to the end of the larger graph. When the pattern is matched, it will be removed from the larger function and replaced by ``replacement``. If there are multiple matches for ``pattern`` in the larger function, each non-overlapping match will be replaced. In the case of a match overlap, the first found match in the set of overlapping matches will be replaced. ("First" here being defined as the first in a topological ordering of the Nodes' use-def relationships. In most cases, the first Node is the parameter that appears directly after ``self``, while the last Node is whatever the function returns.) One important thing to note is that the parameters of the ``pattern`` Callable must be used in the Callable itself, and the parameters of the ``replacement`` Callable must match the pattern. The first rule is why, in the above code block, the ``forward`` function has parameters ``x, w1, w2``, but the ``pattern`` function only has parameters ``w1, w2``. ``pattern`` doesn't use ``x``, so it shouldn't specify ``x`` as a parameter. As an example of the second rule, consider replacing .. code-block:: python def pattern(x, y): return torch.neg(x) + torch.relu(y) with .. code-block:: python def replacement(x, y): return torch.relu(x) In this case, ``replacement`` needs the same number of parameters as ``pattern`` (both ``x`` and ``y``), even though the parameter ``y`` isn't used in ``replacement``. After calling ``subgraph_rewriter.replace_pattern``, the generated Python code looks like this: .. code-block:: python def forward(self, x, w1, w2): stack_1 = torch.stack([w1, w2]) sum_1 = stack_1.sum() stack_2 = torch.stack([w1, w2]) sum_2 = stack_2.sum() max_1 = torch.max(sum_1) add_1 = x + max_1 max_2 = torch.max(sum_2) add_2 = add_1 + max_2 return add_2 """ # Get the module and graph for `gm`, `pattern`, `replacement` original_module = gm original_graph = original_module.graph pattern_module = symbolic_trace(pattern) pattern_graph = pattern_module.graph replacement_module = symbolic_trace(replacement) replacement_graph = replacement_module.graph # Find all possible pattern matches in original_graph. Note that # pattern matches may overlap with each other. matcher = _SubgraphMatcher(pattern_graph) matches: List[Match] = [] # Consider each node as an "anchor" (deepest matching graph node) for anchor in original_graph.nodes: if matcher.matches_subgraph_from_anchor(anchor, original_module, pattern_module): def pattern_is_contained(nodes_map: Dict[Node, Node]) -> bool: # `lookup` represents all the nodes in `original_graph` # that are part of `pattern` lookup: Dict[Node, Node] = {v: k for k, v in nodes_map.items()} for n in lookup.keys(): # Nodes that can "leak"... if not isinstance(lookup[n], Node): continue # Placeholders (by definition) if lookup[n].op == "placeholder": continue # Pattern output (acts as a container) if lookup[n].op == "output": continue # Result contained by pattern output (what we'll # hook in to the new Graph, thus what we'll # potentially use in other areas of the Graph as # an input Node) if (len(lookup[n].users) == 1 and list( lookup[n].users.keys())[0].op == "output"): continue if not isinstance(n, Node): continue for user in n.users: # If this node has users that were not in # `lookup`, then it must leak out of the # pattern subgraph if user not in lookup: return False return True # It's not a match if the pattern leaks out into the rest # of the graph if pattern_is_contained(matcher.nodes_map): # Shallow copy nodes_map matches.append( Match(anchor=anchor, nodes_map=copy.copy({ key: value for key, value in matcher.nodes_map.items() }))) # The set of all nodes in `original_graph` that we've seen thus far # as part of a pattern match replaced_nodes: Set[Node] = set() # As we progressively replace nodes, we'll need to keep track of how the match results should change match_changed_node: Dict[Node, Node] = dict() # Return True if one of the nodes in the current match has already # been used as part of another match def overlaps_with_prev_match(match: Match) -> bool: for pn, gn in match.nodes_map.items(): if not isinstance(pn, Node): continue if pn.op in ["placeholder", "output"]: continue if not isinstance(gn, Node): continue if gn in replaced_nodes and gn.op != "placeholder": return True return False if is_match_filters is None: is_match_filters = [] def is_match(match: Match): # for mypy assert is_match_filters is not None for filter in is_match_filters: if not filter(match, pattern_graph, replacement_graph): return False return True for match in matches: # Skip overlapping matches if overlaps_with_prev_match(match): continue if not is_match(match): continue # Map replacement graph nodes to their copy in `original_graph` val_map: Dict[Node, Node] = {} pattern_placeholders = [ n for n in pattern_graph.nodes if n.op == "placeholder" ] assert len(pattern_placeholders) > 0 replacement_placeholders = [ n for n in replacement_graph.nodes if n.op == "placeholder" ] assert len(pattern_placeholders) == len(replacement_placeholders) placeholder_map = { r: p for r, p in zip(replacement_placeholders, pattern_placeholders) } # node from `original_graph` that matched with the output node # in `pattern` subgraph_output: Node = match.anchor def mark_node_as_replaced(n: Node) -> None: if n not in match.nodes_map.values(): return for n_ in n.all_input_nodes: mark_node_as_replaced(n_) replaced_nodes.add(n) for input_node in subgraph_output.all_input_nodes: mark_node_as_replaced(input_node) # Initialize `val_map` with mappings from placeholder nodes in # `replacement` to their corresponding node in `original_graph` for replacement_node in replacement_placeholders: # Get the `original_graph` placeholder node # corresponding to the current `replacement_node` pattern_node = placeholder_map[replacement_node] original_graph_node = match_changed_node.get( match.nodes_map[pattern_node], match.nodes_map[pattern_node]) # Populate `val_map` val_map[replacement_node] = original_graph_node # Copy the replacement graph over with original_graph.inserting_before(subgraph_output): copied_output = original_graph.graph_copy(replacement_graph, val_map) # Hook the output Node of the replacement subgraph in to the # original Graph at the correct location # CASE 1: We need to hook the replacement subgraph in somewhere # in the middle of the graph. We replace the Node in the # original graph that corresponds to the end of the pattern # subgraph if subgraph_output.op != "output": pattern_outputs = [ n for n in pattern_graph.nodes if n.op == "output" ] assert len(pattern_outputs) > 0 replacement_outputs = [ n for n in replacement_graph.nodes if n.op == "output" ] assert len(replacement_outputs) == len(pattern_outputs) outputs_map = { p: r for r, p in zip(replacement_outputs, pattern_outputs) } for pn, gn in match.nodes_map.items(): if not isinstance(gn, Node): continue if gn.op == "placeholder": continue # Search for the node corresponding to the output of the pattern if pn.op != "output": continue assert subgraph_output == gn # Update all anchor inputs to the new nodes rn = outputs_map[pn] for pn_input, rn_input in zip(pn.args, rn.args): gn_input = match.nodes_map[pn_input] # type: ignore[index] rn_input_in_original_graph = val_map[rn_input] gn_input.replace_all_uses_with(rn_input_in_original_graph) # We store the updated node point in case other nodes want to use it match_changed_node[gn_input] = rn_input_in_original_graph assert subgraph_output.op != "output" # CASE 2: The pattern subgraph match extends to the end of the # original graph, so we need to change the current graph's # output Node to reflect the insertion of the replacement graph. # We'll keep the current output Node, but update its args and # `_input_nodes` as necessary else: subgraph_output.args = ((copied_output, )) if isinstance(copied_output, Node): subgraph_output._input_nodes = {copied_output: None} assert isinstance(copied_output, Node) # Erase the `pattern` nodes for node in reversed(original_graph.nodes): if len(node.users ) == 0 and node.op != "output" and node.op != "placeholder": original_graph.erase_node(node) # Update the passed-in GraphModule to reflect the new state of # `original_graph` gm.recompile() # If `replacement` was an nn.Module, we'll need to make sure that # all the submodules have been copied over correctly if isinstance(replacement, torch.nn.Module): _replace_submodules(gm, replacement) return matches
def serialize_module(fx_module: GraphModule, weights: Dict, name_prefix="") -> Dict: """Recursively Serializes a graph module (fx_module) to a dictionary which is later exported to JSON. It also adds all weights the provided weights dictionary by qualified_name. Dictionary Schema: MODULE { modules: {module_name: MODULE], nodes: [NODE], weights {qualified_name: WEIGHT}, } NODE { shape: [], stride: [], dtype: dtype, is_quantized: bool, target: target, op_code: op_code, name: name, args: [], kwargs: {} } WEIGHT { dtype: dtype, is_quantized: bool, shape: [], QUANTIZATION, } QUANTIZATION { qscheme: qscheme, q_scale: float, q_zero_point: float, q_per_channel_scales, [], q_per_channel_zero_points: [], q_per_channel_axis, int } """ serialized_dict: Dict[str, Any] = {} serialized_dict["modules"] = {} serialized_dict["weights"] = {} serialized_dict["nodes"] = [] submodules = dict(fx_module.named_modules()) prefix = f"{name_prefix}." if name_prefix else "" def get_node_info(node): tensor_meta = get_tensor_meta(node) node_rep = { "shape": serialize_shape(tensor_meta.shape), "dtype": str(tensor_meta.dtype), "requires_grad": str(tensor_meta.requires_grad), "stride": serialize_stride(tensor_meta.stride), "is_quantized": tensor_meta.is_quantized, } if tensor_meta.is_quantized: node_rep["qscheme"] = str(tensor_meta.qparams["qscheme"]) if tensor_meta.qparams["qscheme"] in { torch.per_tensor_affine, torch.per_tensor_symmetric, }: node_rep["q_scale"] = tensor_meta.qparams["scale"] node_rep["q_zero_point"] = tensor_meta.qparams["zero_point"] # Add all extra lowering_info that was provided in node.meta. lowering_info = node.meta.get("lowering_info") if lowering_info is not None: overlapping_keys = node_rep.keys() & lowering_info.keys() assert ( len(overlapping_keys) == 0 ), f"Overlap found between lowering_info and node_rep: {overlapping_keys}" node_rep.update(lowering_info) return node_rep # Note: lift_lowering_attrs_to_nodes is only used to support leaf modules # that cannot currently be symbolically traced into, e.g. batch norm. lift_lowering_attrs_to_nodes(fx_module) for node in fx_module.graph.nodes: node_rep: Dict[str, Any] = {} # Get shape/type info, currently not needed for call_module node # whose target is a GraphModule and output node. if (not (node.op == "call_module" and isinstance(submodules[node.target], GraphModule)) and node.op != "output"): node_rep.update(get_node_info(node)) # Recurse down into any submodules we are calling. if node.op == "call_module": if isinstance(submodules[node.target], GraphModule): serialized_module = serialize_module( getattr(fx_module, node.target), weights, node.target) serialized_dict["modules"][node.target] = serialized_module else: node_rep["parameters"] = serialize_leaf_module( node, serialized_dict["weights"], weights, prefix + node.target, ) if node.op == "call_function": node_rep["target"] = _get_qualified_name(node.target) else: node_rep["target"] = str(node.target) # Make sure we capture all constants. if node.op == "get_attr": # If we are targeting a parent constant we update the target. if node.target.startswith("parent."): qualname = node.target[len("parent."):] node.name = qualname node_rep["target"] = qualname else: qualname = prefix + node.target # Find the actual target parameter/buffer from the fx_module. submod_path, _, target_name = node.target.rpartition(".") submod: Optional[torch.nn.Module] = ( fx_module.get_submodule(submod_path) if submod_path else fx_module) assert submod is not None, f"submod {submod_path} not found" target = getattr(submod, target_name, None) assert target is not None, f"{target_name} not an attr of {submod_path}" # Check that the target is a tensor, and that we haven't added it already from a leaf module. if isinstance(target, torch.Tensor) and qualname not in weights: weight = serialize_weight(target, weights, qualname) _update_weight_fused_dtypes(weight, qualname, node) serialized_dict["weights"].update(weight) weights[qualname] = target elif node.op == "placeholder": ph_type = node.meta.get("ph_type", "") assert ( ph_type == "" or ph_type == "input_ph" or ph_type == "output_ph" ), "When present, placeholder type must be 'input_ph' or 'ouput_ph'" if ph_type == "input_ph": node_rep["ph_type"] = "input_ph" elif ph_type == "output_ph": node_rep["ph_type"] = "output_ph" node_rep["op_code"] = node.op node_rep["name"] = node.name def get_user_info(user_node: Argument) -> Any: return {"is_node": True, "name": str(user_node)} def get_arg_info(arg: Argument) -> Any: if isinstance(arg, torch.fx.Node): return {"is_node": True, "name": str(arg)} elif isinstance(arg, (torch.dtype, torch.memory_format, torch.qscheme)): return str(arg) else: return arg def get_output_arg_info(arg: Node) -> Dict[str, Any]: node_rep: Dict[str, Any] = get_arg_info(arg) node_rep.update(get_node_info(arg)) return node_rep if node.op == "output": node_rep["args"] = map_arg( node.args, get_output_arg_info, ) # If there're multiple outputs then node_rep["args"][0] will be a tuple or # list. In this case we want to unpack the tuple or list. if isinstance(node_rep["args"][0], (tuple, list)): node_rep["args"] = node_rep["args"][0] else: node_rep["args"] = map_aggregate(node.args, get_arg_info) node_rep["kwargs"] = map_aggregate(node.kwargs, get_arg_info) node_rep["users"] = map_aggregate(list(node.users.keys()), get_user_info) serialized_dict["nodes"] += [node_rep] return serialized_dict