[quant][graphmode][fx] Add support for additional_{fusion/quant}_pattern

torch/quantization/fx/fuse.py

@@ -24,7 +24,8 @@ def fuse(self, model, inplace=False, fuse_custom_config_dict=None):
         input_graph = model.graph
         self.modules = dict(input_root.named_modules())
 
-        fusion_patterns = get_default_fusion_patterns()
+        additional_fusion_patterns = fuse_custom_config_dict.get("additional_quant_pattern", {})
+        fusion_patterns = dict(get_default_fusion_patterns(), **additional_fusion_patterns)
         # find fusion
         fusion_pairs = self._find_matches(input_root, input_graph, fusion_patterns)
         self.fused_graph = Graph()

torch/quantization/fx/quantize.py

@@ -325,7 +325,8 @@ def _prepare(self, model, qconfig_dict, inplace, prepare_custom_config_dict, is_
             prepare_custom_config_dict = {}
         if not inplace:
             model = copy.deepcopy(model)
-        self.patterns = get_default_quant_patterns()
+        additional_quant_patterns = prepare_custom_config_dict.get("additional_quant_pattern", {})
+        self.patterns = dict(get_default_quant_patterns, **additional_quant_patterns)
 
         flattened_qconfig_dict = get_flattened_qconfig_dict(qconfig_dict)
         # TODO: support regex as well

torch/quantization/quantize_fx.py

@@ -184,6 +184,17 @@ def prepare_fx(model, qconfig_dict, inplace=False, prepare_custom_config_dict=No
         "additional_qat_module_mapping": {
            FloatModule: QATModule
         },
+
+        # Additional fusion patterns
+        "additioanl_fusion_pattern": {
+           (torch.nn.BatchNorm2d, torch.nn.Conv2d): ConvReluFusionhandler
+        },
+
+        # Additional quantization patterns
+        "additional_quant_pattern": {
+           torch.nn.Conv2d: ConvReluQuantizeHandler,
+           (torch.nn.ReLU, torch.nn.Conv2d): ConvReluQuantizeHandler,
+        }
       }