[quant][be] Cleanup xnnpack_quantizer implementation (#108921)

jerryzh168 · pytorchmergebot · commit b01b934aca3c · 2023-09-12T19:28:41.000Z
Summary: att Test Plan: python test/test_quantization.py TestQuantizePT2E Reviewers: Subscribers: Tasks: Tags: Pull Request resolved: #108921 Approved by: https://github.com/andrewor14
diff --git a/torch/ao/quantization/quantizer/xnnpack_quantizer.py b/torch/ao/quantization/quantizer/xnnpack_quantizer.py
@@ -237,6 +237,30 @@ def module_type_filter(n: Node) -> bool:
 
 class XNNPACKQuantizer(Quantizer):
     supported_config_and_operators = _get_supported_config_and_operators()
+    STATIC_QAT_ONLY_OPS = [
+        "conv2d_bn_relu",
+        "conv2d_bn",
+    ]
+
+    # static quantization ops (both PTQ and QAT)
+    STATIC_OPS = [
+        "linear",
+        "conv2d_relu",
+        "conv2d",
+        "adaptive_avg_pool2d",
+        # TODO: move this to BoltNNQuantizer?
+        "gru_io_only",
+        "max_pool2d",
+        "add_relu",
+        "add",
+        "mul_relu",
+        "mul",
+        "cat",
+    ]
+
+    DYNAMIC_OPS = [
+        "linear",
+    ]
 
     def __init__(self):
         super().__init__()
@@ -315,170 +339,69 @@ def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
         propagate_annotation(model)
         return model
 
-    def _annotate_all_patterns(
+    def _annotate_all_static_patterns(
         self,
         model: torch.fx.GraphModule,
-        config: Optional[QuantizationConfig],
+        quantization_config: Optional[QuantizationConfig],
         filter_fn: Optional[Callable[[Node], bool]] = None,
     ) -> torch.fx.GraphModule:
         # TODO: implement the support for None to be canceling out previous annotations
-        if config is None:
+        if quantization_config is None:
             return model
 
-        # assert config is not None
-
-        self._annotate_linear(model, config, filter_fn)
-        self._annotate_conv2d_patterns(model, config, filter_fn)
-        self._annotate_max_pool2d(model, config, filter_fn)
-        self._annotate_add_patterns(model, config, filter_fn)
-        OP_TO_ANNOTATOR["mul_relu"](model, config, filter_fn)
-        OP_TO_ANNOTATOR["mul"](model, config, filter_fn)
-        OP_TO_ANNOTATOR["cat"](model, config, filter_fn)
-        self._annotate_adaptive_avg_pool2d(model, config, filter_fn)
-        self._annotate_gru_io_only(model, config, filter_fn)
+        if quantization_config.is_qat:
+            for op in self.STATIC_QAT_ONLY_OPS:
+                OP_TO_ANNOTATOR[op](model, quantization_config, filter_fn)
+        for op in self.STATIC_OPS:
+            OP_TO_ANNOTATOR[op](model, quantization_config, filter_fn)
+        return model
+
+    def _annotate_all_dynamic_patterns(
+        self,
+        model: torch.fx.GraphModule,
+        quantization_config: Optional[QuantizationConfig],
+        filter_fn: Optional[Callable[[Node], bool]] = None,
+    ) -> torch.fx.GraphModule:
+        # TODO: implement the support for None to be canceling out previous annotations
+        if quantization_config is None:
+            return model
+
+        for op in self.DYNAMIC_OPS:
+            OP_TO_ANNOTATOR[op](model, quantization_config, filter_fn)
         return model
 
     def _annotate_for_static_quantization_config(
         self, model: torch.fx.GraphModule
     ) -> torch.fx.GraphModule:
         for module_name, config in self.module_name_config.items():
-            self._annotate_all_patterns(
+            self._annotate_all_static_patterns(
                 model, config, _get_module_name_filter(module_name)
             )
 
         for module_type, config in self.module_type_config.items():
-            self._annotate_all_patterns(
+            self._annotate_all_static_patterns(
                 model, config, _get_module_type_filter(module_type)
             )
 
-        self._annotate_all_patterns(model, self.global_config)
+        self._annotate_all_static_patterns(model, self.global_config)
         return model
 
     def _annotate_for_dynamic_quantization_config(
         self, model: torch.fx.GraphModule
     ) -> torch.fx.GraphModule:
         for module_name, config in self.module_name_config.items():
-            self._annotate_linear(model, config, _get_module_name_filter(module_name))
+            self._annotate_all_dynamic_patterns(
+                model, config, _get_module_name_filter(module_name)
+            )
 
         for module_type, config in self.module_type_config.items():
-            self._annotate_linear(model, config, _get_module_type_filter(module_type))
+            self._annotate_all_dynamic_patterns(
+                model, config, _get_module_type_filter(module_type)
+            )
 
-        self._annotate_linear(model, self.global_config)
+        self._annotate_all_dynamic_patterns(model, self.global_config)
         return model
 
-    def _annotate_conv2d_patterns(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> None:
-        if quantization_config is None:
-            return
-
-        if quantization_config.is_qat:
-            self._annotate_conv2d_bn_relu(gm, quantization_config, filter_fn)
-            self._annotate_conv2d_bn(gm, quantization_config, filter_fn)
-        self._annotate_conv2d_relu(gm, quantization_config, filter_fn)
-        self._annotate_conv2d(gm, quantization_config, filter_fn)
-
-    def _annotate_conv2d_bn(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        """
-        Note: This is only used for QAT. In PTQ, batchnorm should already be fused into the conv.
-        """
-        return OP_TO_ANNOTATOR["conv2d_bn"](gm, quantization_config, filter_fn)
-
-    def _annotate_conv2d_bn_relu(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        """
-        Note: This is only used for QAT. In PTQ, batchnorm should already be fused into the conv.
-        """
-        return OP_TO_ANNOTATOR["conv2d_bn_relu"](gm, quantization_config, filter_fn)
-
-    def _annotate_conv2d_relu(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        return OP_TO_ANNOTATOR["conv2d_relu"](gm, quantization_config, filter_fn)
-
-    def _annotate_conv2d(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        return OP_TO_ANNOTATOR["conv2d"](gm, quantization_config, filter_fn)
-
-    def _annotate_linear(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        return OP_TO_ANNOTATOR["linear"](gm, quantization_config, filter_fn)
-
-    def _annotate_adaptive_avg_pool2d(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        return OP_TO_ANNOTATOR["adaptive_avg_pool2d"](
-            gm, quantization_config, filter_fn
-        )
-
-    # TODO: move this to BoltNNQuantizer?
-    def _annotate_gru_io_only(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        return OP_TO_ANNOTATOR["gru_io_only"](gm, quantization_config, filter_fn)
-
-    def _annotate_max_pool2d(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        return OP_TO_ANNOTATOR["max_pool2d"](gm, quantization_config, filter_fn)
-
-    def _annotate_add_patterns(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> None:
-        self._annotate_add_relu(gm, quantization_config, filter_fn)
-        self._annotate_add(gm, quantization_config, filter_fn)
-
-    def _annotate_add_relu(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        return OP_TO_ANNOTATOR["add_relu"](gm, quantization_config, filter_fn)
-
-    def _annotate_add(
-        self,
-        gm: torch.fx.GraphModule,
-        quantization_config: Optional[QuantizationConfig],
-        filter_fn: Optional[Callable[[Node], bool]] = None,
-    ) -> Optional[List[List[Node]]]:
-        return OP_TO_ANNOTATOR["add"](gm, quantization_config, filter_fn)
-
     def validate(self, model: torch.fx.GraphModule) -> None:
         pass
 
