throw exception on unbroadcastable input for binary ops

Author: pinzhenx

torch_ipex/csrc/cpu/DevOPs.cpp

@@ -263,16 +263,22 @@ at::Tensor& AtenIpexCPUDev::dil_add_out(
   CHECK_DNNL_OP_PRE_COND(self);
   CHECK_DNNL_OP_PRE_COND(other);
 
+  TORCH_CHECK(self.sizes().equals(other.sizes()),
+      "dil_add not support broadcast yet");
+  auto inferred_size = self.sizes();
+  if (!result.sizes().equals(inferred_size)) {
+    result.resize_(inferred_size);
+  }
+
   dbl::comm::reorder_to_bf16_for_mix_prec(self);
   dbl::comm::reorder_to_bf16_for_mix_prec(other);
   dbl::comm::reorder_to_bf16_for_mix_prec(result);
 
-  dil::tensor x = dbl::comm::try_gen_dil_tensor(self);
-  dil::tensor y = dbl::comm::try_gen_dil_tensor(other);
+  auto x = dbl::comm::try_gen_dil_tensor(self);
+  auto y = dbl::comm::try_gen_dil_tensor(other);
+  auto z = dbl::comm::try_gen_dil_tensor(result);
 
-  dil::tensor z = dbl::comm::try_gen_dil_tensor(result);
-  const std::vector<float> scales{1.0, alpha.to<float>()};
-  dil::sum::compute(scales, {x, y}, z);
+  dil::sum::compute({1.0, alpha.to<float>()}, {x, y}, z);
 
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(z.is_public_format() || check_tensor_own_whole_storage(result));
   dbl::comm::sync_shape_from_dil_to_aten(result, z);
@@ -284,36 +290,25 @@ at::Tensor AtenIpexCPUDev::dil_add(const at::Tensor& self, const at::Tensor& oth
   CHECK_DNNL_OP_PRE_COND(self);
   CHECK_DNNL_OP_PRE_COND(other);
 
+  TORCH_CHECK(self.sizes().equals(other.sizes()),
+      "dil_add not support broadcast yet");
+
   dbl::comm::reorder_to_bf16_for_mix_prec(self);
   dbl::comm::reorder_to_bf16_for_mix_prec(other);
 
-  dil::tensor x = dbl::comm::try_gen_dil_tensor(self);
-  dil::tensor y = dbl::comm::try_gen_dil_tensor(other);
-
+  auto x = dbl::comm::try_gen_dil_tensor(self);
+  auto y = dbl::comm::try_gen_dil_tensor(other);
   dil::tensor z;
-  const std::vector<float> scales{1.0, alpha.to<float>()};
-  dil::sum::compute(scales, {x, y}, z);
+
+  dil::sum::compute({1.0, alpha.to<float>()}, {x, y}, z);
 
   return dbl::comm::gen_aten_tensor_by(std::move(z));
 }
 
 at::Tensor & AtenIpexCPUDev::dil_add_(at::Tensor& self, const at::Tensor& other, at::Scalar alpha) {
   DEBUG("AtenIpexCPUDev::dil_add_\n");
-  CHECK_DNNL_OP_PRE_COND(self);
-  CHECK_DNNL_OP_PRE_COND(other);
 
-  dbl::comm::reorder_to_bf16_for_mix_prec(self);
-  dbl::comm::reorder_to_bf16_for_mix_prec(other);
-
-  auto dil_self = dbl::comm::try_gen_dil_tensor(self);
-  auto dil_other = dbl::comm::try_gen_dil_tensor(other);
-
-  const std::vector<float> scales{1.0, alpha.to<float>()};
-  dil::sum::compute(scales, {dil_self, dil_other}, dil_self);
-
-  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dil_self.is_public_format() || check_tensor_own_whole_storage(self));
-  dbl::comm::sync_shape_from_dil_to_aten(self, dil_self);
-  return self;
+  return dil_add_out(self, self, other, alpha);
 }
 
 at::Tensor& AtenIpexCPUDev::dil_mul_out(at::Tensor& result, const at::Tensor& self, const at::Tensor& other) {
@@ -322,6 +317,13 @@ at::Tensor& AtenIpexCPUDev::dil_mul_out(at::Tensor& result, const at::Tensor& se
   CHECK_DNNL_OP_PRE_COND(self);
   CHECK_DNNL_OP_PRE_COND(other);
 
+  TORCH_CHECK(self.sizes().equals(other.sizes()),
+      "dil_mul not support broadcast yet");
+  auto inferred_size = self.sizes();
+  if (!result.sizes().equals(inferred_size)) {
+    result.resize_(inferred_size);
+  }
+
   dbl::comm::reorder_to_bf16_for_mix_prec(self);
   dbl::comm::reorder_to_bf16_for_mix_prec(other);
   dbl::comm::reorder_to_bf16_for_mix_prec(result);
@@ -339,21 +341,27 @@ at::Tensor& AtenIpexCPUDev::dil_mul_out(at::Tensor& result, const at::Tensor& se
 
 at::Tensor AtenIpexCPUDev::dil_mul(const at::Tensor& self, const at::Tensor& other) {
   DEBUG("AtenIpexCPUDev::dil_mul\n");
+  CHECK_DNNL_OP_PRE_COND(self);
+  CHECK_DNNL_OP_PRE_COND(other);
+
+  TORCH_CHECK(self.sizes().equals(other.sizes()),
+      "dil_mul not support broadcast yet");
 
   dbl::comm::reorder_to_bf16_for_mix_prec(self);
   dbl::comm::reorder_to_bf16_for_mix_prec(other);
 
-  at::Tensor result = dbl::comm::empty_dil_tensor(self.sizes(), self.options());
+  auto x = dbl::comm::try_gen_dil_tensor(self);
+  auto y = dbl::comm::try_gen_dil_tensor(other);
+  dil::tensor z;
+
+  dil::binary::compute(x, y, z, dil::algorithm::binary_mul);
 
-  return dil_mul_out(result, self, other);
+  return dbl::comm::gen_aten_tensor_by(std::move(z));
 }
 
 at::Tensor& AtenIpexCPUDev::dil_mul_(at::Tensor& self, const at::Tensor& other) {
   DEBUG("AtenIpexCPUDev::dil_mul_\n");
 
-  dbl::comm::reorder_to_bf16_for_mix_prec(self);
-  dbl::comm::reorder_to_bf16_for_mix_prec(other);
-
   return dil_mul_out(self, self, other);
 }
 
@@ -472,7 +480,7 @@ at::Tensor& AtenIpexCPUDev::dil_baddbmm_out(
     result.resize_(inferred_size);
   }
   TORCH_CHECK(self.sizes().equals(inferred_size),
-      "baddbmm not support broadcast yet");
+      "dil_baddbmm not support broadcast yet");
 
   dbl::comm::reorder_to_bf16_for_mix_prec(result);
   dbl::comm::reorder_to_bf16_for_mix_prec(self);
@@ -541,7 +549,7 @@ at::Tensor& AtenIpexCPUDev::dil_addmm_out(
     result.resize_(inferred_size);
   }
   TORCH_CHECK(self.sizes().equals(inferred_size),
-      "addmm not support broadcast yet");
+      "dil_addmm not support broadcast yet");
 
   dbl::comm::reorder_to_bf16_for_mix_prec(result);
   dbl::comm::reorder_to_bf16_for_mix_prec(self);
@@ -610,7 +618,7 @@ at::Tensor& AtenIpexCPUDev::dil_addbmm_out(
     result.resize_(inferred_size);
   }
   TORCH_CHECK(self.sizes().equals(inferred_size),
-      "addbmm not support broadcast yet");
+      "dil_addbmm not support broadcast yet");
 
   dbl::comm::reorder_to_bf16_for_mix_prec(result);
   dbl::comm::reorder_to_bf16_for_mix_prec(self);