Index portable kernels by the tensor's strides instead of assuming row-major

kernels/portable/cpu/op_flip.cpp

@@ -21,22 +21,6 @@ bool check_flip_args(const Tensor& in, IntArrayRef dims, const Tensor& out) {
   return check_dim_list_is_valid(in, dims);
 }
 
-size_t unflip_flat_ix(size_t ix, const Tensor& in, ArrayRef<bool> flip_dim) {
-  size_t ix_coord[kTensorDimensionLimit];
-  indexToCoordinate(in, ix, ix_coord);
-
-  size_t unflip_coord[kTensorDimensionLimit];
-  for (const auto d : c10::irange(in.dim())) {
-    if (flip_dim[d]) {
-      unflip_coord[d] = in.size(d) - ix_coord[d] - 1;
-    } else {
-      unflip_coord[d] = ix_coord[d];
-    }
-  }
-
-  return coordinateToIndex(in, unflip_coord);
-}
-
 } // namespace
 
 Tensor& flip_out(
@@ -72,8 +56,22 @@ Tensor& flip_out(
     const CTYPE* in_data = in.const_data_ptr<CTYPE>();
     CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
 
+    const bool out_is_default = executorch::runtime::is_contiguous_dim_order(
+        out.dim_order().data(), out.dim_order().size());
+
     for (const auto ix : c10::irange(in.numel())) {
-      out_data[ix] = in_data[unflip_flat_ix(ix, in, flip_dim)];
+      // @lint-ignore CLANGTIDY facebook-hte-CArray
+      size_t coord[kTensorDimensionLimit];
+      indexToCoordinate(in, ix, coord);
+
+      // @lint-ignore CLANGTIDY facebook-hte-CArray
+      size_t src_coord[kTensorDimensionLimit];
+      for (const auto d : c10::irange(in.dim())) {
+        src_coord[d] = flip_dim[d] ? in.size(d) - coord[d] - 1 : coord[d];
+      }
+
+      out_data[out_is_default ? ix : coordinateToIndex(out, coord)] =
+          in_data[coordinateToIndex(in, src_coord)];
     }
   });
 

kernels/portable/cpu/op_permute_copy.cpp

@@ -71,8 +71,19 @@ Tensor& permute_copy_out(
     const CTYPE* const in_data = in.const_data_ptr<CTYPE>();
     CTYPE* const out_data = out.mutable_data_ptr<CTYPE>();
 
+    const bool out_is_default = executorch::runtime::is_contiguous_dim_order(
+        out.dim_order().data(), out.dim_order().size());
+
     for (const auto i : c10::irange(out.numel())) {
-      out_data[i] =
+      size_t out_ix = i;
+      if (!out_is_default) {
+        // @lint-ignore CLANGTIDY facebook-hte-CArray
+        size_t out_coord[kTensorDimensionLimit];
+        indexToCoordinate(out, i, out_coord);
+        out_ix = coordinateToIndex(out, out_coord);
+      }
+
+      out_data[out_ix] =
           in_data[executorch::runtime::coordinateToIndexWithTrailingDimsMemo(
               in, in_coord, trailing_dims_memo)];
       increment_coordinate_permuted(in, in_coord, dims);

kernels/portable/cpu/op_roll.cpp

@@ -32,19 +32,6 @@ bool check_roll_args(
   return true;
 }
 
-size_t unshift_flat_ix(size_t ix, const Tensor& in, IntArrayRef dim_shifts) {
-  size_t ix_coord[kTensorDimensionLimit];
-  indexToCoordinate(in, ix, ix_coord);
-
-  size_t shifted_coord[kTensorDimensionLimit];
-  for (const auto d : c10::irange(in.dim())) {
-    shifted_coord[d] =
-        (ix_coord[d] + in.size(d) - dim_shifts[d] % in.size(d)) % in.size(d);
-  }
-
-  return coordinateToIndex(in, shifted_coord);
-}
-
 } // namespace
 
 Tensor& roll_out(
@@ -86,8 +73,23 @@ Tensor& roll_out(
     const CTYPE* in_data = in.const_data_ptr<CTYPE>();
     CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
 
+    const bool out_is_default = executorch::runtime::is_contiguous_dim_order(
+        out.dim_order().data(), out.dim_order().size());
+
     for (const auto ix : c10::irange(out.numel())) {
-      out_data[ix] = in_data[unshift_flat_ix(ix, in, dim_shifts)];
+      // @lint-ignore CLANGTIDY facebook-hte-CArray
+      size_t coord[kTensorDimensionLimit];
+      indexToCoordinate(in, ix, coord);
+
+      // @lint-ignore CLANGTIDY facebook-hte-CArray
+      size_t shifted_coord[kTensorDimensionLimit];
+      for (const auto d : c10::irange(in.dim())) {
+        shifted_coord[d] =
+            (coord[d] + in.size(d) - dim_shifts[d] % in.size(d)) % in.size(d);
+      }
+
+      out_data[out_is_default ? ix : coordinateToIndex(out, coord)] =
+          in_data[coordinateToIndex(in, shifted_coord)];
     }
   });
 

kernels/portable/cpu/op_scatter_add.cpp

@@ -67,9 +67,6 @@ Tensor& scatter_add_out(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(self, src, out), InvalidArgument, out);
 
-  ET_KERNEL_CHECK(
-      ctx, tensor_is_default_dim_order(index), InvalidArgument, out);
-
   if (dim < 0) {
     dim += nonzero_dim(self);
   }

kernels/test/op_flip_test.cpp

@@ -68,3 +68,27 @@ TEST_F(OpFlipOutTest, SmokeTest2Dims) {
   op_flip_out(input, dims, out);
   EXPECT_TENSOR_CLOSE(out, out_expected);
 }
+
+class OpFlipOutDimOrderTest : public OperatorTest {
+ protected:
+  Tensor& op_flip_out(const Tensor& input, IntArrayRef dims, Tensor& out) {
+    return torch::executor::aten::flip_outf(context_, input, dims, out);
+  }
+};
+
+TEST_F(OpFlipOutDimOrderTest, ChannelsLastMatchesContiguous) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor contiguous_in =
+      tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor expected =
+      tf.make({1, 3, 2, 2}, {9, 10, 11, 12, 5, 6, 7, 8, 1, 2, 3, 4});
+  int64_t dims_data[1] = {1};
+  IntArrayRef dims = IntArrayRef(dims_data, 1);
+
+  Tensor in = tf.channels_last_like(contiguous_in);
+  Tensor out = tf.zeros_channels_last({1, 3, 2, 2});
+  op_flip_out(in, dims, out);
+
+  EXPECT_TENSOR_CLOSE(out, tf.channels_last_like(expected));
+}

kernels/test/op_gather_test.cpp

@@ -377,3 +377,19 @@ TEST_F(OpGatherOutTest, InvalidOneDimInputAndZeroDimIndex) {
   ET_EXPECT_KERNEL_FAILURE(
       context_, op_gather_out(self, 0, index, sparse_grad, out));
 }
+
+TEST_F(OpGatherOutTest, ChannelsLastMatchesContiguous) {
+  TensorFactory<ScalarType::Long> tf_index;
+  TensorFactory<ScalarType::Float> tf_data;
+
+  Tensor contiguous_in =
+      tf_data.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor expected = tf_data.make({1, 1, 2, 2}, {5, 6, 7, 8});
+
+  Tensor in = tf_data.channels_last_like(contiguous_in);
+  Tensor index = tf_index.channels_last_like(tf_index.full({1, 1, 2, 2}, 1));
+  Tensor out = tf_data.zeros_channels_last({1, 1, 2, 2});
+  op_gather_out(in, 1, index, false, out);
+
+  EXPECT_TENSOR_CLOSE(out, tf_data.channels_last_like(expected));
+}

kernels/test/op_permute_copy_test.cpp

@@ -475,3 +475,19 @@ TEST_F(OpPermuteCopyTest, DynamicShapeUnbound) {
   op_permute_copy_out(x, perm_aref, out);
   EXPECT_TENSOR_EQ(out, expected);
 }
+
+TEST_F(OpPermuteCopyTest, ChannelsLastMatchesContiguous) {
+  TensorFactory<ScalarType::Float> tf;
+
+  // in[0][c][h][w] laid out channels-last, permuted to (0, 2, 3, 1)
+  Tensor in = tf.channels_last_like(
+      tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+  Tensor expected =
+      tf.make({1, 2, 2, 3}, {1, 5, 9, 2, 6, 10, 3, 7, 11, 4, 8, 12});
+  Tensor out = tf.zeros_channels_last({1, 2, 2, 3});
+  const std::vector<int64_t> dims = {0, 2, 3, 1};
+
+  op_permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out);
+
+  EXPECT_TENSOR_CLOSE(out, tf.channels_last_like(expected));
+}

kernels/test/op_roll_test.cpp

@@ -60,3 +60,35 @@ TEST_F(OpRollOutTest, SmokeTest) {
   ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
+
+class OpRollOutDimOrderTest : public OperatorTest {
+ protected:
+  Tensor& op_roll_out(
+      const Tensor& input,
+      ArrayRef<int64_t> shifts,
+      ArrayRef<int64_t> dims,
+      Tensor& out) {
+    return torch::executor::aten::roll_outf(context_, input, shifts, dims, out);
+  }
+};
+
+TEST_F(OpRollOutDimOrderTest, ChannelsLastMatchesContiguous) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor contiguous_in =
+      tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor expected =
+      tf.make({1, 3, 2, 2}, {9, 10, 11, 12, 1, 2, 3, 4, 5, 6, 7, 8});
+  const std::vector<int64_t> shifts = {1};
+  const std::vector<int64_t> dims = {1};
+
+  Tensor in = tf.channels_last_like(contiguous_in);
+  Tensor out = tf.zeros_channels_last({1, 3, 2, 2});
+  op_roll_out(
+      in,
+      ArrayRef<int64_t>(shifts.data(), shifts.size()),
+      ArrayRef<int64_t>(dims.data(), dims.size()),
+      out);
+
+  EXPECT_TENSOR_CLOSE(out, tf.channels_last_like(expected));
+}

kernels/test/op_scatter_add_test.cpp

@@ -402,3 +402,21 @@ TEST_F(OpScatterAddOutTest, DynamicShapeUnbound) {
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
+
+TEST_F(OpScatterAddOutTest, ChannelsLastMatchesContiguous) {
+  TensorFactory<ScalarType::Long> tf_index;
+  TensorFactory<ScalarType::Float> tf_data;
+
+  Tensor contiguous_in =
+      tf_data.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor expected =
+      tf_data.make({1, 3, 2, 2}, {1, 2, 3, 4, 15, 16, 17, 18, 9, 10, 11, 12});
+
+  Tensor self = tf_data.channels_last_like(contiguous_in);
+  Tensor index = tf_index.channels_last_like(tf_index.full({1, 1, 2, 2}, 1));
+  Tensor src = tf_data.channels_last_like(tf_data.full({1, 1, 2, 2}, 10));
+  Tensor out = tf_data.zeros_channels_last({1, 3, 2, 2});
+  op_scatter_add_out(self, 1, index, src, out);
+
+  EXPECT_TENSOR_CLOSE(out, tf_data.channels_last_like(expected));
+}

kernels/test/op_scatter_test.cpp

@@ -669,3 +669,38 @@ TEST_F(OpScatterSrcOutTest, InvalidOneDimInputAndZeroDimIndex) {
 }
 
 GENERATE_SCALAR_OVERFLOW_TESTS(OpScatterValueOutTest)
+
+TEST_F(OpScatterSrcOutTest, ChannelsLastMatchesContiguous) {
+  TensorFactory<ScalarType::Long> tf_index;
+  TensorFactory<ScalarType::Float> tf_data;
+
+  Tensor contiguous_in =
+      tf_data.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor expected =
+      tf_data.make({1, 3, 2, 2}, {1, 2, 3, 4, 99, 99, 99, 99, 9, 10, 11, 12});
+
+  Tensor self = tf_data.channels_last_like(contiguous_in);
+  Tensor index = tf_index.channels_last_like(tf_index.full({1, 1, 2, 2}, 1));
+  Tensor src = tf_data.channels_last_like(tf_data.full({1, 1, 2, 2}, 99));
+  Tensor out = tf_data.zeros_channels_last({1, 3, 2, 2});
+  op_scatter_src_out(self, 1, index, src, out);
+
+  EXPECT_TENSOR_CLOSE(out, tf_data.channels_last_like(expected));
+}
+
+TEST_F(OpScatterValueOutTest, ChannelsLastMatchesContiguous) {
+  TensorFactory<ScalarType::Long> tf_index;
+  TensorFactory<ScalarType::Float> tf_data;
+
+  Tensor contiguous_in =
+      tf_data.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor expected =
+      tf_data.make({1, 3, 2, 2}, {1, 2, 3, 4, 99, 99, 99, 99, 9, 10, 11, 12});
+
+  Tensor self = tf_data.channels_last_like(contiguous_in);
+  Tensor index = tf_index.channels_last_like(tf_index.full({1, 1, 2, 2}, 1));
+  Tensor out = tf_data.zeros_channels_last({1, 3, 2, 2});
+  op_scatter_value_out(self, 1, index, 99.0, out);
+
+  EXPECT_TENSOR_CLOSE(out, tf_data.channels_last_like(expected));
+}

runtime/core/exec_aten/util/tensor_util.h

@@ -981,8 +981,9 @@ inline size_t coordinateToIndex(
     const executorch::aten::Tensor& tensor,
     const size_t* const coordinate) {
   size_t index = 0;
+  const auto strides = tensor.strides();
   for (const auto d : c10::irange(tensor.dim())) {
-    index += coordinate[d] * getTrailingDims(tensor, d);
+    index += coordinate[d] * static_cast<size_t>(strides[d]);
   }
   return index;
 }
@@ -995,11 +996,9 @@ inline size_t coordinateToIndex(
 inline void memoizeTrailingDims(
     const executorch::aten::Tensor& tensor,
     size_t trailing_dims_memo[kTensorDimensionLimit]) {
-  const auto tensorDim = tensor.dim();
-  size_t dims = 1;
-  for (int ii = tensorDim - 1; ii >= 0; --ii) {
-    trailing_dims_memo[ii] = dims;
-    dims *= static_cast<size_t>(tensor.size(ii));
+  const auto strides = tensor.strides();
+  for (const auto d : c10::irange(tensor.dim())) {
+    trailing_dims_memo[d] = static_cast<size_t>(strides[d]);
   }
 }