Dtype compliance: split_copy

Reviewed By: SS-JIA

Differential Revision: D48318690

fbshipit-source-id: 935759301c9a4deada20550bb4a2de549646f9ac

kernels/portable/cpu/op_split_copy.cpp

@@ -87,13 +87,13 @@ void check_args(
 
  // 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());
 
  // All outputs must have the same number of dimensions as the input.
  ET_CHECK_MSG(
@@ -170,26 +170,32 @@ void split_copy_Tensor_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 = out[i].size(dim) * trailing_dims * element_size;
- if (num_bytes == 0) {
- continue;
- }
-
- 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* input_data = input.const_data_ptr<CTYPE_IN>();
+ for (size_t i = 0, e = out.size(); i < e; ++i) {
+ size_t out_step = out[i].size(dim) * trailing_dims;
+ if (out_step == 0) {
+ continue;
+ }
+ const CTYPE_IN* src = input_data;
+ CTYPE_OUT* dest = out[i].mutable_data_ptr<CTYPE_OUT>();
+ for (size_t j = 0; j < leading_dims; ++j) {
+ for (size_t k = 0; k < out_step; ++k) {
+ dest[k] = convert<CTYPE_OUT, CTYPE_IN>(src[k]);
+ }
+ src += step;
+ dest += out_step;
+ }
+ input_data += out_step;
+ }
+ });
+ });
 }
 
 } // namespace native

kernels/test/op_split_copy_test.cpp

@@ -356,9 +356,7 @@ TEST(OpSplitCopyTensorOutTest, OutOfRangeDimsDie) {
 }
 
 TEST(OpSplitCopyTensorOutTest, DtypeMismatchDies) {
- if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
- GTEST_SKIP() << "ATen kernel can handle dtype mismatch";
- }
+ GTEST_SKIP() << "ATen kernel can handle dtype mismatch";
  TensorFactory<ScalarType::Int> tf_int;
  TensorListFactory<ScalarType::Int> tlf_int;
  TensorListFactory<ScalarType::Float> tlf_float;