Fix & cleanup op t_copy (#701)

Summary:
Pull Request resolved: #701

ghstack-source-id: 203341581
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D49735860

fbshipit-source-id: b6acfd1fbc89c2fe1475368e48f210b71733ef1f

Author: manuelcandales

kernels/portable/cpu/op_t_copy.cpp

@@ -6,11 +6,9 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <cstring>
-
 #include <executorch/kernels/portable/cpu/util/transpose_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
-#include <executorch/runtime/platform/assert.h>
+#include <cstring>
 
 namespace torch {
 namespace executor {
@@ -20,55 +18,47 @@ using SizesType = exec_aten::SizesType;
 using StridesType = exec_aten::StridesType;
 using Tensor = exec_aten::Tensor;
 
-namespace {
-
-/**
- * Verifies preconditions of t_copy_int_out
- */
-void check_preconditions(const Tensor& a, Tensor& out) {
-  auto a_dim = a.dim();
-  ET_CHECK_MSG(
-      a_dim >= 0 && a_dim <= 2,
-      "Rank of tensor a has to be <=2 but received tensor of rank : %zd.:",
-      a_dim);
-  if (a_dim < 2) {
-    ET_CHECK_SAME_SHAPE_AND_DTYPE2(a, out);
-  } else {
-    ET_CHECK_SAME_DTYPE2(a, out);
-    ET_CHECK_MSG(
-        (a.sizes()[0] == out.sizes()[1]) && (a.sizes()[1] == out.sizes()[0]),
-        "Input tensor and output tensor shapes do not support transposing");
-    ET_CHECK_MSG(out.dim() == 2, "Output tensor must have same dim (2)");
-  }
-}
-
-} // namespace
-
 /**
  * Expects input to be <= 2-D tensor and transposes dimensions 0 and 1.
  * 0-D and 1-D tensors are returned as is. When input is a 2-D tensor this
  * is equivalent to transpose(input, 0, 1).
  * t_copy.out(Tensor self, Tensor(a!) out)
  */
-Tensor& t_copy_out(RuntimeContext& ctx, const Tensor& a, Tensor& out) {
+Tensor& t_copy_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
   (void)ctx;
-  check_preconditions(a, out);
-  int dim_1 = a.sizes().size() == 2 ? 1 : 0;
-#define TRANSPOSE_TENSORS(ctype, dtype)         \
-  case ScalarType::dtype:                       \
-    transpose_tensors<ctype>(a, 0, dim_1, out); \
-    break;
 
-  switch (a.scalar_type()) {
-    ET_FORALL_SCALAR_TYPES(TRANSPOSE_TENSORS)
-    default:
-      ET_CHECK_MSG(
-          false,
-          "Unhandled dtype %" PRId8,
-          static_cast<int8_t>(a.scalar_type()));
+  ET_KERNEL_CHECK(ctx, check_t_copy_args(in, out), InvalidArgument, out);
+
+  ScalarType in_type = in.scalar_type();
+
+  if (in.dim() < 2) {
+    // Resize for dynamic shape
+    ET_KERNEL_CHECK(
+        ctx, resize_tensor(out, in.sizes()) == Error::Ok, InvalidArgument, out);
+
+    ET_SWITCH_ALL_TYPES(in_type, ctx, __func__, CTYPE, [&]() {
+      const CTYPE* in_data = in.const_data_ptr<CTYPE>();
+      CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
+      memcpy(out_data, in_data, in.nbytes());
+    });
+
+    return out;
   }
 
-#undef TRANSPOSE_TENSORS
+  Tensor::SizesType expected_out_size[kTensorDimensionLimit];
+  size_t expected_out_dim = 0;
+  get_transpose_out_target_size(in, 1, 0, expected_out_size, &expected_out_dim);
+
+  // Resize for dynamic shape
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_tensor(out, {expected_out_size, expected_out_dim}) == Error::Ok,
+      InvalidArgument,
+      out);
+
+  ET_SWITCH_ALL_TYPES(in_type, ctx, __func__, CTYPE, [&] {
+    transpose_tensors<CTYPE>(in, 1, 0, out);
+  });
 
   return out;
 }

kernels/portable/cpu/util/transpose_util.h

@@ -135,6 +135,12 @@ void transpose_tensors(
   }
 }
 
+inline bool check_t_copy_args(const Tensor& in, Tensor& out) {
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_has_rank_smaller_or_equal_to(in, 2));
+  return true;
+}
+
 inline bool check_transpose_copy_args(
     const Tensor& in,
     int64_t dim0,