Reject non-default dim order in the portable kernels that copy by block

kernels/portable/cpu/op_cat.cpp

@@ -28,6 +28,12 @@ Tensor& cat_out(
 
   ET_KERNEL_CHECK(ctx, check_cat_args(tensors, dim, out), InvalidArgument, out);
 
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(out), InvalidArgument, out);
+  for (size_t i = 0; i < tensors.size(); ++i) {
+    ET_KERNEL_CHECK(
+        ctx, tensor_is_default_dim_order(tensors[i]), InvalidArgument, out);
+  }
+
   Tensor::SizesType expected_out_size[kTensorDimensionLimit];
   size_t expected_out_dim = 0;
   get_cat_out_target_size(tensors, dim, expected_out_size, &expected_out_dim);

kernels/portable/cpu/op_constant_pad_nd.cpp

@@ -252,6 +252,8 @@ Tensor& constant_pad_nd_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);
+
   // resize out tensor for dynamic shapes
   ET_KERNEL_CHECK_MSG(
       ctx,

kernels/portable/cpu/op_cumsum.cpp

@@ -103,6 +103,8 @@ Tensor& cumsum_out(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(self, out), InvalidArgument, out);
 
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(self), InvalidArgument, out);
+
   ET_KERNEL_CHECK(
       ctx, resize_tensor(out, self.sizes()) == Error::Ok, InvalidArgument, out);
 

kernels/portable/cpu/op_pixel_unshuffle.cpp

@@ -81,6 +81,11 @@ Tensor& pixel_unshuffle_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);
+
   // @lint-ignore CLANGTIDY facebook-hte-CArray
   Tensor::SizesType expected_out_size[kTensorDimensionLimit];
   size_t expected_out_dim = 0;

kernels/portable/cpu/op_slice_scatter.cpp

@@ -44,6 +44,9 @@ Tensor& slice_scatter_out(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(input, out), InvalidArgument, out);
 
+  ET_KERNEL_CHECK(
+      ctx, tensor_is_default_dim_order(input), InvalidArgument, out);
+
   if (input.numel() == 0) {
     return out;
   }

kernels/portable/cpu/op_split_copy.cpp

@@ -49,6 +49,8 @@ void split_copy_Tensor_out(
   for (size_t i = 0; i < out.size(); ++i) {
     ET_KERNEL_CHECK(
         ctx, tensors_have_same_dim_order(input, out[i]), InvalidArgument, );
+    ET_KERNEL_CHECK(
+        ctx, tensor_is_default_dim_order(out[i]), InvalidArgument, );
   }
 
   const size_t leading_dims = getLeadingDims(input, dim);

kernels/portable/cpu/op_split_with_sizes_copy.cpp

@@ -43,6 +43,8 @@ void split_with_sizes_copy_out(
   for (const auto i : c10::irange(out.size())) {
     ET_KERNEL_CHECK(
         ctx, tensors_have_same_dim_order(in, out[i]), InvalidArgument, );
+    ET_KERNEL_CHECK(
+        ctx, tensor_is_default_dim_order(out[i]), InvalidArgument, );
   }
 
   // If out is empty, then nothing needs to be done after checking the args.

kernels/portable/cpu/op_topk.cpp

@@ -170,6 +170,14 @@ std::tuple<Tensor&, Tensor&> topk_values(
   ET_KERNEL_CHECK(
       ctx, check_topk_args(in, k, dim, values, indices), InvalidArgument, out);
 
+  ET_KERNEL_CHECK(
+      ctx,
+      tensors_have_same_dim_order(in, values, indices),
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
+
   if (dim < 0) {
     dim += nonzero_dim(in);
   }

kernels/test/CMakeLists.txt

@@ -261,6 +261,7 @@ set(all_test_sources
     "op_pdist_forward_test.cpp"
     "op_permute_copy_test.cpp"
     "op_pixel_shuffle_test.cpp"
+    "op_pixel_unshuffle_test.cpp"
     "op_prod_test.cpp"
     "op_rand_test.cpp"
     "op_randn_test.cpp"

kernels/test/op_cat_test.cpp

@@ -464,3 +464,17 @@ TEST_F(OpCatOutTest, DynamicShapeUnbound) {
   op_cat_out(x, 0, out);
   EXPECT_TENSOR_EQ(out, expected);
 }
+
+TEST_F(OpCatOutTest, NonDefaultDimOrderDies) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor a = tf.channels_last_like(
+      tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+  Tensor b = tf.channels_last_like(
+      tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+  Tensor out = tf.zeros_channels_last({1, 6, 2, 2});
+  std::vector<Tensor> inputs = {a, b};
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_, op_cat_out(TensorList(inputs.data(), inputs.size()), 1, out));
+}

kernels/test/op_constant_pad_nd_test.cpp

@@ -484,3 +484,17 @@ TEST_F(OpConstantPadNDOutTest, IncorrectOutputShapeFail) {
 }
 
 GENERATE_SCALAR_OVERFLOW_TESTS(OpConstantPadNDOutTest)
+
+TEST_F(OpConstantPadNDOutTest, NonDefaultDimOrderDies) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor self = tf.channels_last_like(
+      tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+  Tensor out = tf.zeros_channels_last({1, 3, 2, 4});
+  const std::vector<int64_t> padding = {1, 1};
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_,
+      op_constant_pad_nd_out(
+          self, IntArrayRef(padding.data(), padding.size()), 0.0, out));
+}

kernels/test/op_cumsum_test.cpp

@@ -286,3 +286,14 @@ TEST_F(OpCumSumOutTest, DISABLED_DynamicShapeUnbound) {
   Tensor ret = op_cumsum_out(x, 1, ScalarType::Float, out);
   EXPECT_TENSOR_CLOSE(out, expected_result);
 }
+
+TEST_F(OpCumSumOutTest, NonDefaultDimOrderDies) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor self = tf.channels_last_like(
+      tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+  Tensor out = tf.zeros_channels_last({1, 3, 2, 2});
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_, op_cumsum_out(self, 1, ScalarType::Float, out));
+}

kernels/test/op_pixel_unshuffle_test.cpp

@@ -126,3 +126,13 @@ TEST_F(OpPixelUnshuffleOutTest, NegativeUpscaleFactorDies) {
   // Using a negative upscale factor should exit with an error code.
   ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_unshuffle_out(a, -3, out));
 }
+
+TEST_F(OpPixelUnshuffleOutTest, NonDefaultDimOrderDies) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor in = tf.channels_last_like(tf.make(
+      {1, 1, 4, 4}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}));
+  Tensor out = tf.zeros_channels_last({1, 4, 2, 2});
+
+  ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_unshuffle_out(in, 2, out));
+}

kernels/test/op_slice_scatter_test.cpp

@@ -884,3 +884,15 @@ TEST_F(OpSliceScatterTensorOutTest, LargeEndValue) {
   EXPECT_TENSOR_EQ(ret, out);
   EXPECT_TENSOR_EQ(ret, expected);
 }
+
+TEST_F(OpSliceScatterTensorOutTest, NonDefaultDimOrderDies) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor self = tf.channels_last_like(
+      tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+  Tensor src = tf.zeros_channels_last({1, 1, 2, 2});
+  Tensor out = tf.zeros_channels_last({1, 3, 2, 2});
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_, op_slice_scatter_out(self, src, 1, 0, 1, 1, out));
+}

kernels/test/op_split_copy_test.cpp

@@ -576,3 +576,18 @@ TEST_F(OpSplitCopyTensorOutTest, DISABLED_DynamicShapeUnbound) {
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
+
+TEST_F(OpSplitCopyTensorOutTest, NonDefaultDimOrderDies) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor self = tf.channels_last_like(tf.make(
+      {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}));
+  std::vector<Tensor> outs = {
+      tf.zeros_channels_last({1, 2, 2, 2}),
+      tf.zeros_channels_last({1, 2, 2, 2})};
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_,
+      op_split_copy_tensor_out(
+          self, 2, 1, TensorList(outs.data(), outs.size())));
+}

kernels/test/op_split_with_sizes_copy_test.cpp

@@ -115,3 +115,24 @@ TEST_F(OpSplitWithSizesCopyOutTest, DynamicShape) {
   test_tensor_shape_dynamism(
       executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND);
 }
+
+TEST_F(OpSplitWithSizesCopyOutTest, NonDefaultDimOrderDies) {
+  torch::executor::testing::TensorFactory<executorch::aten::ScalarType::Float>
+      tf;
+
+  executorch::aten::Tensor self = tf.channels_last_like(tf.make(
+      {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}));
+  std::vector<executorch::aten::Tensor> outs = {
+      tf.zeros_channels_last({1, 2, 2, 2}),
+      tf.zeros_channels_last({1, 2, 2, 2})};
+  const std::vector<int64_t> split_sizes = {2, 2};
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_,
+      op_split_with_sizes_copy_out(
+          self,
+          executorch::aten::ArrayRef<int64_t>(
+              split_sizes.data(), split_sizes.size()),
+          1,
+          executorch::aten::TensorList(outs.data(), outs.size())));
+}

kernels/test/op_topk_test.cpp

@@ -173,3 +173,21 @@ TEST_F(OpTopkValuesTest, NonPartialSort) {
     EXPECT_TENSOR_EQ(indices, indices_expected);
   }
 }
+
+TEST_F(OpTopkValuesTest, NonDefaultDimOrderDies) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tf_long;
+
+  Tensor in = tf.channels_last_like(tf.make(
+      {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}));
+  Tensor values = tf.zeros_channels_last({1, 2, 2, 2});
+  Tensor indices = tf_long.zeros_channels_last({1, 2, 2, 2});
+
+  TempMemoryAllocator allocator = TempMemoryAllocator();
+  executorch::ET_RUNTIME_NAMESPACE::KernelRuntimeContext context(
+      nullptr, &allocator);
+  torch::executor::aten::topk_outf(
+      context, in, 2, 1, true, true, values, indices);
+
+  EXPECT_NE(context.failure_state(), torch::executor::Error::Ok);
+}