[Intel GPU][pt2e] Enable quantized grouped convolution at XPU (#148522)

ZhiweiYan-96 · pytorchmergebot · commit 00cd6c07b98e · 2025-03-07T01:57:45.000Z
# Motivation&Details This PR fix a bug that blocked quantized group convolution before. The bug is caused by that, grouped convolution requires setting weight scale mask on both group dimension and output channel dimension. This PR fixs the wrong mask in integration and add grouped conv in UT. # UT ` python test/inductor/test_mkldnn_pattern_matcher.py -k test_qconv2d_xpu` # Runtime exemplification ```onednn_verbose,v1,primitive,exec,gpu:0,convolution,jit:ir,forward_training,src:s8::blocked:acdb::f0 wei:s8::blocked:abcde::f0 bia:f32::blocked:a::f0 dst:f32::blocked:acdb::f0,attr-scratchpad:user attr-scales:src0:0:f32+dst:0:f32+wei:3:f32 attr-zero-points:src0:0:s32,alg:convolution_direct,g4mb1_ic128oc128_ih4oh2kh3sh1dh0ph0_iw4ow2kw3sw1dw0pw0,0.0529785`` The verbose shows that we successfully run into quantized convolution, where weight is `abcde` format(group conv). Pull Request resolved: #148522 Approved by: https://github.com/EikanWang, https://github.com/liangan1, https://github.com/jansel ghstack dependencies: #148423
diff --git a/aten/src/ATen/native/mkldnn/xpu/detail/QConv.cpp b/aten/src/ATen/native/mkldnn/xpu/detail/QConv.cpp
@@ -131,6 +131,8 @@ at::Tensor quantized_convolution(
   // quant, aka mask=0. Per-channel quantization on activation is not
   // supported in conv.
   mask_weight = weight_zero_points.numel() > 1 ? 1 : 0;
+  if (groups > 1 && weight_zero_points.numel() > 1)
+    mask_weight = (2 ^ 0) | (2 ^ 1); // 2^0 (group) | 2^1 (output channel)
   dnnl::primitive_attr pattr;
 
   bool src_need_zp = (act_scale != 0);
diff --git a/test/inductor/test_mkldnn_pattern_matcher.py b/test/inductor/test_mkldnn_pattern_matcher.py
@@ -988,9 +988,12 @@ def __init__(
                 super().__init__()
                 self.conv = torch.nn.Conv2d(3, 128, kernel_size=3, stride=1)
                 self.conv2 = torch.nn.Conv2d(128, 128, kernel_size=3, stride=1)
+                self.conv3 = torch.nn.Conv2d(
+                    128, 128, kernel_size=3, stride=1, groups=4
+                )
 
             def forward(self, x):
-                return self.conv2(self.conv(x))
+                return self.conv3(self.conv2(self.conv(x)))
 
         mod = M().eval().to(device=device)
         v = (
@@ -1005,14 +1008,14 @@ def matcher_check_fn():
             #    int8_mixed_bf16: [dequant_node, optional(convert_element_type_4),
             #     dequantize_per_channel, optional(convert_element_type_3), clone, convolution]
             self.assertEqual(
-                counters["inductor"]["qconv2d_weight_prepack_matcher_count"], 2
+                counters["inductor"]["qconv2d_weight_prepack_matcher_count"], 3
             )
             self.assertEqual(
                 counters["inductor"]["qconv2d_weight_prepack_matcher_nodes"],
-                12 if int8_mixed_bf16 else 8,
+                18 if int8_mixed_bf16 else 12,
             )
             self.assertEqual(
-                counters["inductor"]["qconv2d_unary_lower_count"], 0 if TEST_ACL else 2
+                counters["inductor"]["qconv2d_unary_lower_count"], 0 if TEST_ACL else 3
             )
 
         self._test_common(
diff --git a/torch/_inductor/fx_passes/quantization.py b/torch/_inductor/fx_passes/quantization.py
@@ -1387,10 +1387,8 @@ def _inner(match):
                 meta_value is None
                 or (meta_value.device.type != "cpu" and meta_value.device.type != "xpu")
                 or meta_value.dim() != 4
-                or (meta_value.device.type == "xpu" and match.kwargs["groups"] != 1)
             ):
                 # Only support conv2d now
-                # Grouped quantized convolution is not supported at XPU backend
                 return False
 
         assert dtype in [torch.float32, torch.bfloat16]
