| /* |
| * Copyright (c) Meta Platforms, Inc. and affiliates. |
| * All rights reserved. |
| * |
| * This source code is licensed under the BSD-style license found in the |
| * LICENSE file in the root directory of this source tree. |
| */ |
| |
| #include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator |
| #include <executorch/kernels/test/TestUtil.h> |
| #include <executorch/kernels/test/supported_features.h> |
| #include <executorch/runtime/core/exec_aten/exec_aten.h> |
| #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h> |
| #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h> |
| #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h> |
| |
| #include <gtest/gtest.h> |
| #include <cstdint> |
| |
| using namespace ::testing; |
| using exec_aten::ArrayRef; |
| using exec_aten::ScalarType; |
| using exec_aten::Tensor; |
| using exec_aten::TensorList; |
| using torch::executor::testing::TensorFactory; |
| |
| class OpStackOutTest : public OperatorTest { |
| protected: |
| Tensor& op_stack_out(TensorList tensors, int64_t dim, Tensor& out) { |
| return torch::executor::aten::stack_outf(context_, tensors, dim, out); |
| } |
| |
| template <class CTYPE, exec_aten::ScalarType DTYPE> |
| void test_dtype() { |
| TensorFactory<DTYPE> tf; |
| |
| // Will be stackd along out.dim(1). Use different input values so we can see |
| // where each output value came from. |
| Tensor x = tf.ones({3, 4}); |
| Tensor y = tf.zeros({3, 4}); |
| std::vector<Tensor> inputs = {x, y}; |
| |
| Tensor out = tf.ones({3, 2, 4}); |
| op_stack_out( |
| ArrayRef<Tensor>(inputs.data(), inputs.size()), /*dim=*/1, out); |
| |
| // The two tensors x and y are stacked along the 1st dimension with the |
| // order [x, y], so the x and y should be equal to expected[:, 0, :] and |
| // expected[:, 1, :] e.g. expected[i, 0, j] = x[i, j] and expected[i, 1, j] |
| // = y[i, j] for any i in [-x.size(0), x.size(0)-1] and j in [-x.size(1), |
| // x.size(1)-1] |
| // clang-format off |
| Tensor expected = tf.make( |
| {3, 2, 4}, |
| { |
| // all ones below are from x, |
| // and all zeros are from y. |
| // [0, :, :] |
| 1, 1, 1, 1, // [0, 0, :] |
| 0, 0, 0, 0, // [0, 1, :] |
| |
| // [1, :, :] |
| 1, 1, 1, 1, // [1, 0, :] |
| 0, 0, 0, 0, // [1, 1, :] |
| |
| // [2, :, :] |
| 1, 1, 1, 1, // [2, 0, :] |
| 0, 0, 0, 0, // [2, 1, :] |
| }); |
| // clang-format on |
| |
| EXPECT_TENSOR_EQ(out, expected); |
| } |
| |
| // Running stacking experiments along given dim. |
| void run_stack_tests( |
| const std::vector<Tensor>& inputs, |
| int64_t dim, |
| const Tensor& expected) { |
| ArrayRef<Tensor> inputs_array(inputs.data(), inputs.size()); |
| |
| TensorFactory<ScalarType::Double> tf; |
| const std::vector<int32_t> out_size( |
| expected.sizes().begin(), expected.sizes().end()); |
| Tensor out = tf.zeros(out_size); |
| |
| // Should always return the provided out Tensor. |
| Tensor ret = op_stack_out(inputs_array, dim, out); |
| EXPECT_TENSOR_EQ(out, ret); |
| EXPECT_TENSOR_EQ(out, expected); |
| |
| ret = op_stack_out(inputs_array, /*dim=*/dim - out.dim(), out); |
| EXPECT_TENSOR_EQ(out, ret); |
| EXPECT_TENSOR_EQ(out, expected); |
| } |
| }; |
| |
| TEST_F(OpStackOutTest, InsertFront) { |
| TensorFactory<ScalarType::Double> tf; |
| |
| // clang-format off |
| Tensor x = tf.make( |
| {3, 4}, |
| { |
| 1., 2., 3., 4., // [0, :] |
| 5., 6., 7., 8., // [1, :] |
| 9., 10., 11., 12., // [2, :] |
| |
| }); |
| Tensor y = tf.make( |
| {3, 4}, |
| { |
| -1., -2., -3., -4., // [0, :] |
| -5., -6., -7., -8., // [1, :] |
| -9., -10., -11., -12., // [2, :] |
| }); |
| // clang-format on |
| |
| const std::vector<Tensor> inputs = {x, y}; |
| |
| // Try to stack the two tensors in the front (dim = 0) |
| // The size of expected output tensor should follow the below rules: |
| // - For any input tensor, its size(i) == output.size(i) if i < dim, and its |
| // size(i) == output.size(i+1) if i >= dim |
| // - For the stack dimension (output[dim]), its size should be the number of |
| // input tensors |
| std::vector<int32_t> expected_size = {2, 3, 4}; |
| |
| // The two tensors x and y are stacked along the 0th dimension with the order |
| // [x, y], so the x and y should be equal to expected[0, :, :] and expected[1, |
| // :, :] e.g. expected[0, i, j] = x[i, j] and expected[1, i, j] = y[i, j] for |
| // any i in [-x.size(0), x.size(0)-1] and j in [-x.size(1), x.size(1)-1] |
| // clang-format off |
| Tensor expected = tf.make( |
| expected_size, |
| { |
| // [0, :, :] equals to x |
| 1., 2., 3., 4., // [0, 0, :] |
| 5., 6., 7., 8., // [0, 1, :] |
| 9., 10., 11., 12., // [0, 2, :] |
| |
| // [1, :, :] equals to y |
| -1., -2., -3., -4., // [1, 0, :] |
| -5., -6., -7., -8., // [1, 1, :] |
| -9., -10., -11., -12., // [1, 2, :] |
| }); |
| // clang-format on |
| |
| run_stack_tests(inputs, /*dim=*/0, expected); |
| } |
| |
| TEST_F(OpStackOutTest, InsertMiddle) { |
| TensorFactory<ScalarType::Double> tf; |
| |
| // Two tensors with same size. Stack them on multiple dimensions |
| // clang-format off |
| Tensor x = tf.make( |
| {3, 4}, |
| { |
| 1., 2., 3., 4., // [0, :] |
| 5., 6., 7., 8., // [1, :] |
| 9., 10., 11., 12., // [2, :] |
| |
| }); |
| Tensor y = tf.make( |
| {3, 4}, |
| { |
| -1., -2., -3., -4., // [0, :] |
| -5., -6., -7., -8., // [1, :] |
| -9., -10., -11., -12., // [2, :] |
| }); |
| // clang-format on |
| |
| const std::vector<Tensor> inputs = {x, y}; |
| |
| // Try to stack the two tensors in the middle (dim = 1) |
| // The size of expected output tensor should follow the below rules: |
| // - For any input tensor, its size(i) == output.size(i) if i < dim, and its |
| // size(i) == output.size(i+1) if i >= dim |
| // - For the stack dimension (output[dim]), its size should be the number of |
| // input tensors |
| std::vector<int32_t> expected_size = {3, 2, 4}; |
| |
| // The two tensors x and y are stacked along the 1st dimension with the order |
| // [x, y], so the x and y should be equal to expected[:, 0, :] and expected[:, |
| // 1, :] e.g. expected[i, 0, j] = x[i, j] and expected[i, 1, j] = y[i, j] for |
| // any i in [-x.size(0), x.size(0)-1] and j in [-x.size(1), x.size(1)-1] |
| // clang-format off |
| Tensor expected = tf.make( |
| expected_size, |
| { |
| // [0, :, :] |
| 1., 2., 3., 4., // [0, 0, :] = x[0, :] |
| -1., -2., -3., -4., // [0, 1, :] = y[0, :] |
| |
| // [1, :, :] |
| 5., 6., 7., 8., // [1, 0, :] = x[1, :] |
| -5., -6., -7., -8., // [1, 1, :] = y[1, :] |
| |
| // [2, :, :] |
| 9., 10., 11., 12., // [2, 0, :] = x[2, :] |
| -9., -10., -11., -12., // [2, 1, :] = y[2, :] |
| }); |
| // clang-format on |
| |
| run_stack_tests(inputs, /*dim=*/1, expected); |
| } |
| |
| TEST_F(OpStackOutTest, InsertEnd) { |
| TensorFactory<ScalarType::Double> tf; |
| |
| // Two tensors with same size. Stack them on multiple dimensions |
| // clang-format off |
| Tensor x = tf.make( |
| {3, 4}, |
| { |
| 1., 2., 3., 4., // [0, :] |
| 5., 6., 7., 8., // [1, :] |
| 9., 10., 11., 12., // [2, :] |
| |
| }); |
| Tensor y = tf.make( |
| {3, 4}, |
| { |
| -1., -2., -3., -4., // [0, :] |
| -5., -6., -7., -8., // [1, :] |
| -9., -10., -11., -12., // [2, :] |
| }); |
| // clang-format on |
| |
| const std::vector<Tensor> inputs = {x, y}; |
| |
| // Try to stack the two tensors at the end (dim = 2) |
| // The size of expected output tensor should follow the below rules: |
| // - For any input tensor, its size(i) == output.size(i) if i < dim, and its |
| // size(i) == output.size(i+1) if i >= dim |
| // - For the stack dimension (output[dim]), its size should be the number of |
| // input tensors |
| std::vector<int32_t> expected_size = {3, 4, 2}; |
| |
| // The two tensors x and y are stacked along the 2nd dimension with the order |
| // [x, y], so the x and y should be equal to expected[:, :, 0] and expected[:, |
| // :, 1] e.g. expected[i, j, 0] = x[i, j] and expected[i, j, 1] = y[i, j] for |
| // any i in [-x.size(0), x.size(0)-1] and j in [-x.size(1), x.size(1)-1] |
| // clang-format off |
| Tensor expected = tf.make( |
| expected_size, |
| { |
| // All values in the first column are from x, |
| // and the second column are from y |
| |
| // [0, :, :] |
| 1., -1., // [0, 0, :] |
| 2., -2., // [0, 1, :] |
| 3., -3., // [0, 2, :] |
| 4., -4., // [0, 3, :] |
| |
| // [1, :, :] |
| 5., -5., // [1, 0, :] |
| 6., -6., // [1, 1, :] |
| 7., -7., // [1, 2, :] |
| 8., -8., // [1, 3, :] |
| |
| // [2, :, :] |
| 9., -9., // [2, 0, :] |
| 10., -10., // [2, 1, :] |
| 11., -11., // [2, 2, :] |
| 12., -12., // [2, 3, :] |
| }); |
| // clang-format on |
| |
| run_stack_tests(inputs, /*dim=*/2, expected); |
| } |
| |
| /// A generic smoke test that works for any dtype that supports ones() and |
| /// zeros(). |
| TEST_F(OpStackOutTest, AllDtypesSupported) { |
| #define TEST_ENTRY(ctype, dtype) test_dtype<ctype, ScalarType::dtype>(); |
| ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY); |
| #undef TEST_ENTRY |
| // TODO: Also add tests for half, complex, quantized, and other types. Easiest |
| // way to do that would be to make TensorFactory support zeros() and ones() |
| // for those types. |
| } |
| |
| TEST_F(OpStackOutTest, NoInputTensorsWithEmptyOutTensorFails) { |
| TensorFactory<ScalarType::Int> tf; |
| |
| // Make an empty out tensor and demonstrate that it's empty. |
| Tensor out = tf.make({0}, {}); |
| EXPECT_EQ(out.numel(), 0); |
| |
| // Pass an empty list of input tensors. |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, op_stack_out(ArrayRef<Tensor>(), /*dim=*/0, out)); |
| } |
| |
| TEST_F(OpStackOutTest, AllEmptyInputTensors) { |
| TensorFactory<ScalarType::Int> tf; |
| |
| // Using empty tensors as input. |
| Tensor empty = tf.make({0, 10, 3}, {}); |
| EXPECT_EQ(empty.numel(), 0); |
| std::vector<Tensor> inputs = {empty, empty, empty}; |
| |
| Tensor x = tf.ones({2, 2}); |
| |
| // Output whose shape is appropriate for stacking along out.dim(0). |
| Tensor out = tf.make({3, 0, 10, 3}, {}); |
| EXPECT_EQ(out.numel(), 0); |
| |
| Tensor ret = op_stack_out( |
| ArrayRef<Tensor>(inputs.data(), inputs.size()), /*dim=*/0, out); |
| EXPECT_EQ(ret.numel(), 0); |
| // Success if it doesn't assert on the weird-shaped empty input and the |
| // empty_out is still a empty array |
| } |
| |
| TEST_F(OpStackOutTest, DimOutOfBoundDies) { |
| TensorFactory<ScalarType::Int> tf; |
| |
| // Stack a single tensor with size [1, 1]. The size of output would always be |
| // [1, 1, 1] no matter stack on which dimension. |
| Tensor x = tf.ones({1, 1}); |
| ArrayRef<Tensor> inputs(&x, 1); |
| |
| Tensor out = tf.zeros({1, 1, 1}); |
| |
| // Some invalid dim values. |
| const std::vector<int64_t> invalid_dims = {3, 4, 5, -4, -5, -6}; |
| for (int64_t dim : invalid_dims) { |
| ET_EXPECT_KERNEL_FAILURE(context_, op_stack_out(inputs, dim, out)); |
| } |
| } |
| |
| TEST_F(OpStackOutTest, MismatchedDtypesDies) { |
| TensorFactory<ScalarType::Int> tf_int; |
| TensorFactory<ScalarType::Float> tf_float; |
| Tensor out = tf_int.zeros({1, 2, 2}); |
| |
| // Size is compatible to the output, but a mismatched dtype. |
| std::vector<Tensor> inputs = {tf_float.ones({2, 2})}; |
| |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, |
| op_stack_out( |
| ArrayRef<Tensor>(inputs.data(), inputs.size()), /*dim=*/0, out)); |
| } |
| |
| TEST_F(OpStackOutTest, OutMatchNumelWithExtraDimAtEndDies) { |
| if (torch::executor::testing::SupportedFeatures::get()->is_aten) { |
| GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions"; |
| } |
| TensorFactory<ScalarType::Int> tf; |
| Tensor out = tf.zeros({1, 2, 2, 1}); |
| |
| // Same dtype and numel as the output, but a mixmatched size (output.dim() |
| // should always one greater than input.dim()) |
| std::vector<Tensor> inputs = {tf.ones({2, 2})}; |
| |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, |
| op_stack_out( |
| ArrayRef<Tensor>(inputs.data(), inputs.size()), /*dim=*/0, out)); |
| } |
| |
| TEST_F(OpStackOutTest, OutMatchNumelLackDimAtFrontDies) { |
| if (torch::executor::testing::SupportedFeatures::get()->is_aten) { |
| GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions"; |
| } |
| TensorFactory<ScalarType::Int> tf; |
| Tensor out = tf.zeros({2, 2}); |
| |
| // Same dtype and numel as the output, but a mixmatched size (output.dim() |
| // should always one greater than input.dim()) |
| std::vector<Tensor> inputs = {tf.ones({2, 2})}; |
| |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, |
| op_stack_out( |
| ArrayRef<Tensor>(inputs.data(), inputs.size()), /*dim=*/0, out)); |
| } |
| |
| TEST_F(OpStackOutTest, OutRegularMismatchDimDies) { |
| if (torch::executor::testing::SupportedFeatures::get()->is_aten) { |
| GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions"; |
| } |
| TensorFactory<ScalarType::Int> tf; |
| |
| // Should be {2, 2, 3} to match the inputs when calling stack() with dim 0. |
| Tensor out = tf.zeros({2, 4, 5}); |
| |
| std::vector<Tensor> inputs = { |
| tf.ones({2, 3}), |
| tf.ones({2, 3}), |
| }; |
| |
| ET_EXPECT_KERNEL_FAILURE( |
| context_, |
| op_stack_out( |
| ArrayRef<Tensor>(inputs.data(), inputs.size()), /*dim=*/0, out)); |
| } |
| |
| /* %python |
| import torch |
| torch.manual_seed(0) |
| x = [torch.randint(10, (2, 3)), |
| torch.randint(10, (2, 3)), |
| torch.randint(10, (2, 3)), |
| torch.randint(10, (2, 3))] |
| res = torch.stack(x, 0) |
| op = "op_stack_out" |
| opt_extra_params = "0," |
| dtype = "ScalarType::Int" |
| check = "EXPECT_TENSOR_EQ" */ |
| |
| TEST_F(OpStackOutTest, DynamicShapeUpperBoundSameAsExpected) { |
| /* %python |
| out_args = "{4, 2, 3}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND" |
| %rewrite(unary_op_tensor_list_in) */ |
| |
| TensorFactory<ScalarType::Int> tf; |
| |
| std::vector<Tensor> xv = { |
| tf.make({2, 3}, {4, 9, 3, 0, 3, 9}), |
| tf.make({2, 3}, {7, 3, 7, 3, 1, 6}), |
| tf.make({2, 3}, {6, 9, 8, 6, 6, 8}), |
| tf.make({2, 3}, {4, 3, 6, 9, 1, 4})}; |
| TensorList x(xv.data(), xv.size()); |
| Tensor expected = tf.make({4, 2, 3}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6, |
| 6, 9, 8, 6, 6, 8, 4, 3, 6, 9, 1, 4}); |
| |
| Tensor out = |
| tf.zeros({4, 2, 3}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); |
| op_stack_out(x, 0, out); |
| EXPECT_TENSOR_EQ(out, expected); |
| } |
| |
| TEST_F(OpStackOutTest, DynamicShapeUpperBoundLargerThanExpected) { |
| if (!torch::executor::testing::SupportedFeatures::get()->output_resize) { |
| GTEST_SKIP() << "Dynamic shape not supported"; |
| } |
| /* %python |
| out_args = "{5, 5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND" |
| %rewrite(unary_op_tensor_list_in) */ |
| |
| TensorFactory<ScalarType::Int> tf; |
| |
| std::vector<Tensor> xv = { |
| tf.make({2, 3}, {4, 9, 3, 0, 3, 9}), |
| tf.make({2, 3}, {7, 3, 7, 3, 1, 6}), |
| tf.make({2, 3}, {6, 9, 8, 6, 6, 8}), |
| tf.make({2, 3}, {4, 3, 6, 9, 1, 4})}; |
| TensorList x(xv.data(), xv.size()); |
| Tensor expected = tf.make({4, 2, 3}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6, |
| 6, 9, 8, 6, 6, 8, 4, 3, 6, 9, 1, 4}); |
| |
| Tensor out = |
| tf.zeros({5, 5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); |
| op_stack_out(x, 0, out); |
| EXPECT_TENSOR_EQ(out, expected); |
| } |
| |
| TEST_F(OpStackOutTest, DynamicShapeUnbound) { |
| if (!torch::executor::testing::SupportedFeatures::get()->output_resize) { |
| GTEST_SKIP() << "Dynamic shape not supported"; |
| } |
| /* %python |
| out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND" |
| %rewrite(unary_op_tensor_list_in) */ |
| |
| TensorFactory<ScalarType::Int> tf; |
| |
| std::vector<Tensor> xv = { |
| tf.make({2, 3}, {4, 9, 3, 0, 3, 9}), |
| tf.make({2, 3}, {7, 3, 7, 3, 1, 6}), |
| tf.make({2, 3}, {6, 9, 8, 6, 6, 8}), |
| tf.make({2, 3}, {4, 3, 6, 9, 1, 4})}; |
| TensorList x(xv.data(), xv.size()); |
| Tensor expected = tf.make({4, 2, 3}, {4, 9, 3, 0, 3, 9, 7, 3, 7, 3, 1, 6, |
| 6, 9, 8, 6, 6, 8, 4, 3, 6, 9, 1, 4}); |
| |
| Tensor out = tf.zeros( |
| {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND); |
| op_stack_out(x, 0, out); |
| EXPECT_TENSOR_EQ(out, expected); |
| } |
