diff --git a/kernels/portable/cpu/op_flip.cpp b/kernels/portable/cpu/op_flip.cpp index 41ec6663714..8a381fc561f 100644 --- a/kernels/portable/cpu/op_flip.cpp +++ b/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 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* out_data = out.mutable_data_ptr(); + 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)]; } }); diff --git a/kernels/portable/cpu/op_permute_copy.cpp b/kernels/portable/cpu/op_permute_copy.cpp index 719f8fcb445..6343d1cfd70 100644 --- a/kernels/portable/cpu/op_permute_copy.cpp +++ b/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* const out_data = out.mutable_data_ptr(); + 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); diff --git a/kernels/portable/cpu/op_roll.cpp b/kernels/portable/cpu/op_roll.cpp index 4d314b3d191..c510aa90ac4 100644 --- a/kernels/portable/cpu/op_roll.cpp +++ b/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* out_data = out.mutable_data_ptr(); + 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)]; } }); diff --git a/kernels/portable/cpu/op_scatter_add.cpp b/kernels/portable/cpu/op_scatter_add.cpp index 690c31342a9..e1cb5d82f7f 100644 --- a/kernels/portable/cpu/op_scatter_add.cpp +++ b/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); } diff --git a/kernels/test/op_flip_test.cpp b/kernels/test/op_flip_test.cpp index be06e397be2..d7f34aab605 100644 --- a/kernels/test/op_flip_test.cpp +++ b/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 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)); +} diff --git a/kernels/test/op_gather_test.cpp b/kernels/test/op_gather_test.cpp index c1d6762af38..ede490f74d8 100644 --- a/kernels/test/op_gather_test.cpp +++ b/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 tf_index; + TensorFactory 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)); +} diff --git a/kernels/test/op_permute_copy_test.cpp b/kernels/test/op_permute_copy_test.cpp index 3273ea35481..33af440612e 100644 --- a/kernels/test/op_permute_copy_test.cpp +++ b/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 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 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)); +} diff --git a/kernels/test/op_roll_test.cpp b/kernels/test/op_roll_test.cpp index 4407e395db6..342c8e13503 100644 --- a/kernels/test/op_roll_test.cpp +++ b/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 shifts, + ArrayRef dims, + Tensor& out) { + return torch::executor::aten::roll_outf(context_, input, shifts, dims, out); + } +}; + +TEST_F(OpRollOutDimOrderTest, ChannelsLastMatchesContiguous) { + TensorFactory 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 shifts = {1}; + const std::vector 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(shifts.data(), shifts.size()), + ArrayRef(dims.data(), dims.size()), + out); + + EXPECT_TENSOR_CLOSE(out, tf.channels_last_like(expected)); +} diff --git a/kernels/test/op_scatter_add_test.cpp b/kernels/test/op_scatter_add_test.cpp index c259dae7420..ce067f9869b 100644 --- a/kernels/test/op_scatter_add_test.cpp +++ b/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 tf_index; + TensorFactory 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)); +} diff --git a/kernels/test/op_scatter_test.cpp b/kernels/test/op_scatter_test.cpp index 79f106e3578..d4634c86e5e 100644 --- a/kernels/test/op_scatter_test.cpp +++ b/kernels/test/op_scatter_test.cpp @@ -669,3 +669,38 @@ TEST_F(OpScatterSrcOutTest, InvalidOneDimInputAndZeroDimIndex) { } GENERATE_SCALAR_OVERFLOW_TESTS(OpScatterValueOutTest) + +TEST_F(OpScatterSrcOutTest, ChannelsLastMatchesContiguous) { + TensorFactory tf_index; + TensorFactory 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 tf_index; + TensorFactory 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)); +} diff --git a/runtime/core/exec_aten/util/tensor_util.h b/runtime/core/exec_aten/util/tensor_util.h index b9cbab4c1ef..61a05e4370f 100644 --- a/runtime/core/exec_aten/util/tensor_util.h +++ b/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(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(tensor.size(ii)); + const auto strides = tensor.strides(); + for (const auto d : c10::irange(tensor.dim())) { + trailing_dims_memo[d] = static_cast(strides[d]); } }