llga: no data type promotion needed for binary op of `Double' (#2622)

* llga: no data type promotion needed for binary op when one data type is 'Double' * llga: no promotion for data types not supported by oneDNN Graph * fix clang-format error --------- Co-authored-by: Chunyuan WU <[email protected]> Co-authored-by: sanchitintel <[email protected]> Co-authored-by: WeizhuoZhang-intel <[email protected]>
intel · Mar 7, 2024 · f4ee125 · f4ee125
1 parent 55efdde
commit f4ee125
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 24 deletions.
diff --git a/csrc/cpu/jit/codegen/LlgaTensorImpl.cpp b/csrc/cpu/jit/codegen/LlgaTensorImpl.cpp
@@ -56,6 +56,26 @@ dnnl::graph::tensor llga_from_aten_tensor(const at::Tensor& tensor) {
       tensor.data_ptr()};
 }
 
+using data_type = dnnl::graph::logical_tensor::data_type;
+data_type getLlgaDataType(at::ScalarType dt) {
+  switch (dt) {
+    case at::ScalarType::Float:
+      return data_type::f32;
+    case at::ScalarType::BFloat16:
+      return data_type::bf16;
+    case at::ScalarType::Bool:
+      return data_type::boolean;
+    case at::kInt:
+      return data_type::s32;
+    case at::ScalarType::QInt8:
+      return data_type::s8;
+    case at::ScalarType::QUInt8:
+      return data_type::u8;
+    default:
+      return data_type::undef;
+  }
+}
+
 at::Tensor LlgaTensorImpl::llga_to_aten_tensor(LlgaTensorImpl* llgaImpl) {
   auto aten_tensor = at::detail::make_tensor<TensorImpl>(
       std::move(llgaImpl->storage_),
@@ -81,27 +101,6 @@ at::Tensor LlgaTensorImpl::llga_to_aten_tensor(
   return aten_tensor;
 }
 
-using data_type = dnnl::graph::logical_tensor::data_type;
-
-data_type LlgaTensorDesc::getLlgaDataType(at::ScalarType dt) const {
-  switch (dt) {
-    case at::ScalarType::Float:
-      return data_type::f32;
-    case at::ScalarType::BFloat16:
-      return data_type::bf16;
-    case at::ScalarType::Bool:
-      return data_type::boolean;
-    case at::kInt:
-      return data_type::s32;
-    case at::ScalarType::QInt8:
-      return data_type::s8;
-    case at::ScalarType::QUInt8:
-      return data_type::u8;
-    default:
-      return data_type::undef;
-  }
-}
-
 LlgaTensorDesc LlgaTensorDesc::supplementTensorInfo(const at::Tensor& t) const {
   if (t.is_mkldnn()) {
     // if input tensor is of mkldnn, it's originated from an upstream

diff --git a/csrc/cpu/jit/codegen/LlgaTensorImpl.h b/csrc/cpu/jit/codegen/LlgaTensorImpl.h
@@ -13,6 +13,8 @@ namespace jit {
 namespace fuser {
 namespace onednn {
 
+dnnl::graph::logical_tensor::data_type getLlgaDataType(at::ScalarType dt);
+
 struct LlgaTensorDesc {
   using desc = dnnl::graph::logical_tensor;
 
@@ -106,9 +108,6 @@ struct LlgaTensorDesc {
 
   at::ScalarType aten_scalar_type() const;
 
-  dnnl::graph::logical_tensor::data_type getLlgaDataType(
-      at::ScalarType dt) const;
-
   const std::vector<int64_t>& sizes() const {
     return sizes_;
   }

diff --git a/csrc/cpu/jit/codegen/onednn/prepare_binary.cpp b/csrc/cpu/jit/codegen/onednn/prepare_binary.cpp
@@ -1,6 +1,7 @@
 #include "prepare_binary.h"
 #include <torch/csrc/jit/passes/dead_code_elimination.h>
 #include <torch/csrc/jit/passes/shape_analysis.h>
+#include "../LlgaTensorImpl.h"
 #include "utils.h"
 
 namespace torch_ipex {
@@ -9,6 +10,7 @@ namespace fuser {
 namespace onednn {
 
 using namespace torch::jit;
+using data_type = dnnl::graph::logical_tensor::data_type;
 
 void handleBinaryOpInputs(Node* node, int first_input, int second_input) {
   if (node->input(first_input)->type()->isSubtypeOf(TensorType::get()) &&
@@ -33,6 +35,10 @@ void handleBinaryOpInputs(Node* node, int first_input, int second_input) {
       // https://pytorch.org/docs/stable/tensor_attributes.html#type-promotion-doc
       // clang-format on
       auto promotedDtype = dtypeOfFirstInput;
+      // This tensor won't be added to oneDNN graph due to unsupported data
+      // type, so no need to do promotion for it.
+      if (getLlgaDataType(promotedDtype) == data_type::undef)
+        return;
       utils::convertInputTo0DTensor(node, second_input, promotedDtype);
       // dtype might have changed, so needs to be updated in IR as well
       utils::modifyDtypeOfNode(node, promotedDtype);
@@ -53,6 +59,10 @@ void handleBinaryOpInputs(Node* node, int first_input, int second_input) {
           // Type promotion is required
           auto promotedDtype =
               c10::promoteTypes(dtypeOfFirstInput, dtypeOfSecondInput);
+          // This tensor won't be added to oneDNN graph due to unsupported data
+          // type, so no need to do promotion for it.
+          if (getLlgaDataType(promotedDtype) == data_type::undef)
+            return;
           int input_to_replace;
           if (promotedDtype == dtypeOfFirstInput) {
             input_to_replace = second_input;
@@ -65,6 +75,10 @@ void handleBinaryOpInputs(Node* node, int first_input, int second_input) {
           utils::mark_original_output_dtype(node);
           utils::modifyDtypeOfNode(node, promotedDtype);
         } else {
+          // This tensor won't be added to oneDNN graph due to unsupported data
+          // type, so no need to do promotion for it.
+          if (getLlgaDataType(dtypeOfFirstInput) == data_type::undef)
+            return;
           // both dtypes are same
           // IR info of dtypes is missing sometimes in JIT IR,
           // and we shouldn't treat those tensors as FP32 tensors by default.

diff --git a/tests/cpu/test_jit_llga_fuser.py b/tests/cpu/test_jit_llga_fuser.py
@@ -450,6 +450,27 @@ def forward(self, x, y):
         self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)
         self.assertFused(graph, ["aten::matmul", "aten::div"])
 
+    @llga_fp32_bf16_test_env
+    def test_bmm_div_with_double_dt(self):
+        class M(nn.Module):
+            def __init__(self):
+                super(M, self).__init__()
+                self.divisor = torch.randn(1, dtype=torch.float64)
+
+            def forward(self, x, y):
+                return x.matmul(y) / self.divisor
+
+        x = torch.randn(128, 16, 384, 64)
+        y = torch.randn(128, 16, 64, 384)
+        m = M()
+
+        graph, _ = self.checkTrace(m, [x, y])
+        self.assertGraphContainsExactly(graph, LLGA_FUSION_GROUP, 1)
+        # no need to do data type promotion for `Double` which llga doesn't
+        # support
+        self.assertGraphContainsExactly(graph, "aten::to", 0)
+        self.assertFused(graph, ["aten::matmul"])
+
     @llga_fp32_bf16_test_env
     def test_bmm_div_add(self):
         class M(nn.Module):