| import torch |
| from torch.fx import ( |
| Graph, |
| GraphModule, |
| Node, |
| ) |
| from torch.fx.subgraph_rewriter import replace_pattern_with_filters |
| import torch.nn.functional as F |
| from torch.nn.utils.fusion import fuse_conv_bn_weights |
| # TODO[jerryzh168]: move this to a more general util function |
| from torch.ao.quantization.fx.prepare import ( |
| _is_activation_post_process_node, |
| ) |
| import copy |
| import operator |
| from typing import Any, Callable, Dict, Optional, Tuple |
| |
| __all__ = [ |
| "fold_bn_weights_into_conv_node", |
| "get_aten_graph_module", |
| "remove_tensor_overload_for_qdq_ops", |
| ] |
| |
| def _get_tensor_constant_from_node(node, m): |
| if node is None: |
| return None |
| assert node.op == "get_attr" |
| return getattr(m, node.target) |
| |
| def _get_all_arguments(orig_args, orig_kwargs, args_schema): |
| all_args = [] |
| for i, schema in enumerate(args_schema): |
| if schema.name in orig_kwargs: |
| all_args.append(orig_kwargs[schema.name]) |
| elif not schema.kwarg_only and i < len(orig_args): |
| all_args.append(orig_args[i]) |
| else: |
| all_args.append(schema.default_value) |
| return all_args |
| |
| def fold_bn_weights_into_conv_node( |
| conv_node: Node, |
| conv_weight_node: Node, |
| conv_bias_node: Optional[Node], |
| bn_node: Node, |
| m: GraphModule |
| ) -> None: |
| # conv args: input, weight, bias, stride, padding, dilation, transposed, ... |
| conv_w = _get_tensor_constant_from_node(conv_weight_node, m) |
| conv_b = _get_tensor_constant_from_node(conv_bias_node, m) |
| transpose = conv_node.args[6] |
| |
| # eval bn args: input, weight, bias, running mean, running var, momentum, eps |
| # train bn args: input, weight, bias, running mean, running var, training, momentum, eps |
| bn_args_schema = bn_node.target._schema.arguments # type: ignore[union-attr] |
| bn_args = _get_all_arguments(bn_node.args, bn_node.kwargs, bn_args_schema) |
| bn_w = _get_tensor_constant_from_node(bn_args[1], m) |
| bn_b = _get_tensor_constant_from_node(bn_args[2], m) |
| bn_rm = _get_tensor_constant_from_node(bn_args[3], m) |
| bn_rv = _get_tensor_constant_from_node(bn_args[4], m) |
| if bn_node.target == torch.ops.aten._native_batch_norm_legit_no_training.default: |
| eps_arg_index = 6 |
| elif bn_node.target == torch.ops.aten._native_batch_norm_legit.default: |
| eps_arg_index = 7 |
| else: |
| raise ValueError("BN node target is unexpected ", bn_node.target) |
| bn_eps = bn_args[eps_arg_index] |
| |
| fused_weight, fused_bias = fuse_conv_bn_weights(conv_w, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b, transpose=transpose) |
| |
| # update the weight and bias for conv |
| conv_args = list(conv_node.args) |
| # calling data since the fused_weight and fused_bias are nn.Parameter |
| weight_attr_name = conv_weight_node.target |
| assert isinstance(weight_attr_name, str) |
| setattr(m, weight_attr_name, fused_weight) |
| if conv_bias_node is not None: |
| bias_attr_name = conv_bias_node.target |
| else: |
| bias_attr_name = weight_attr_name + "_bias" |
| with m.graph.inserting_before(conv_node): |
| get_bias_node = m.graph.get_attr(bias_attr_name) |
| # NOTE: here we assume the bias of conv is not quantized! |
| conv_args[2] = get_bias_node |
| setattr(m, bias_attr_name, fused_bias) # type: ignore[arg-type] |
| conv_node.args = tuple(conv_args) |
| |
| # native_batch_norm has 3 outputs, we expect getitem calls on the output |
| # and we want to replace the uses of getitem 0 with the output of conv |
| # |
| # Before: |
| # conv -> bn - (first output) -> users1 |
| # \ - (second output) -> users2 |
| # \ - (third output) -> users3 |
| # After: |
| # conv -> (first output) -> users1 |
| # bn - |
| # \ - (second output) -> users2 |
| # \ - (third output) -> users3 |
| # if users2 and users3 are empty then bn will be removed through dead code elimination |
| |
| for user in bn_node.users: |
| if user.op != "call_function" or user.target != operator.getitem or user.args[1] != 0: |
| continue |
| user.replace_all_uses_with(conv_node) |
| |
| # fuse conv bn weights, inplace modification of the graph_module and graph |
| def _fuse_conv_bn_(m: GraphModule) -> None: |
| for n in m.graph.nodes: |
| if n.op != "call_function" or n.target != torch.ops.aten._native_batch_norm_legit_no_training.default: |
| continue |
| bn_node = n |
| n = bn_node.args[0] |
| if n.op != "call_function" or n.target != torch.ops.aten.convolution.default: |
| continue |
| conv_node = n |
| conv_weight_node = conv_node.args[1] |
| conv_bias_node = conv_node.args[2] |
| fold_bn_weights_into_conv_node(conv_node, conv_weight_node, conv_bias_node, bn_node, m) |
| |
| m.graph.eliminate_dead_code() |
| m.recompile() |
| |
| # TODO: remove hack when we have better support for pattern matching |
| # move around the observer for addmm |
| def _rearrange_weight_observer_for_decomposed_linear( |
| model: GraphModule, |
| ) -> None: |
| """ |
| Linear is decomposed to `t - addmm` (w/ bias) or `t - mm` (w/o bias) |
| before: |
| weight - t - observer \ |
| input - observer - addmm/mm |
| after: |
| weight - observer - t \ |
| input - observer - addmm/mm |
| """ |
| aten = torch.ops.aten |
| op_to_weight_obs_index = { |
| aten.addmm.default : 2, |
| aten.mm.default : 1, |
| } |
| named_modules = dict(model.named_modules(remove_duplicate=False)) |
| for node in model.graph.nodes: |
| if node.target not in (aten.addmm.default, aten.mm.default): |
| continue |
| root_node = node |
| maybe_weight_obs = root_node.args[op_to_weight_obs_index[root_node.target]] |
| if not _is_activation_post_process_node(maybe_weight_obs, named_modules): |
| continue |
| transpose_node = maybe_weight_obs.args[0] |
| if transpose_node.target != torch.ops.aten.t.default: |
| continue |
| # swap the order of transpose and observation |
| |
| maybe_weight_obs.replace_input_with(transpose_node, transpose_node.args[0]) |
| # remove the transpose node |
| with model.graph.inserting_after(maybe_weight_obs): |
| args = list(transpose_node.args) |
| args[0] = maybe_weight_obs |
| new_transpose_node = model.graph.create_node( |
| "call_function", |
| torch.ops.aten.t.default, |
| tuple(args), |
| transpose_node.kwargs |
| ) |
| root_node.replace_input_with(maybe_weight_obs, new_transpose_node) |
| |
| model.graph.eliminate_dead_code() |
| model.graph.lint() |
| model.recompile() |
| |
| def _get_node_name_to_scope(model: GraphModule) -> Dict[str, Tuple[str, type]]: |
| # TODO: move this information to fx node itself |
| node_name_to_scope: Dict[str, Tuple[str, type]] = {} |
| for n in model.graph.nodes: |
| nn_module_stack = n.meta.get("nn_module_stack", None) |
| current_scope = ("", type(None)) |
| if nn_module_stack: |
| bt = list(nn_module_stack.values())[-1] |
| current_scope = (bt[0].split(".")[-1], bt[1]) |
| node_name_to_scope[n.name] = current_scope |
| return node_name_to_scope |
| |
| def get_aten_graph_module( |
| pattern: Callable, |
| example_inputs: Tuple[Any, ...], |
| **kwargs, |
| ) -> GraphModule: |
| """ |
| Convert the pattern to an FX graph with decomposed aten ops. |
| """ |
| # Avoid circular imports |
| import torch._dynamo |
| aten_pattern, _ = torch._dynamo.export( |
| pattern, |
| *copy.deepcopy(example_inputs), |
| aten_graph=True, |
| tracing_mode="real", |
| **kwargs, |
| ) |
| aten_pattern.graph.eliminate_dead_code() |
| aten_pattern.recompile() |
| return aten_pattern |
| |
| def remove_tensor_overload_for_qdq_ops(match_pattern: GraphModule) -> None: |
| """ Remove .tensor overload for quantize/dequantize ops so that we can |
| use the match_pattern that we get from torchdynamo export to match the output of convert_pt2e |
| """ |
| _MAP = { |
| torch.ops.quantized_decomposed.quantize_per_tensor.default: torch.ops.quantized_decomposed.quantize_per_tensor, |
| torch.ops.quantized_decomposed.dequantize_per_tensor.default: torch.ops.quantized_decomposed.dequantize_per_tensor, |
| torch.ops.quantized_decomposed.quantize_per_tensor.tensor: torch.ops.quantized_decomposed.quantize_per_tensor, |
| torch.ops.quantized_decomposed.dequantize_per_tensor.tensor: torch.ops.quantized_decomposed.dequantize_per_tensor, |
| torch.ops.quantized_decomposed.quantize_per_tensor.tensor2: torch.ops.quantized_decomposed.quantize_per_tensor, |
| torch.ops.quantized_decomposed.dequantize_per_tensor.tensor2: torch.ops.quantized_decomposed.dequantize_per_tensor, |
| torch.ops.quantized_decomposed.quantize_per_channel.default: torch.ops.quantized_decomposed.quantize_per_channel, |
| torch.ops.quantized_decomposed.dequantize_per_channel.default: torch.ops.quantized_decomposed.dequantize_per_channel, |
| torch.ops.aten.clamp.Tensor: torch.ops.aten.clamp, |
| } |
| for n in match_pattern.graph.nodes: |
| if n.op != "call_function": |
| continue |
| if n.target in _MAP: |
| n.target = _MAP[n.target] |
| |
| def _is_dropout_filter( |
| match: "InternalMatch", # type: ignore[name-defined] |
| original_graph: Graph, |
| pattern_graph: Graph, |
| ) -> bool: |
| """ |
| Match filter for the subgraph rewriter that returns True if the matched |
| graph includes all the ops used in the aten dropout pattern. |
| """ |
| ops_to_match = { |
| torch.ops.aten.empty_like.default, |
| torch.ops.aten.bernoulli_.float, |
| torch.ops.aten.div_.Scalar, |
| torch.ops.aten.mul.Tensor, |
| } |
| for n in match.nodes_map.values(): |
| if n.target in ops_to_match: |
| ops_to_match.remove(n.target) |
| return len(ops_to_match) == 0 |
| |
| def _replace_dropout_for_eval(m: GraphModule): |
| """ |
| Replace the aten training dropout pattern with a noop, intended for eval. |
| |
| For models with dropout torch ops (nn.Dropout, F.dropout), calling model.eval() |
| effectively turns these dropout ops into noops. For exported models, however, |
| this is not done automatically, since the aten dropout patterns previously generated |
| for training remain in the graph. Here we rewrite these dropout patterns with noops |
| to avoid incorrectly applying further dropout during eval. |
| |
| See https://github.com/pytorch/pytorch/issues/103681. |
| """ |
| def dropout_train(x): |
| return F.dropout(x, p=0.5, training=True) |
| |
| def dropout_eval(x): |
| return F.dropout(x, p=0.5, training=False) |
| |
| example_inputs = (torch.randn(1),) |
| match_pattern = get_aten_graph_module(dropout_train, example_inputs) |
| replacement_pattern = get_aten_graph_module(dropout_eval, example_inputs) |
| |
| # Note: The match pattern looks like: |
| # |
| # empty_like_default = torch.ops.aten.empty_like.default(x) |
| # bernoulli__float = torch.ops.aten.bernoulli_.float(empty_like_default) |
| # div__scalar = torch.ops.aten.div_.Scalar(bernoulli__float, 0.5) |
| # mul_tensor = torch.ops.aten.mul.Tensor(x, div__scalar) |
| # |
| # We need to use `ignore_literals=True` here to handle arbitrary dropout |
| # probability (not just 0.5). However, without a match filter, this would |
| # also match any mul op, since `div__scalar` is also a literal, e.g.: |
| # |
| # mul_tensor = torch.ops.aten.mul.Tensor(x, 0.8) |
| # |
| # Therefore, we need both `ignore_literals=True` and `_is_dropout_filter` |
| # to make sure we are in fact replacing the dropout pattern. |
| |
| replace_pattern_with_filters( |
| m, |
| match_pattern, |
| replacement_pattern, |
| match_filters=[_is_dropout_filter], |
| ignore_literals=True, |
| ) |
| m.recompile() |
