[quant][graphmode] Refactor fusion to use the new Pattern format (pytorch#68770)

Summary:
Pull Request resolved: pytorch#68770

Previous fusion only works for a sequnce of ops, which is not general enough for fusion patterns
that is defined by a subgraph, this PR refactors that to make it more general

Test Plan:
```
python test/test_quantization.py TestFuseFx
```

Imported from OSS

Reviewed By: vkuzo

Differential Revision: D32602637

fbshipit-source-id: a7897c62081b9d71c67fb56e78484cf68deaacf6

Author: jerryzh168

torch/ao/quantization/fuser_method_mappings.py

@@ -2,9 +2,11 @@
 import torch.nn.intrinsic as nni
 
 from typing import Union, Callable, Tuple, Dict, Optional, Type
+from torch.ao.quantization.utils import Pattern
 
 from torch.ao.quantization.utils import get_combined_dict
 
+
 def fuse_conv_bn(conv, bn):
     r"""Given the conv and bn modules, fuses them and returns the fused module
 
@@ -104,7 +106,7 @@ def fuse_linear_bn(linear, bn):
     else:
         return nn.utils.fusion.fuse_linear_bn_eval(linear, bn)
 
-DEFAULT_OP_LIST_TO_FUSER_METHOD : Dict[Tuple, Union[nn.Sequential, Callable]] = {
+DEFAULT_OP_LIST_TO_FUSER_METHOD: Dict[Tuple, Union[nn.Sequential, Callable]] = {
     (nn.Conv1d, nn.BatchNorm1d): fuse_conv_bn,
     (nn.Conv1d, nn.BatchNorm1d, nn.ReLU): fuse_conv_bn_relu,
     (nn.Conv2d, nn.BatchNorm2d): fuse_conv_bn,
@@ -131,3 +133,39 @@ def get_fuser_method(op_list, additional_fuser_method_mapping=None):
     fuser_method = all_mappings.get(op_list, None)
     assert fuser_method is not None, "did not find fuser method for: {} ".format(op_list)
     return fuser_method
+
+def reverse2(f):
+    return lambda x, y: f(y, x)
+
+def reverse3(f):
+    def reversed(x, w):
+        y, z = w
+        return f(z, y, x)
+    return reversed
+
+DEFAULT_PATTERN_TO_FUSER_METHOD: Dict[Pattern, Union[nn.Sequential, Callable]] = {
+    (nn.BatchNorm1d, nn.Conv1d): reverse2(fuse_conv_bn),
+    (nn.ReLU, (nn.BatchNorm1d, nn.Conv1d)): reverse3(fuse_conv_bn_relu),
+    (nn.BatchNorm2d, nn.Conv2d): reverse2(fuse_conv_bn),
+    (nn.ReLU, (nn.BatchNorm2d, nn.Conv2d)): reverse3(fuse_conv_bn_relu),
+    (nn.BatchNorm3d, nn.Conv2d): reverse2(fuse_conv_bn),
+    (nn.ReLU, (nn.BatchNorm3d, nn.Conv3d)): reverse3(fuse_conv_bn_relu),
+    (nn.ReLU, nn.Conv1d): reverse2(nni.ConvReLU1d),
+    (nn.ReLU, nn.Conv2d): reverse2(nni.ConvReLU2d),
+    (nn.ReLU, nn.Conv3d): reverse2(nni.ConvReLU3d),
+    (nn.BatchNorm1d, nn.Linear): reverse2(fuse_linear_bn),
+    (nn.ReLU, nn.Linear): reverse2(nni.LinearReLU),
+    (nn.ReLU, nn.BatchNorm2d): reverse2(nni.BNReLU2d),
+    (nn.ReLU, nn.BatchNorm3d): reverse2(nni.BNReLU3d),
+}
+
+def get_fuser_method_new(op_pattern, fuser_method_mapping=None):
+    """ This will be made defult after we deparate the get_fuser_method
+    Would like to implement this first and have a separate PR for deprecation
+    """
+    if fuser_method_mapping is None:
+        fuser_method_mapping = DEFAULT_PATTERN_TO_FUSER_METHOD
+
+    fuser_method = fuser_method_mapping.get(op_pattern, None)
+    assert fuser_method is not None, "did not find fuser method for: {} ".format(op_pattern)
+    return fuser_method

torch/ao/quantization/fx/fuse.py

@@ -1,5 +1,3 @@
-from typing import Dict, Any
-
 from torch.fx import (
     GraphModule,
     Node,
@@ -19,11 +17,9 @@
 
 from .fusion_patterns import *  # noqa: F401,F403
 
-from typing import Callable, Tuple
-from typing import Optional
-
-from .quantization_types import Pattern
+from typing import Callable, Tuple, Dict, Any, Optional, List
 
+from .quantization_types import Pattern, NodePattern
 
 class Fuser:
     def fuse(
@@ -50,11 +46,18 @@ def load_arg(a):
             return map_arg(a, lambda node: env[node.name])
 
         for node in input_graph.nodes:
-            root_node, obj = fusion_pairs.get(node.name, (None, None))
-            if root_node is node:
+            maybe_last_node, pattern, matched_node_pattern, obj = \
+                fusion_pairs.get(node.name, (None, None, None, None))
+            if maybe_last_node is node:
                 assert obj is not None
-                env[node.name] = obj.fuse(self, load_arg, fuse_custom_config_dict)
-            elif root_node is None:
+                # TODO: currently we hard code the root node, which only works for
+                # a tuple of two nodes, we want to make this more general to
+                # support more complex patterns
+                root_node = matched_node_pattern[-1]  # type: ignore[index]
+                env[node.name] = obj.fuse(
+                    self, load_arg, root_node, matched_node_pattern,  # type: ignore[arg-type]
+                    fuse_custom_config_dict)
+            elif maybe_last_node is None:
                 env[node.name] = self.fused_graph.node_copy(node, load_arg)
             # node matched in patterns and is not root is removed here
 
@@ -65,25 +68,30 @@ def load_arg(a):
     def _find_matches(
             self, root: GraphModule, graph: Graph,
             patterns: Dict[Pattern, Callable]
-    ) -> Dict[str, Tuple[Node, FuseHandler]]:
+    ) -> Dict[str, Tuple[Node, Pattern, NodePattern, FuseHandler]]:
         modules = dict(root.named_modules())
-        match_map : Dict[str, Tuple[Node, FuseHandler]] = {}  # node name -> (root_node, match_value)
+        match_map : Dict[str, Tuple[Node, Pattern, NodePattern, FuseHandler]] = {}  # node name -> (root_node, match_value)
 
-        def apply_match(pattern, node, match):
+        def apply_match(pattern, node, match, matched_node_pattern):
             if isinstance(pattern, tuple):
                 s, *args = pattern
-                apply_match(s, node, match)
+                current_node_pattern: List[Node] = []
+                apply_match(s, node, match, current_node_pattern)
                 for subpattern, arg in zip(args, node.args):
-                    apply_match(subpattern, arg, match)
+                    apply_match(subpattern, arg, match, current_node_pattern)
+                matched_node_pattern.append(tuple(current_node_pattern))
             else:
                 # the first pattern matches will take precedence
                 if node.name not in match_map:
-                    match_map[node.name] = match
+                    matched_node_pattern.append(node)
+                    root_node, pattern, handler = match
+                    match_map[node.name] = (root_node, pattern, matched_node_pattern, handler)
 
         for node in reversed(graph.nodes):
             if node.name not in match_map:
                 for pattern, value in patterns.items():
+                    matched_node_pattern: List[Node] = []
                     if is_match(modules, node, pattern):
-                        apply_match(pattern, node, (node, value(self, node)))
+                        apply_match(pattern, node, (node, pattern, value(self, node)), matched_node_pattern)
 
         return match_map

torch/ao/quantization/fx/fusion_patterns.py

@@ -4,10 +4,12 @@
     register_fusion_pattern,
 )
 from .utils import _parent_name
-from .quantization_types import QuantizerCls
+from .quantization_types import QuantizerCls, NodePattern
 from ..fuser_method_mappings import get_fuser_method
+from ..fuser_method_mappings import get_fuser_method_new
 from abc import ABC, abstractmethod
 from typing import Any, Callable, Dict
+from .match_utils import MatchAllNode
 
 # ---------------------
 # Fusion Pattern Registrations
@@ -21,7 +23,11 @@ def __init__(self, quantizer: QuantizerCls, node: Node):
         pass
 
     @abstractmethod
-    def fuse(self, quantizer: QuantizerCls, load_arg: Callable,
+    def fuse(self,
+             quantizer: QuantizerCls,
+             load_arg: Callable,
+             root_node: Node,
+             matched_node_pattern: NodePattern,
              fuse_custom_config_dict: Dict[str, Any]) -> Node:
         pass
 
@@ -61,7 +67,11 @@ def __init__(self, quantizer: QuantizerCls, node: Node):
         self.conv_or_linear_node = node
         self.conv_or_linear = quantizer.modules[self.conv_or_linear_node.target]
 
-    def fuse(self, quantizer: QuantizerCls, load_arg: Callable,
+    def fuse(self,
+             quantizer: QuantizerCls,
+             load_arg: Callable,
+             root_node: Node,
+             matched_node_pattern: NodePattern,
              fuse_custom_config_dict: Dict[str, Any]) -> Node:
         additional_fuser_method_mapping = fuse_custom_config_dict.get("additional_fuser_method_mapping", {})
         op_list = []
@@ -116,23 +126,36 @@ def __init__(self, quantizer: QuantizerCls, node: Node):
         self.module_node = node
         self.module = quantizer.modules[self.module_node.target]
 
-    def fuse(self, quantizer: QuantizerCls, load_arg: Callable,
+    def fuse(self, quantizer: QuantizerCls,
+             load_arg: Callable,
+             root_node: Node,
+             matched_node_pattern: NodePattern,
              fuse_custom_config_dict: Dict[str, Any]) -> Node:
         additional_fuser_method_mapping = fuse_custom_config_dict.get("additional_fuser_method_mapping", {})
-        op_list = []
+        assert root_node.op == "call_module", "Expecting module node to be a call_module Node"
+        root_module = quantizer.modules[root_node.target]
+        assert len(additional_fuser_method_mapping) == 0, "Fusion implementation is "
+        "undergoing changes, additoinal_fuser_method_mapping is not supported currently."
+        def get_module(n):
+            if n.op == "call_module":
+                return quantizer.modules[n.target]
+            elif n.op == "call_function" and n.target == torch.nn.functional.relu:
+                relu = torch.nn.ReLU()
+                relu.training = root_module.training
+                return relu
+            return MatchAllNode
+
+        matched_modules = tuple(map(get_module, matched_node_pattern))
         # since relu can be used multiple times, we'll need to create a relu module for each match
-        if self.relu_node.op == 'call_module':
-            relu = torch.nn.ReLU(quantizer.modules[self.relu_node.target].inplace)
-        else:
-            # TODO: get inplace argument from functional
-            relu = torch.nn.ReLU()
-        relu.training = self.module.training
-        op_list.append(relu)
-        op_list.append(self.module)
 
-        op_list.reverse()
-        op_type_list = tuple(type(m) for m in op_list)
-        module_parent_name, module_name = _parent_name(self.module_node.target)
-        fuser_method = get_fuser_method(op_type_list, additional_fuser_method_mapping)
-        setattr(quantizer.modules[module_parent_name], module_name, fuser_method(*op_list))
-        return quantizer.fused_graph.node_copy(self.module_node, load_arg)
+        def get_type(m):
+            return type(m)
+
+        matched_module_types = tuple(map(get_type, matched_modules))
+        module_parent_name, module_name = _parent_name(root_node.target)
+        fuser_method = get_fuser_method_new(matched_module_types)
+        # TODO: change the signature for fuser_method to take matched module patterns
+        # as input
+        fused_module = fuser_method(*matched_modules)
+        setattr(quantizer.modules[module_parent_name], module_name, fused_module)
+        return quantizer.fused_graph.node_copy(root_node, load_arg)

torch/ao/quantization/fx/match_utils.py

@@ -22,6 +22,7 @@
 MatchResult = Tuple[Node, List[Node], Optional[Pattern], QuantizeHandler,
                     QConfigAny]
 
+# TODO: maybe rename this to MatchInputNode
 class MatchAllNode:
     """ A node pattern that matches all nodes
     """

torch/ao/quantization/fx/quantization_types.py

@@ -1,6 +1,9 @@
-from typing import Any
+from typing import Any, Tuple, Union
+from torch.fx import Node
 from ..utils import Pattern  # noqa: F401
 
+NodePattern = Union[Tuple[Node, Node], Tuple[Node, Tuple[Node, Node]], Any]
+
 # This is the Quantizer class instance from torch/quantization/fx/quantize.py.
 # Define separately to prevent circular imports.
 # TODO(future PR): improve this.

torch/ao/quantization/utils.py

@@ -10,7 +10,7 @@
 # Type for fusion patterns, it can be more complicated than the following actually,
 # see pattern.md for docs
 # TODO: not sure if typing supports recursive data types
-Pattern = Union[Callable, Tuple[Callable, Callable], Tuple[Callable, Callable, Callable]]
+Pattern = Union[Callable, Tuple[Callable, Callable], Tuple[Callable, Tuple[Callable, Callable]], Any]
 
 def get_combined_dict(default_dict, additional_dict):
     d = default_dict.copy()