Dtype compliance: slice_scatter

Reviewed By: SS-JIA

Differential Revision: D48318692

fbshipit-source-id: d39ae3a2dfe8b1e260ab14ce66144862c7f26867

Author: manuelcandales

kernels/portable/cpu/op_slice_scatter.cpp

@@ -39,8 +39,8 @@ void check_input_args(
   // Input and output tensors should be the same shape
   ET_CHECK_SAME_SHAPE2(input, output);
 
-  // All tensors should have the same dtype
-  ET_CHECK_SAME_DTYPE3(input, src, output);
+  // Input and output tensors should have the same shape
+  ET_CHECK_SAME_DTYPE2(input, output);
 
   // The input.dim() shall equal to src.dim()
   ET_CHECK_MSG(
@@ -158,28 +158,36 @@ Tensor& slice_scatter_out(
   check_input_args(input, src, dim, num_values, step, out);
 
   size_t dim_length = input.size(dim);
-
   size_t leading_dims = getLeadingDims(input, dim);
   size_t trailing_dims = getTrailingDims(input, dim);
 
-  size_t length_per_step = trailing_dims * input.element_size();
-
-  const char* in_data = input.const_data_ptr<char>();
-  const char* src_data = src.const_data_ptr<char>();
-
-  char* out_data = out.mutable_data_ptr<char>();
-
   // To start, copy the input into the output
-  memcpy(out_data, in_data, input.nbytes());
-
-  for (int i = 0; i < leading_dims; i++) {
-    char* dst = out_data + (i * dim_length + start) * length_per_step;
-    for (int j = 0; j < num_values; j++) {
-      memcpy(dst, src_data, length_per_step);
-      src_data += length_per_step;
-      dst += step * length_per_step;
-    }
-  }
+  memcpy(out.mutable_data_ptr(), input.const_data_ptr(), input.nbytes());
+
+  ScalarType in_type = input.scalar_type();
+  ScalarType src_type = src.scalar_type();
+
+  ET_SWITCH_REAL_TYPES_AND(Bool, in_type, ctx, __func__, CTYPE, [&]() {
+    ET_SWITCH_REAL_TYPES_AND(Bool, src_type, ctx, __func__, CTYPE_SRC, [&]() {
+      CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
+      const CTYPE_SRC* src_data = src.const_data_ptr<CTYPE_SRC>();
+
+      size_t src_offset = 0;
+
+      for (int i = 0; i < leading_dims; i++) {
+        size_t out_offset = (i * dim_length + start) * trailing_dims;
+        for (int j = 0; j < num_values; j++) {
+          for (size_t k = 0; k < trailing_dims; ++k) {
+            out_data[out_offset + k] =
+                convert<CTYPE, CTYPE_SRC>(src_data[src_offset + k]);
+          }
+          src_offset += trailing_dims;
+          out_offset += step * trailing_dims;
+        }
+      }
+    });
+  });
+
   return out;
 }
 

kernels/test/op_slice_scatter_test.cpp

@@ -666,12 +666,9 @@ void test_dtype() {
   EXPECT_TENSOR_EQ(ret, expect_ret);
 }
 
-TEST(OpSliceCopyTensorOutTest, AllDtypesSupported) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+TEST(OpSliceCopyTensorOutTest, AllRealDtypesSupported) {
 #define TEST_ENTRY(ctype, dtype) test_dtype<ctype, ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
+  ET_FORALL_REAL_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
   // TODO: Also add tests for half, complex, quantized, and other types. Easiest
   // way to do that would be to make TensorFactory support zeros() and ones()
@@ -755,22 +752,6 @@ TEST(OpSliceCopyTensorOutTest, MismatchedOutDtypesDies) {
       input, src, /*dim=*/0, /*start=*/0, /*end=*/1, /*step=*/1, out));
 }
 
-TEST(OpSliceCopyTensorOutTest, MismatchedSrcDtypesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle src with mismatched dtype";
-  }
-  TensorFactory<ScalarType::Int> tf_int;
-  TensorFactory<ScalarType::Float> tf_float;
-  Tensor input = tf_int.zeros({1, 2, 2});
-  Tensor src = tf_float.zeros({1, 2, 2});
-
-  // Size is compatible to the output, but a mismatched dtype.
-  Tensor out = tf_int.ones({1, 2, 2});
-
-  ET_EXPECT_KERNEL_FAILURE(op_slice_scatter_out(
-      input, src, /*dim=*/0, /*start=*/0, /*end=*/1, /*step=*/1, out));
-}
-
 TEST(OpSliceCopyTensorOutTest, OutSizeMismatchDimDies) {
   if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
     GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";