Inductor cpp wrapper: fix QMaxPool (#112379)

chunyuan-w · pytorchmergebot · commit 46a34e8c755f · 2023-11-06T02:07:51.000Z
Based on the `Argument types` section in this [file](https://github.com/pytorch/pytorch/tree/cb942ef2b12134bfaa1727295380fe00ebb537c0/aten/src/ATen/native#func), for non-inplace `Tensor` type in schema, it should be mapped to C++ argument of type `const Tensor&`. For `quantized_max_pool1d` and `quantized_max_pool2d`, the type of the `qx` input is `Tensor` type in the schema, thus modified the C++ type to be `const Tensor&`: https://github.com/pytorch/pytorch/blob/cb942ef2b12134bfaa1727295380fe00ebb537c0/aten/src/ATen/native/quantized/library.cpp#L222-L223 Pull Request resolved: #112379 Approved by: https://github.com/jgong5, https://github.com/jansel ghstack dependencies: #112373, #112378
diff --git a/aten/src/ATen/native/quantized/cpu/Pooling.cpp b/aten/src/ATen/native/quantized/cpu/Pooling.cpp
@@ -726,7 +726,7 @@ template <uint32_t kSpatialDim>
 class QMaxPool_arr_args final {
  public:
   static Tensor run(
-      Tensor qx,
+      const Tensor& qx,
       std::vector<int64_t> kernel_size,
       std::vector<int64_t> stride,
       std::vector<int64_t> padding,
diff --git a/aten/src/ATen/native/quantized/cudnn/Pooling.cpp b/aten/src/ATen/native/quantized/cudnn/Pooling.cpp
@@ -227,7 +227,7 @@ template <uint32_t kSpatialDim>
 class QMaxPool_arr_args final {
  public:
   static Tensor run(
-      Tensor qx,
+      const Tensor& qx,
       std::vector<int64_t> kernel_size,
       std::vector<int64_t> stride,
       std::vector<int64_t> padding,
diff --git a/test/inductor/test_cpp_wrapper.py b/test/inductor/test_cpp_wrapper.py
@@ -224,6 +224,19 @@ class BaseTest(NamedTuple):
             test_mkldnn_pattern_matcher.TestPatternMatcher(),
             condition=torch.backends.mkldnn.is_available(),
         ),
+        BaseTest(
+            "test_qconv2d_maxpool2d_linear_dynamic",
+            "cpu",
+            test_mkldnn_pattern_matcher.TestDynamicPatternMatcher(),
+            condition=torch.backends.mkldnn.is_available(),
+            func_inputs=[
+                [
+                    "op_qconv2d_pointwise.call",
+                    "op_quantized_max_pool2d_.call",
+                    "op_qlinear_pointwise.call",
+                ]
+            ],
+        ),
         BaseTest(
             "test_qlinear",
             "cpu",
diff --git a/test/inductor/test_mkldnn_pattern_matcher.py b/test/inductor/test_mkldnn_pattern_matcher.py
@@ -1415,7 +1415,7 @@ def forward(self, x):
             match_nodes = 12
             self._test_common(mod, (v,), match_count, match_nodes, rtol=1e-2, atol=1e-2)
 
-    def test_qconv2d_maxpool2d_linear_dynamic(self):
+    def test_qconv2d_maxpool2d_linear_dynamic_cpu(self, include_ops=None):
         r"""
         This testcase will quantize a single Conv2d->Maxpool2d->Linear module
         with dynamic batch size input.
@@ -1444,11 +1444,12 @@ def forward(self, x):
 
         mod = M().eval()
         v = torch.randn((2, 3, 8, 8), dtype=torch.float32, requires_grad=False).add(1)
-        include_ops = [
-            "torch.ops.onednn.qconv2d_pointwise",
-            "torch.ops.quantized.max_pool2d",
-            "torch.ops.onednn.qlinear_pointwise",
-        ]
+        if include_ops is None:
+            include_ops = [
+                "torch.ops.onednn.qconv2d_pointwise",
+                "torch.ops.quantized.max_pool2d",
+                "torch.ops.onednn.qlinear_pointwise",
+            ]
         exclude_ops = []
         self._test_code_common(
             mod,
