Revert ATen _clone_dim_order implementation and refactoring

Author: keyprocedure

kernels/aten/cpu/op__clone_dim_order.cpp

@@ -6,7 +6,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/kernels/aten/cpu/util/copy_ops_util.h>
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 
@@ -26,6 +25,71 @@ using OptionalArrayRef = executorch::aten::OptionalArrayRef<T>;
 template <typename T>
 using Optional = std::optional<T>;
 
+namespace {
+Optional<MemoryFormat> get_memory_format(OptionalArrayRef<int64_t> dim_order) {
+  if (!dim_order.has_value()) {
+    return executorch::aten::nullopt;
+  }
+  if (is_contiguous_dim_order(
+          dim_order.value().data(), dim_order.value().size())) {
+    return MemoryFormat::Contiguous;
+  } else if (is_channels_last_dim_order(
+                 dim_order.value().data(), dim_order.value().size())) {
+    return MemoryFormat::ChannelsLast;
+  } else {
+    ET_ASSERT_UNREACHABLE();
+  }
+}
+
+bool check__to_dim_order_copy_args(
+    const Tensor& input,
+    bool non_blocking,
+    executorch::aten::OptionalArrayRef<int64_t> dim_order,
+    Tensor& out) {
+  // Right now we only support blocking data transfer
+  ET_LOG_AND_RETURN_IF_FALSE(non_blocking == false);
+
+  // dim_order is set, the target dim_order will be either contiguous or
+  // channels_last memory format
+  if (dim_order.has_value()) {
+    executorch::aten::ArrayRef<int64_t> dim_order_ref = dim_order.value();
+
+    // dim order size shall equal to input dim
+    ET_LOG_AND_RETURN_IF_FALSE(dim_order_ref.size() == input.dim());
+
+    ET_LOG_AND_RETURN_IF_FALSE(
+        is_channels_last_dim_order(
+            dim_order.value().data(), dim_order.value().size()) ||
+        is_contiguous_dim_order(
+            dim_order.value().data(), dim_order.value().size()));
+
+    // Out Aten tensor shall have same memory format stride as dim_order
+    const size_t kMaxNumOfDimensions = 16;
+    ET_LOG_AND_RETURN_IF_FALSE(kMaxNumOfDimensions >= out.dim());
+    executorch::aten::StridesType target_strides[kMaxNumOfDimensions];
+    dim_order_to_stride_nocheck(
+        out.sizes().data(),
+        dim_order_ref.data(),
+        dim_order_ref.size(),
+        target_strides);
+    ET_LOG_AND_RETURN_IF_FALSE(out.dim() == dim_order_ref.size());
+    for (size_t i = 0; i < dim_order_ref.size(); i++) {
+      ET_LOG_AND_RETURN_IF_FALSE(target_strides[i] == out.strides()[i]);
+    }
+
+  } else { // dim_order is not set, preserve the dim order of input
+
+    auto out_strides = out.strides();
+    auto input_strides = input.strides();
+    ET_LOG_AND_RETURN_IF_FALSE(input_strides.size() == out_strides.size());
+    for (size_t i = 0; i < input_strides.size(); i++) {
+      ET_LOG_AND_RETURN_IF_FALSE(input_strides[i] == out_strides[i]);
+    }
+  }
+  return true;
+}
+} // namespace
+
 // _to_dim_order_copy.out(Tensor self, *, bool non_blocking=False, int[]?
 // dim_order=None, Tensor(a!) out) -> Tensor(a!)
 Tensor& _to_dim_order_copy_out(

kernels/aten/cpu/op__to_dim_order_copy.cpp

@@ -15,28 +15,6 @@ namespace torch {
 namespace executor {
 
 using Tensor = executorch::aten::Tensor;
-using MemoryFormat = executorch::aten::MemoryFormat;
-
-/**
- * Determines the memory format (Contiguous or ChannelsLast) corresponding to
- * the dim_order. Provides support for bridging torch.memory_format with
- * ExecuTorch's dim_order.
- */
-std::optional<MemoryFormat> get_memory_format(
-    executorch::aten::OptionalArrayRef<int64_t> dim_order) {
-  if (!dim_order.has_value()) {
-    return executorch::aten::nullopt;
-  }
-  if (is_contiguous_dim_order(
-          dim_order.value().data(), dim_order.value().size())) {
-    return MemoryFormat::Contiguous;
-  } else if (is_channels_last_dim_order(
-                 dim_order.value().data(), dim_order.value().size())) {
-    return MemoryFormat::ChannelsLast;
-  } else {
-    ET_ASSERT_UNREACHABLE();
-  }
-}
 
 bool check__to_dim_order_copy_args(
     const Tensor& input,

kernels/aten/cpu/util/copy_ops_util.cpp

@@ -13,9 +13,6 @@
 namespace torch {
 namespace executor {
 
-std::optional<MemoryFormat> get_memory_format(
-    executorch::aten::OptionalArrayRef<int64_t> dim_order);
-
 bool check__to_dim_order_copy_args(
     const Tensor& input,
     bool non_blocking,

kernels/aten/cpu/util/copy_ops_util.h

@@ -11,8 +11,3 @@
   kernels:
     - arg_meta: null
       kernel_name: torch::executor::_to_dim_order_copy_out
-
-- func: dim_order_ops::_clone_dim_order.out(Tensor self, *, bool non_blocking=False, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!)
-  kernels:
-    - arg_meta: null
-      kernel_name: torch::executor::_clone_dim_order_out