[quant][fx][graphmode] Add support for conv add pattern in backend_config_dict (pytorch#69778)

Summary:
Pull Request resolved: pytorch#69778

This PR extends fusion pattern support from simple sequence of ops to a simple
subgraph like conv - add
```
x - conv ---\
y ---------add ---- ouptut
```
where input x, y and output are observed/quantized

Test Plan:
```
python test/fx2trt/test_quant_trt.py TestQuantizeFxTRTOps.test_conv_add
```

Imported from OSS

Reviewed By: vkuzo

Differential Revision: D33024528

fbshipit-source-id: 5c770c82c8f693fabdac5c69343942a9dfda84ef

Author: jerryzh168

test/fx2trt/test_quant_trt.py

@@ -17,6 +17,9 @@
 from torch.ao.quantization._quantize_fx_do_not_use import (
     _convert_fx_do_not_use,
 )
+from torch.ao.quantization.fx.match_utils import (
+    MatchAllNode,
+)
 from torch.testing._internal.common_quantization import (
     QuantizationTestCase,
 )
@@ -27,6 +30,8 @@
 from torch.testing._internal.common_quantization import NodeSpec as ns
 import unittest
 import itertools
+import copy
+import operator
 
 def lower_to_trt(model, inputs, shape_ranges):
     """ Lower a quantized model to TensorRT
@@ -375,5 +380,51 @@ def forward(self, x):
         }
         self.checkGraphModuleNodes(quantized, expected_node_occurrence=node_occurrence)
 
+    def test_conv_add(self):
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv = torch.nn.Conv2d(3, 3, 3)
+
+            def forward(self, x, y):
+                return self.conv(x) + y
+
+        from torch.ao.quantization.fx.backend_config_dict.observation_type import ObservationType
+        weighted_op_qint8_dtype_config = {
+            # optional, input activation dtype
+            "input_dtype": torch.qint8,
+            # optional, weight dtype
+            "weight_dtype": torch.qint8,
+            # optional, bias dtype
+            "bias_dtype": torch.float,
+            # optional, output activation dtype
+            "output_dtype": torch.qint8
+        }
+
+        conv_add_config = {
+            "pattern": (operator.add, torch.nn.Conv2d, MatchAllNode),
+            "observation_type": ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT,
+            "dtype_configs": [
+                weighted_op_qint8_dtype_config,
+            ],
+            "root_module": torch.nn.Conv2d,
+            "reference_quantized_module_for_root": torch.nn.quantized._reference.Conv2d,
+        }
+
+        m = M().eval()
+        modified_backend_config_dict = copy.deepcopy(self.trt_backend_config_dict)
+        modified_backend_config_dict["configs"].insert(0, conv_add_config)
+        m = prepare_fx(m, {"": self.qconfig}, backend_config_dict=modified_backend_config_dict)
+        node_occurrence = {
+            ns.call_module(torch.ao.quantization.HistogramObserver): 3,
+        }
+        self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
+        m = _convert_fx_do_not_use(m, is_reference=True, backend_config_dict=modified_backend_config_dict)
+        node_occurrence = {
+            ns.call_function(torch.quantize_per_tensor): 3,
+            ns.call_method("dequantize"): 3,
+        }
+        self.checkGraphModuleNodes(m, expected_node_occurrence=node_occurrence)
+
 if __name__ == "__main__":
     run_tests()

torch/ao/quantization/fx/_convert_do_not_use.py

@@ -15,9 +15,6 @@
 )
 from .backend_config_dict.utils import get_quantized_reference_module_mapping
 
-from .match_utils import (
-    find_matches,
-)
 from .graph_module import (
     QuantizedGraphModule,
 )
@@ -103,10 +100,6 @@ def _convert_do_not_use(
     custom_module_classes = get_custom_module_class_keys(
         convert_custom_config_dict,
         "observed_to_quantized_custom_module_class")
-    matches = find_matches(
-        model.graph, modules, patterns,
-        qconfig_map,
-        custom_module_classes=custom_module_classes)
 
     if model._equalization_qconfig_map is not None:
         # If we want to do equalization then do the following:

torch/ao/quantization/fx/backend_config_dict/fuse_handler.py

@@ -1,5 +1,5 @@
 from ..fusion_patterns import DefaultFuseHandler
 
-# TODO: move ModuleReLUFusion here
+# TODO: move DefaultFuseHandler
 def get_fuse_handler_cls():
     return DefaultFuseHandler

torch/ao/quantization/fx/prepare.py

@@ -174,8 +174,10 @@ def is_pattern_dtype_config_supported_by_backend(
     pattern_to_dtype_configs = get_pattern_to_dtype_configs(backend_config_dict)
     dtype_configs: List[Dict[str, torch.dtype]] = pattern_to_dtype_configs.get(pattern, [])
 
-    input_node = matched_nodes[0]
-    output_node = matched_nodes[-1]
+    # TODO: this only checks one input and one output, need to generalize to multiple
+    # inputs/output
+    input_node = matched_nodes[-1]
+    output_node = matched_nodes[0]
     for dtype_config in dtype_configs:
         # check if arg dtype are supported
         supported = True