pytorch · ZhiweiYan-96 · Sep 9, 2024 · Oct 9, 2024 · Oct 17, 2024 · Oct 21, 2024
diff --git a/aten/src/ATen/native/mkldnn/xpu/qconv.cpp b/aten/src/ATen/native/mkldnn/xpu/qconv.cpp
@@ -1,13 +1,24 @@
 #include <ATen/core/op_registration/op_registration.h>
 #include <ATen/native/mkldnn/xpu/detail/oneDNN.h>
 #include <c10/core/MemoryFormat.h>
+#include <c10/core/ScalarType.h>
 #include <torch/library.h>
 
-#include <iostream>
-
 using namespace at::native::onednn;
 namespace at::native::xpu {
 
+static inline c10::ScalarType qconv_decide_out_dtype(
+    const at::Tensor& act,
+    const std::optional<c10::ScalarType> output_dtype) {
+  bool fp32_output = output_dtype.has_value() && (output_dtype == c10::kFloat);
+  bool bfloat16_output =
+      output_dtype.has_value() && (output_dtype == c10::kBFloat16);
+  auto dst_dtype = fp32_output
+      ? c10::kFloat
+      : (bfloat16_output ? c10::kBFloat16 : act.scalar_type());
+  return dst_dtype;
+}
+
 at::Tensor qconv_prepack_xpu(
     at::Tensor weight,
     at::Tensor weight_scales,
@@ -75,8 +86,9 @@ class QConvoneDNNXPU final {
         stride.vec(),
         dilation.vec());
 
-    Tensor output = at::empty(
-        dst_tz, act.options().dtype(output_dtype).memory_format(mfmt));
+    auto dst_dtype = qconv_decide_out_dtype(act, output_dtype);
+    Tensor output =
+        at::empty(dst_tz, act.options().dtype(dst_dtype).memory_format(mfmt));
 
     return quantized_convolution(
         act,
@@ -155,11 +167,11 @@ class QConvoneDNNXPU final {
         stride.vec(),
         dilation.vec());
 
+    auto dst_dtype = qconv_decide_out_dtype(act, output_dtype);
     bool has_accum_postop_sum = binary_attr == "sum";
     Tensor output = has_accum_postop_sum
         ? accum
-        : at::empty(
-              dst_tz, act.options().dtype(output_dtype).memory_format(mfmt));
+        : at::empty(dst_tz, act.options().dtype(dst_dtype).memory_format(mfmt));
 
     output = quantized_convolution(
         act,

diff --git a/test/inductor/test_mkldnn_pattern_matcher.py b/test/inductor/test_mkldnn_pattern_matcher.py
@@ -156,17 +156,24 @@ def _test_common(
     ):
         counters.clear()
         torch._dynamo.reset()
-        if (
-            check_autocast == torch.bfloat16
-            and torch.ops.mkldnn._is_mkldnn_bf16_supported()
+        has_xpu = any(
+            isinstance(input, torch.Tensor) and input.device.type == "xpu"
+            for input in inputs
+        )
+        device_type = "xpu" if has_xpu else "cpu"
+        if check_autocast == torch.bfloat16 and (
+            torch.ops.mkldnn._is_mkldnn_bf16_supported() or has_xpu
         ):
-            maybe_autocast = torch.amp.autocast("cpu", dtype=torch.bfloat16)
+            maybe_autocast = torch.amp.autocast(
+                device_type=device_type, dtype=torch.bfloat16
+            )
             atol, rtol = 1e-2, 1e-2
-        elif (
-            check_autocast == torch.float16
-            and torch.ops.mkldnn._is_mkldnn_fp16_supported()
+        elif check_autocast == torch.float16 and (
+            torch.ops.mkldnn._is_mkldnn_fp16_supported() or has_xpu
         ):
-            maybe_autocast = torch.amp.autocast("cpu", dtype=torch.float16)
+            maybe_autocast = torch.amp.autocast(
+                device_type=device_type, dtype=torch.float16
+            )
             atol, rtol = 1e-2, 1e-2
         else:
             assert check_autocast == torch.float32
@@ -1044,6 +1051,16 @@ def test_qconv2d_int8_mixed_bf16(self):
         """
         self._qconv2d_test_helper(int8_mixed_bf16=True)
 
+    @skipIfNoDynamoSupport
+    @skipIfNoONEDNNBF16
+    @skipIfNoONEDNN
+    @skipIfNoXPU
+    def test_qconv2d_int8_mixed_bf16_xpu(self):
+        r"""
+        This testcase will quantize a single Conv2d module with int8_mixed_bf16 quantization.
+        """
+        self._qconv2d_test_helper(device="xpu", int8_mixed_bf16=True)
+
     def _qconv2d_unary_test_helper(
         self,
         device="cpu",
@@ -1122,7 +1139,7 @@ def test_qconv2d_relu_xpu(self):
     @skipIfNoDynamoSupport
     @skipIfNoONEDNNBF16
     @skipIfNoONEDNN
-    def test_qconv2d_relu_int8_mixed_bf16(self):
+    def test_qconv2d_relu_int8_mixed_bf16_xpu(self):
         r"""
         This testcase will quantize Conv2d->ReLU pattern with int8_mixed_bf16 quantization.
         """
@@ -1178,6 +1195,24 @@ def test_qconv2d_hardtanh_int8_mixed_bf16_cpu(self):
             qconv2d_unary_matcher_nodes=11,
         )
 
+    @skipIfNoDynamoSupport
+    @skipIfNoONEDNNBF16
+    @skipIfNoONEDNN
+    @skipIfNoXPU
+    def test_qconv2d_hardtanh_int8_mixed_bf16_xpu(self):
+        r"""
+        This testcase will quantize Conv2d->Hardtanh pattern.
+        Match.nodes:
+            [qconv2d_pointwise_default, convert_element_type, clamp_min, clamp_max, convert_element_type, quantize_per_tensor]
+            [qconv2d_pointwise_default, convert_element_type, clamp_min, clamp_max, convert_element_type]
+        """
+        self._qconv2d_unary_test_helper(
+            device="xpu",
+            unary_op=torch.nn.Hardtanh(),
+            int8_mixed_bf16=True,
+            qconv2d_unary_matcher_nodes=11,
+        )
+
     @skipIfNoDynamoSupport
     @skipIfNoONEDNN
     def test_qconv2d_hardswish_cpu(self):
@@ -1212,6 +1247,25 @@ def test_qconv2d_hardswish_int8_mixed_bf16_cpu(self):
             qconv2d_unary_matcher_nodes=17,
         )
 
+    @skipIfNoDynamoSupport
+    @skipIfNoONEDNNBF16
+    @skipIfNoONEDNN
+    @skipIfNoXPU
+    def test_qconv2d_hardswish_int8_mixed_bf16_xpu(self):
+        r"""
+        This testcase will quantize Conv2d->Hardswish pattern.
+        Match.nodes:
+            [qconv2d_pointwise_default, convert_element_type, add, clamp_min,
+             clamp_max, mul, div, convert_element_type, quantize_per_tensor]
+            [qconv2d_pointwise_default, convert_element_type, add, clamp_min, clamp_max, mul, div, convert_element_type]
+        """
+        self._qconv2d_unary_test_helper(
+            device="xpu",
+            unary_op=torch.nn.Hardswish(),
+            int8_mixed_bf16=True,
+            qconv2d_unary_matcher_nodes=17,
+        )
+
     @skipIfNoDynamoSupport
     @skipIfNoONEDNN
     def test_qconv2d_silu_cpu(self):
@@ -1246,6 +1300,25 @@ def test_qconv2d_silu_int8_mixed_bf16_cpu(self):
             qconv2d_unary_matcher_nodes=11,
         )
 
+    @skipIfNoDynamoSupport
+    @skipIfNoONEDNNBF16
+    @skipIfNoONEDNN
+    @skipIfNoXPU
+    def test_qconv2d_silu_int8_mixed_bf16_xpu(self):
+        r"""
+        This testcase will quantize Conv2d->SiLU pattern.
+        Match.nodes:
+            [qconv2d_pointwise_default, convert_element_type, sigmoid, mul,
+             convert_element_type, quantize_per_tensor]
+            [qconv2d_pointwise_default, convert_element_type, sigmoid, mul, convert_element_type]
+        """
+        self._qconv2d_unary_test_helper(
+            device="xpu",
+            unary_op=torch.nn.SiLU(),
+            int8_mixed_bf16=True,
+            qconv2d_unary_matcher_nodes=11,
+        )
+
     def _qconv2d_add_test_helper(
         self, device="cpu", use_relu=False, int8_mixed_bf16=False
     ):
@@ -1441,6 +1514,13 @@ def test_qconv2d_add_int8_mixed_bf16(self):
         self._qconv2d_add_test_helper(int8_mixed_bf16=True)
         self._qconv2d_add_test_helper2(int8_mixed_bf16=True)
 
+    @skipIfNoDynamoSupport
+    @skipIfNoONEDNNBF16
+    @skipIfNoONEDNN
+    @skipIfNoXPU
+    def test_qconv2d_add_int8_mixed_bf16_xpu(self):
+        self._qconv2d_add_test_helper(device="xpu", int8_mixed_bf16=True)
+
     @skipIfNoDynamoSupport
     @skipIfNoONEDNN
     def test_qconv2d_add_relu_cpu(self):
@@ -1461,6 +1541,13 @@ def test_qconv2d_add_relu_int8_mixed_bf16(self):
         self._qconv2d_add_test_helper(use_relu=True, int8_mixed_bf16=True)
         self._qconv2d_add_test_helper2(use_relu=True, int8_mixed_bf16=True)
 
+    @skipIfNoDynamoSupport
+    @skipIfNoONEDNNBF16
+    @skipIfNoONEDNN
+    @skipIfNoXPU
+    def test_qconv2d_add_relu_int8_mixed_bf16_xpu(self):
+        self._qconv2d_add_test_helper(device="xpu", use_relu=True, int8_mixed_bf16=True)
+
     @skipIfNoDynamoSupport
     @skipIfNoONEDNN
     def test_qconv2d_add_broadcast_shapes_cpu(self):