Dtype compliance: unbind_copy

Reviewed By: SS-JIA

Differential Revision: D48318691

fbshipit-source-id: 6d01021e2e567addeb5faaf5e9714c9d34139528

Author: manuelcandales

kernels/portable/cpu/op_unbind_copy.cpp

@@ -41,13 +41,13 @@ void check_args(const Tensor& input, int64_t dim, TensorList out) {
 
   // Validate each output.
   for (size_t i = 0; i < out.size(); ++i) {
-    // All output dtypes must match the input type.
+    // All output dtypes must be the same.
     ET_CHECK_MSG(
-        out[i].scalar_type() == input.scalar_type(),
-        "out[%zu] dtype %hhd != input dtype %hhd",
+        out[i].scalar_type() == out[0].scalar_type(),
+        "out[%zu] dtype %hhd != out[0] dtype %hhd",
         i,
         out[i].scalar_type(),
-        input.scalar_type());
+        out[0].scalar_type());
 
     // output tensor must have # of dims = input.dim() -1
     ET_CHECK_MSG(
@@ -97,25 +97,29 @@ void unbind_copy_int_out(
 
   const size_t leading_dims = getLeadingDims(input, dim);
   const size_t trailing_dims = getTrailingDims(input, dim);
-
-  const size_t element_size = input.element_size();
-  const size_t step = input.size(dim) * trailing_dims * element_size;
-
-  const char* input_data = input.const_data_ptr<char>();
-  for (size_t i = 0, e = out.size(); i < e; ++i) {
-    size_t num_bytes = trailing_dims * element_size;
-    // num_bytes should not be zero because trailing_dims
-    // will at least return 1
-
-    const char* src = input_data;
-    char* dest = out[i].mutable_data_ptr<char>();
-    for (size_t j = 0; j < leading_dims; ++j) {
-      memcpy(dest, src, num_bytes);
-      src += step;
-      dest += num_bytes;
-    }
-    input_data += num_bytes;
-  }
+  const size_t step = input.size(dim) * trailing_dims;
+
+  ScalarType in_type = input.scalar_type();
+  ScalarType out_type = out[0].scalar_type();
+
+  ET_SWITCH_REAL_TYPES_AND(Bool, in_type, ctx, __func__, CTYPE_IN, [&]() {
+    ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, __func__, CTYPE_OUT, [&]() {
+      const CTYPE_IN* const input_data = input.const_data_ptr<CTYPE_IN>();
+      for (size_t i = 0, e = out.size(); i < e; ++i) {
+        size_t input_offset = i * trailing_dims;
+        CTYPE_OUT* const dest = out[i].mutable_data_ptr<CTYPE_OUT>();
+        size_t dest_offset = 0;
+        for (size_t j = 0; j < leading_dims; ++j) {
+          for (size_t k = 0; k < trailing_dims; ++k) {
+            dest[dest_offset + k] =
+                convert<CTYPE_OUT, CTYPE_IN>(input_data[input_offset + k]);
+          }
+          input_offset += step;
+          dest_offset += trailing_dims;
+        }
+      }
+    });
+  });
 }
 
 } // namespace native

kernels/test/op_unbind_copy_test.cpp

@@ -48,8 +48,8 @@ Tensor make1x2x3(TensorFactory<DTYPE>& tf) {
 
 template <ScalarType DTYPE>
 void test_unbind_dim0() {
-  TensorFactory<ScalarType::Int> tf;
-  TensorListFactory<ScalarType::Int> tlf;
+  TensorFactory<DTYPE> tf;
+  TensorListFactory<DTYPE> tlf;
 
   // clang-format off
   std::vector<Tensor> expected_out = {
@@ -117,8 +117,8 @@ void test_unbind_dim1() {
 
 template <ScalarType DTYPE>
 void test_unbind_dim2() {
-  TensorFactory<ScalarType::Int> tf;
-  TensorListFactory<ScalarType::Int> tlf;
+  TensorFactory<DTYPE> tf;
+  TensorListFactory<DTYPE> tlf;
 
   // Splitting on dim=N with split_size=2 will produce a list of tensors where
   // the max dim[N] is 2, and the other dims are the same as the input.
@@ -164,24 +164,21 @@ void test_unbind_dim2() {
   EXPECT_TENSOR_LISTS_EQ(expected_out, out2);
 }
 
-TEST(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim0AllSupportedDtypes) {
+TEST(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim0AllRealDtypes) {
 #define TEST_ENTRY(ctype, dtype) test_unbind_dim0<ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
+  ET_FORALL_REAL_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
-TEST(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim1AllSupportedDTypes) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+TEST(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim1AllRealDTypes) {
 #define TEST_ENTRY(ctype, dtype) test_unbind_dim1<ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
+  ET_FORALL_REAL_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
-TEST(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim2) {
+TEST(OpUnbindCopyIntOutTest, Unbind1x2x3OnDim2AllRealDTypes) {
 #define TEST_ENTRY(ctype, dtype) test_unbind_dim2<ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
+  ET_FORALL_REAL_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }