Update

[ghstack-poisoned]
pytorch · ZhiweiYan-96 · Sep 26, 2024 · Oct 9, 2024 · Oct 9, 2024 · Oct 17, 2024
commit dae8179f3cc84c1c1f930b1ab3523dcf8ac13a65
diff --git a/test/inductor/test_mkldnn_pattern_matcher.py b/test/inductor/test_mkldnn_pattern_matcher.py
@@ -1705,7 +1705,7 @@ def test_qlinear_gelu_int8_mixed_bf16(self, device="cpu"):
                 (torch.randn((2, 4)).to(device=device),), gelu, device=device, int8_mixed_bf16=True
             )
 
-    def _qlinear_add_cpu_test_helper(self, device="cpu", use_relu=False, int8_mixed_bf16=False):
+    def _qlinear_add_cpu_test_helper(self, device="cpu", use_relu=False, int8_mixed_bf16=False, is_qat=True, is_dynamic=True):
         r"""
         This testcase will quantize two consecutive Linear->Add(->relu) patterns as:
                  X
@@ -1809,14 +1809,18 @@ def matcher_check_fn():
                 #     (4 if is_dynamic else 5) + 2 * use_relu + to_bf16_after_binary,
                 # )
 
-            # is_qat_list = [False, True]
-            is_qat_list = [False]
-            # is_dynamic_list = [False, True]
-            is_dynamic_list = [False]
-            cases = itertools.product(is_qat_list, is_dynamic_list)
-            for is_qat, is_dynamic in cases:
-                print("int8_mixed_bf16:", int8_mixed_bf16)
-                self._test_common(
+            self._test_common(
+                mod,
+                    (v,),
+                check_quantization=True,
+                check_autocast=torch.bfloat16 if int8_mixed_bf16 else torch.float,
+                matcher_check_fn=matcher_check_fn,
+                is_qat=is_qat,
+                is_dynamic=is_dynamic,
+            )
+            if torch._inductor.config.cpp_wrapper:
+                # For CPP wrapper
+                self._test_code_common(
                     mod,
                     (v,),
                     [
@@ -1845,54 +1849,31 @@ def matcher_check_fn():
                     check_quantization=True,
                     num_include_ops=[2, 2],
                 )
-                # if torch._inductor.config.cpp_wrapper:
-                #     # For CPP wrapper
-                #     self._test_code_common(
-                #         mod,
-                #         (v,),
-                #         [
-                #             "op_qlinear_pointwise.call",
-                #             "op_qlinear_pointwise_binary.call",
-                #         ],
-                #         [],
-                #         check_quantization=True,
-                #         num_include_ops=[2, 2],
-                #     )
-                # else:
-                #     # For python wrapper
-                #     self._test_code_common(
-                #         mod,
-                #         (v,),
-                #         [
-                #             "torch.ops.onednn.qlinear_pointwise.tensor",
-                #             "torch.ops.onednn.qlinear_pointwise.binary",
-                #         ],
-                #         [],
-                #         check_quantization=True,
-                #         num_include_ops=[2, 2],
-                #     )
 
     @skipIfNoDynamoSupport
     @skipIfNoONEDNN
-    def test_qlinear_add_mkldnn(self, device="cpu"):
-        self._qlinear_add_cpu_test_helper(device=device)
-
-    @skipIfNoDynamoSupport
-    @skipIfNoONEDNNBF16
-    @skipIfNoONEDNN
-    def test_qlinear_add_int8_mixed_bf16(self, device="cpu"):
-        self._qlinear_add_cpu_test_helper(device=device, int8_mixed_bf16=True)
-
-    @skipIfNoDynamoSupport
-    @skipIfNoONEDNN
-    def test_qlinear_add_relu(self, device="cpu"):
-        self._qlinear_add_cpu_test_helper(device=device, use_relu=True)
+    @parametrize("use_relu", [True, False])
+    @parametrize("is_qat", [True, False])
+    @parametrize("is_dynamic", [True, False])
+    def test_qlinear_add_cpu(self, use_relu, is_qat, is_dynamic):
+        self._qlinear_add_cpu_test_helper(
+            use_relu=use_relu, is_qat=is_qat, is_dynamic=is_dynamic
+        )
 
     @skipIfNoDynamoSupport
     @skipIfNoONEDNNBF16
     @skipIfNoONEDNN
-    def test_qlinear_add_relu_int8_mixed_bf16(self, device="cpu"):
-        self._qlinear_add_cpu_test_helper(device=device, use_relu=True, int8_mixed_bf16=True)
+    @parametrize("use_relu", [True, False])
+    @parametrize("is_qat", [True, False])
+    @parametrize("is_dynamic", [True, False])
+    def test_qlinear_add_int8_mixed_bf16(self, device, use_relu, is_qat, is_dynamic):
+        self._qlinear_add_cpu_test_helper(
+            device=device,
+            int8_mixed_bf16=True,
+            use_relu=use_relu,
+            is_qat=is_qat,
+            is_dynamic=is_dynamic,
+        )
 
     def _qlinear_dequant_promotion_cpu_test_helper(
-    def _qlinear_dequant_promotion_cpu_test_helper(
+    def _qlinear_dequant_promotion_test_helper(
-    def _qlinear_dequant_promotion_cpu_test_helper(
+    def _qlinear_dequant_promotion_test_helper(
         self,