a-start ops | add dim order regulation

kernels/portable/cpu/op_add.cpp

@@ -79,6 +79,8 @@ Tensor& add_out(
       out);
 
   ET_KERNEL_CHECK(ctx, tensor_is_realhb_type(out), InvalidArgument, out);
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(a, b, out), InvalidArgument, out);
 
   ScalarType a_type = a.scalar_type();
   ScalarType b_type = b.scalar_type();
@@ -131,6 +133,8 @@ Tensor& add_scalar_out(
       "Failed to resize output tensor.");
 
   ET_KERNEL_CHECK(ctx, tensor_is_realhb_type(out), InvalidArgument, out);
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(a, out), InvalidArgument, out);
 
   ScalarType a_type = a.scalar_type();
   ScalarType b_type = utils::get_scalar_dtype(b);

kernels/portable/cpu/op_addmm.cpp

@@ -45,6 +45,14 @@ Tensor& addmm_out(
   ET_KERNEL_CHECK(
       ctx, tensor_is_broadcastable_to(in, out), InvalidArgument, out);
 
+  ET_KERNEL_CHECK(
+      ctx,
+      tensors_have_same_dim_order(in, mat1, mat2, out),
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
+
   ScalarType alpha_dtype = utils::get_scalar_dtype(alpha);
   ScalarType beta_dtype = utils::get_scalar_dtype(beta);
   ET_SWITCH_REAL_TYPES_AND(

kernels/portable/cpu/op_alias_copy.cpp

@@ -28,6 +28,8 @@ Tensor& alias_copy_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
       "Failed to resize output tensor.");
 
   ET_KERNEL_CHECK(ctx, tensors_have_same_dtype(in, out), InvalidArgument, out);
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
 
   if (in.nbytes() > 0) {
     // Note that this check is important. It's valid for a tensor with numel 0

kernels/portable/cpu/op_allclose.cpp

@@ -104,6 +104,9 @@ Tensor& allclose_out(
       out.scalar_type() == ScalarType::Bool,
       "Out tensor must be type Bool; saw type %" PRId8,
       static_cast<int8_t>(out.scalar_type()));
+  ET_CHECK_MSG(
+      tensors_have_same_dim_order(self, other, out),
+      "self, other and out tensors should have same dim order");
   ET_CHECK_MSG(
       out.numel() == 1,
       "Out tensor must be a single element; saw %zu elements",

kernels/portable/cpu/op_amax.cpp

@@ -39,6 +39,9 @@ Tensor& amax_out(
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   ET_SWITCH_REAL_TYPES_AND(
       Bool, in.scalar_type(), ctx, "amax.out", CTYPE, [&]() {
         CTYPE* out_data = out.mutable_data_ptr<CTYPE>();

kernels/portable/cpu/op_amin.cpp

@@ -39,6 +39,9 @@ Tensor& amin_out(
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   ET_SWITCH_REAL_TYPES_AND(
       Bool, in.scalar_type(), ctx, "amin.out", CTYPE, [&]() {
         CTYPE* out_data = out.mutable_data_ptr<CTYPE>();

kernels/portable/cpu/op_any.cpp

@@ -22,6 +22,9 @@ Tensor& any_all_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
   ET_KERNEL_CHECK(
       ctx, resize_tensor(out, {}) == Error::Ok, InvalidArgument, out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   ScalarType in_type = in.scalar_type();
   ScalarType out_type = out.scalar_type();
   constexpr auto name = "any.all_out";
@@ -68,6 +71,9 @@ Tensor& any_dims_out(
         out);
   }
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   ScalarType in_type = in.scalar_type();
   ScalarType out_type = out.scalar_type();
   constexpr auto name = "any.dims_out";
@@ -122,6 +128,9 @@ Tensor& any_out(
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   ScalarType in_type = in.scalar_type();
   ScalarType out_type = out.scalar_type();
   constexpr auto name = "any.out";

kernels/portable/cpu/op_arange.cpp

@@ -27,6 +27,8 @@ Tensor& arange_out(RuntimeContext& ctx, const Scalar& end, Tensor& out) {
   ET_KERNEL_CHECK(
       ctx, check_arange_args(0.0, end_val, 1.0, out), InvalidArgument, out);
 
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(out), InvalidArgument, out);
+
   size_t size = static_cast<size_t>(std::ceil(end_val));
 
   Tensor::SizesType out_length = static_cast<Tensor::SizesType>(size);
@@ -73,6 +75,8 @@ Tensor& arange_start_out(
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(out), InvalidArgument, out);
+
   double size_d = (d_end - d_start) / d_step;
   size_t size = static_cast<size_t>(std::ceil(size_d));
 

kernels/portable/cpu/op_argmax.cpp

@@ -40,6 +40,9 @@ Tensor& argmax_out(
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "argmax.out", CTYPE, [&] {
     long* out_data = out.mutable_data_ptr<long>();
 

kernels/portable/cpu/op_argmin.cpp

@@ -40,6 +40,9 @@ Tensor& argmin_out(
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "argmin.out", CTYPE, [&] {
     long* out_data = out.mutable_data_ptr<long>();
 

kernels/portable/cpu/op_as_strided_copy.cpp

@@ -37,6 +37,11 @@ Tensor& as_strided_copy_out(
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
+
   if (in.numel() == 0) {
     return out;
   }

kernels/portable/cpu/op_atan2.cpp

@@ -26,6 +26,9 @@ atan2_out(RuntimeContext& ctx, const Tensor& a, const Tensor& b, Tensor& out) {
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(a, b, out), InvalidArgument, out);
+
   ScalarType a_type = a.scalar_type();
   ScalarType b_type = b.scalar_type();
   ScalarType out_type = out.scalar_type();

kernels/portable/cpu/op_avg_pool2d.cpp

@@ -44,6 +44,11 @@ Tensor& avg_pool2d_out(
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
+
   size_t output_ndim = 0;
   exec_aten::SizesType output_sizes[kTensorDimensionLimit];
   get_avg_pool2d_out_target_size(

kernels/portable/cpu/pattern/binary_ufunc_realb_realb_to_realb_logical.cpp

@@ -27,6 +27,9 @@ Tensor& binary_ufunc_realb_realb_to_realb_logical(
       InvalidArgument,
       out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(a, b, out), InvalidArgument, out);
+
   ScalarType a_type = a.scalar_type();
   ScalarType b_type = b.scalar_type();
   ScalarType out_type = out.scalar_type();

kernels/portable/cpu/pattern/unary_ufunc_realh.cpp

@@ -33,6 +33,9 @@ Tensor& unary_ufunc_realh(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_shape_and_dtype(in, out), InvalidArgument, out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   ET_SWITCH_REALH_TYPES(in.scalar_type(), ctx, __func__, CTYPE, [&] {
     apply_unary_map_fn(
         [fn](const CTYPE val_in) { return static_cast<CTYPE>(fn(val_in)); },

kernels/portable/cpu/pattern/unary_ufunc_realhb_to_bool.cpp

@@ -38,6 +38,9 @@ Tensor& unary_ufunc_realhb_to_bool(
       "Expected out tensor to have dtype Bool, but got %" PRId8 " instead.",
       static_cast<int8_t>(out.scalar_type()));
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   const auto in_type = in.scalar_type();
 
   ET_SWITCH_REALHB_TYPES(in_type, ctx, __func__, CTYPE_IN, [&] {

kernels/portable/cpu/pattern/unary_ufunc_realhb_to_floath.cpp

@@ -32,6 +32,9 @@ Tensor& unary_ufunc_realhb_to_floath(
       out,
       "Failed to resize output tensor.");
 
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
   const auto in_type = in.scalar_type();
   const auto out_type = out.scalar_type();
 

runtime/core/exec_aten/util/tensor_util.h

@@ -1153,33 +1153,80 @@ bool tensor_has_valid_dim_order(exec_aten::Tensor t);
 bool tensor_is_default_or_channels_last_dim_order(exec_aten::Tensor t);
 
 /**
- * Asserts that two tensors have the same dim_order
+ * Checks whether a tensor has the default dimension order.
+ * Logs an error message if the tensor does not meet the expected criteria.
  *
- * Note that this macro only tests dim order, but not others like actual data,
- * sizes, etc. Also this macro does not support ATen mode since we do not
- * support dim order in ATen mode.
+ * @param t The tensor to check the dimension order of.
+ * @return True if the tensor has the default dimension order, false otherwise.
+ */
+bool tensor_is_default_dim_order(exec_aten::Tensor t);
+
+/**
+ * Checks whether a tensor has the channels last dimension order.
+ * Logs an error message if the tensor does not meet the expected criteria.
  *
- * TODO(T183094318): Add dim order and related function support for ATen mode.
+ * @param t The tensor to check the dimension order of.
+ * @return True if the tensor has the channels last dimension order, false
+ * otherwise.
  */
+bool tensor_is_channels_last_dim_order(exec_aten::Tensor t);
 
+/**
+ * Asserts that four tensors have the same dim_order
+ *
+ * Note that this macro only tests dim order, but not others like actual data,
+ * sizes, etc.
+ *
+ */
 bool tensors_have_same_dim_order(
+    const exec_aten::ArrayRef<exec_aten::Tensor> tensor_list);
+
+/**
+ * Asserts that two tensors have the same dim_order
+ *
+ * Note that this macro only tests dim order, but not others like actual data,
+ * sizes, etc.
+ */
+
+inline bool tensors_have_same_dim_order(
     const exec_aten::Tensor& a,
-    const exec_aten::Tensor& b);
+    const exec_aten::Tensor& b) {
+  exec_aten::Tensor tensor_list[2] = {a, b};
+  return tensors_have_same_dim_order(tensor_list);
+}
 
 /**
  * Asserts that three tensors have the same dim_order
  *
  * Note that this macro only tests dim order, but not others like actual data,
- * sizes, etc. Also this macro does not support ATen mode since we do not
- * support dim order in ATen mode.
+ * sizes, etc.
  *
- * TODO(T183094318): Add dim order and related function support for ATen mode.
  */
 
-bool tensors_have_same_dim_order(
+inline bool tensors_have_same_dim_order(
     const exec_aten::Tensor& a,
     const exec_aten::Tensor& b,
-    const exec_aten::Tensor& c);
+    const exec_aten::Tensor& c) {
+  exec_aten::Tensor tensor_list[3] = {a, b, c};
+  return tensors_have_same_dim_order(tensor_list);
+}
+
+/**
+ * Asserts that four tensors have the same dim_order
+ *
+ * Note that this macro only tests dim order, but not others like actual data,
+ * sizes, etc.
+ *
+ */
+
+inline bool tensors_have_same_dim_order(
+    const exec_aten::Tensor& a,
+    const exec_aten::Tensor& b,
+    const exec_aten::Tensor& c,
+    const exec_aten::Tensor& d) {
+  exec_aten::Tensor tensor_list[4] = {a, b, c, d};
+  return tensors_have_same_dim_order(tensor_list);
+}
 
 /**
  * Given an n-dimensional coordinate array and an array of tensor strides,
@@ -1232,6 +1279,7 @@ using ::executorch::runtime::tensor_is_bits_type;
 using ::executorch::runtime::tensor_is_bool_type;
 using ::executorch::runtime::tensor_is_complex_type;
 using ::executorch::runtime::tensor_is_contiguous;
+using ::executorch::runtime::tensor_is_default_dim_order;
 using ::executorch::runtime::tensor_is_default_or_channels_last_dim_order;
 using ::executorch::runtime::tensor_is_floating_type;
 using ::executorch::runtime::tensor_is_integral_type;