diff --git a/kernels/aten/cpu/op__empty_dim_order.cpp b/kernels/aten/cpu/op__empty_dim_order.cpp new file mode 100644 index 0000000000..f11f853daa --- /dev/null +++ b/kernels/aten/cpu/op__empty_dim_order.cpp @@ -0,0 +1,110 @@ +/* + * 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 +#include + +#include +#include + +namespace torch { +namespace executor { +namespace native { + +using exec_aten::IntArrayRef; +using exec_aten::Tensor; +using OptionalIntArrayRef = exec_aten::OptionalArrayRef; +using DimOrderArrayRef = exec_aten::ArrayRef; +// Out Aten tensor shall have same memory format stride as dim_order +const size_t kMaxNumOfDimensions = 16; + +namespace { + +inline bool _check__empty_out_dim_order( + OptionalIntArrayRef dim_order, + Tensor& out) { + exec_aten::ArrayRef dim_order_ref; + std::vector dim_order_vec; + + if (dim_order.has_value()) { + // out tensor's dim order shall equal to input dim order + dim_order_ref = exec_aten::ArrayRef( + dim_order.value().data(), dim_order.value().size()); + } else { // dim_order is not set, out tensor should be contiguous dim order + for (int i = 0; i < out.dim(); i++) { + dim_order_vec.push_back(i); + } + dim_order_ref = exec_aten::ArrayRef(dim_order_vec); + } + + // dim order size shall equal to input dim + ET_LOG_AND_RETURN_IF_FALSE(dim_order_ref.size() == out.dim()); + + ET_LOG_AND_RETURN_IF_FALSE( + is_channels_last_dim_order(dim_order_ref.data(), dim_order_ref.size()) || + is_contiguous_dim_order(dim_order_ref.data(), dim_order_ref.size())); + + ET_LOG_AND_RETURN_IF_FALSE(kMaxNumOfDimensions >= out.dim()); + exec_aten::StridesType target_strides[kMaxNumOfDimensions]; + dim_order_to_stride_nocheck( + out.sizes().data(), + dim_order_ref.data(), + dim_order_ref.size(), + target_strides); + + for (size_t i = 0; i < dim_order_ref.size(); i++) { + ET_LOG_AND_RETURN_IF_FALSE(target_strides[i] == out.strides()[i]); + } + + return true; +} + +} // namespace + +/* + * Empty out tensor with specified dim order + * + * _empty_dim_order.out(SymInt[] size, *, int[]? dim_order=None, Tensor(a!) out) + * -> Tensor(a!) + */ +Tensor& _empty_dim_order_out( + KernelRuntimeContext& context, + IntArrayRef size, + OptionalIntArrayRef dim_order, + Tensor& out) { + (void)context; + + // Check if dim_order is valid + ET_KERNEL_CHECK( + context, + _check__empty_out_dim_order(dim_order, out), + InvalidArgument, + out); + + // Resize for dynamic shape + ET_KERNEL_CHECK_MSG( + context, + resize_tensor(out, size) == Error::Ok, + InvalidArgument, + out, + "Failed to resize output tensor."); + + return out; +} + +Tensor& _empty_dim_order_out( + IntArrayRef size, + OptionalIntArrayRef dim_order, + Tensor& out) { + executorch::runtime::KernelRuntimeContext ctx{}; + return _empty_dim_order_out(ctx, size, dim_order, out); +} + +} // namespace native +} // namespace executor +} // namespace torch diff --git a/kernels/aten/cpu/targets.bzl b/kernels/aten/cpu/targets.bzl index bdd93bda9e..bb7083c1f0 100644 --- a/kernels/aten/cpu/targets.bzl +++ b/kernels/aten/cpu/targets.bzl @@ -9,6 +9,9 @@ load("@fbsource//xplat/executorch/kernels/portable:op_registration_util.bzl", "d # ops, and must be split. They can, however, share common code via a library dep # if necessary. _EDGE_DIALECT_OPS = ( + op_target( + name = "op__empty_dim_order", + ), op_target( name = "op__to_dim_order_copy", deps = [ diff --git a/kernels/aten/edge_dialect_aten_op.yaml b/kernels/aten/edge_dialect_aten_op.yaml index 016f8dbfab..d9de3f6dde 100644 --- a/kernels/aten/edge_dialect_aten_op.yaml +++ b/kernels/aten/edge_dialect_aten_op.yaml @@ -2,6 +2,11 @@ # # This yaml file contains operators that are defined by ExecuTorch and used in ATen mode. +- func: dim_order_ops::_empty_dim_order.out(int[] size, *, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: torch::executor::_empty_dim_order_out + - func: dim_order_ops::_to_dim_order_copy.out(Tensor self, *, bool non_blocking=False, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null diff --git a/kernels/portable/cpu/op__empty_dim_order.cpp b/kernels/portable/cpu/op__empty_dim_order.cpp new file mode 100644 index 0000000000..a4d733662f --- /dev/null +++ b/kernels/portable/cpu/op__empty_dim_order.cpp @@ -0,0 +1,82 @@ +/* + * 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 +#include + +#include +#include + +namespace torch { +namespace executor { +namespace native { + +using exec_aten::Tensor; +using OptionalIntArrayRef = exec_aten::OptionalArrayRef; +using DimOrderArrayRef = exec_aten::ArrayRef; + +namespace { + +bool _check__empty_out_dim_order(OptionalIntArrayRef dim_order, Tensor& out) { + DimOrderArrayRef out_dim_order = out.dim_order(); + + if (dim_order.has_value()) { + // out tensor's dim order shall equal to input dim order + IntArrayRef dim_order_ref = dim_order.value(); + + ET_LOG_AND_RETURN_IF_FALSE( + is_channels_last_dim_order( + dim_order.value().data(), dim_order.value().size()) || + is_contiguous_dim_order( + dim_order.value().data(), dim_order.value().size())); + + // Out tensor shall have same dim order as dim_order + ET_LOG_AND_RETURN_IF_FALSE(out_dim_order.size() == dim_order_ref.size()); + for (size_t i = 0; i < dim_order_ref.size(); i++) { + ET_LOG_AND_RETURN_IF_FALSE(out_dim_order[i] == dim_order_ref[i]); + } + } else { // dim_order is not set, out tensor should be contiguous memory + // format + ET_LOG_AND_RETURN_IF_FALSE( + is_contiguous_dim_order(out_dim_order.data(), out_dim_order.size())); + } + return true; +} + +} // namespace + +/* + * Empty out tensor with specified dim order + * + * _empty_dim_order.out(SymInt[] size, *, int[]? dim_order=None, Tensor(a!) out) + * -> Tensor(a!) + */ +Tensor& _empty_dim_order_out( + KernelRuntimeContext& context, + IntArrayRef size, + OptionalIntArrayRef dim_order, + Tensor& out) { + (void)context; + + // Check if dim_order is valid + _check__empty_out_dim_order(dim_order, out); + + // Resize for dynamic shape + ET_KERNEL_CHECK_MSG( + context, + resize_tensor(out, size) == Error::Ok, + InvalidArgument, + out, + "Failed to resize output tensor."); + + return out; +} + +} // namespace native +} // namespace executor +} // namespace torch diff --git a/kernels/portable/functions.yaml b/kernels/portable/functions.yaml index a5d60eb59e..266b5e446f 100644 --- a/kernels/portable/functions.yaml +++ b/kernels/portable/functions.yaml @@ -937,6 +937,11 @@ - arg_meta: null kernel_name: torch::executor::zeros_out +- func: dim_order_ops::_empty_dim_order.out(int[] size, *, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: torch::executor::_empty_dim_order_out + - func: dim_order_ops::_to_dim_order_copy.out(Tensor self, *, bool non_blocking=False, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null diff --git a/kernels/test/op__empty_dim_order_test.cpp b/kernels/test/op__empty_dim_order_test.cpp new file mode 100644 index 0000000000..2857bac045 --- /dev/null +++ b/kernels/test/op__empty_dim_order_test.cpp @@ -0,0 +1,161 @@ +/* + * 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 // Declares the operator +#include +#include +#include +#include +#include + +#include + +using namespace ::testing; +using exec_aten::DimOrderType; +using exec_aten::IntArrayRef; +using exec_aten::optional; +using exec_aten::OptionalArrayRef; +using exec_aten::ScalarType; +using exec_aten::Tensor; +using torch::executor::testing::TensorFactory; + +class OpEmptyDimOrderOutTest : public OperatorTest { + protected: + Tensor& op_empty_dim_order_out( + IntArrayRef size, + OptionalArrayRef dim_order, + Tensor& out) { + return torch::executor::dim_order_ops::_empty_dim_order_outf( + context_, size, dim_order, out); + } + + template + void test_op_empty_dim_order_out(std::vector&& size_int32_t) { + TensorFactory tf; + std::vector sizes(size_int32_t.begin(), size_int32_t.end()); + auto aref = exec_aten::ArrayRef(sizes.data(), sizes.size()); + OptionalArrayRef dim_order; + Tensor out = tf.ones(size_int32_t); + + op_empty_dim_order_out(aref, dim_order, out); + } + + void too_short_dim_order_die() { + TensorFactory tf; + + int64_t sizes[3] = {3, 2, 4}; + auto sizes_aref = exec_aten::ArrayRef(sizes); + + int64_t raw_dim_order[2] = {0, 1}; + auto dim_order = OptionalArrayRef(raw_dim_order); + Tensor out = + tf.ones({3, 2, 4}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); + ET_EXPECT_KERNEL_FAILURE( + context_, op_empty_dim_order_out(sizes_aref, dim_order, out)); + } + + void illegal_dim_order_die() { + TensorFactory tf; + + int64_t sizes[2] = {3, 2}; + auto sizes_aref = exec_aten::ArrayRef(sizes); + + int64_t raw_dim_order[2] = {1, 2}; + auto dim_order = OptionalArrayRef(raw_dim_order); + Tensor out = + tf.ones({3, 2}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); + ET_EXPECT_KERNEL_FAILURE( + context_, op_empty_dim_order_out(sizes_aref, dim_order, out)); + } + + void wrong_dim_order_die() { + TensorFactory tf; + + int64_t sizes[4] = {3, 2, 4, 5}; + auto sizes_aref = exec_aten::ArrayRef(sizes); + + // should be {0, 2, 3, 1} + int64_t raw_dim_order[4] = {0, 1, 2, 3}; + auto dim_order = OptionalArrayRef(raw_dim_order); + Tensor out = tf.full_channels_last( + {3, 2, 4, 5}, 1, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); + ET_EXPECT_KERNEL_FAILURE( + context_, op_empty_dim_order_out(sizes_aref, dim_order, out)); + } +}; + +#define GENERATE_TEST(_, DTYPE) \ + TEST_F(OpEmptyDimOrderOutTest, DTYPE##Tensors) { \ + test_op_empty_dim_order_out({2, 3, 4}); \ + test_op_empty_dim_order_out({2, 0, 4}); \ + test_op_empty_dim_order_out({}); \ + } + +ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_TEST) + +TEST_F(OpEmptyDimOrderOutTest, DynamicShapeUpperBoundSameAsExpected) { + TensorFactory tf; + + int64_t sizes[2] = {3, 2}; + auto sizes_aref = exec_aten::ArrayRef(sizes); + OptionalArrayRef dim_order; + Tensor out = + tf.ones({3, 2}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); + op_empty_dim_order_out(sizes_aref, dim_order, out); +} + +TEST_F(OpEmptyDimOrderOutTest, ContiguousDimOrderSuccees) { + TensorFactory tf; + + int64_t sizes[2] = {3, 2}; + auto sizes_aref = exec_aten::ArrayRef(sizes); + + int64_t raw_dim_order[2] = {0, 1}; + auto dim_order = OptionalArrayRef(raw_dim_order); + Tensor out = + tf.ones({3, 2}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); + op_empty_dim_order_out(sizes_aref, dim_order, out); +} + +TEST_F(OpEmptyDimOrderOutTest, ChannelsLastsDimOrderSuccees) { + TensorFactory tf; + + int64_t sizes[4] = {3, 2, 4, 5}; + auto sizes_aref = exec_aten::ArrayRef(sizes); + + int64_t raw_dim_order[4] = {0, 2, 3, 1}; + auto dim_order = OptionalArrayRef(raw_dim_order); + Tensor out = tf.full_channels_last( + {3, 2, 4, 5}, 1, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); + op_empty_dim_order_out(sizes_aref, dim_order, out); +} + +TEST_F(OpEmptyDimOrderOutTest, DynamicShapeUpperBoundLargerThanExpected) { + TensorFactory tf; + + int64_t sizes[2] = {3, 2}; + auto sizes_aref = exec_aten::ArrayRef(sizes); + OptionalArrayRef dim_order; + Tensor out = + tf.ones({10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND); + op_empty_dim_order_out(sizes_aref, dim_order, out); +} + +TEST_F(OpEmptyDimOrderOutTest, DynamicShapeUnbound) { + if (!torch::executor::testing::SupportedFeatures::get()->output_resize) { + GTEST_SKIP() << "Dynamic shape unbound not supported"; + } + TensorFactory tf; + + int64_t sizes[2] = {3, 2}; + auto sizes_aref = exec_aten::ArrayRef(sizes); + OptionalArrayRef dim_order; + Tensor out = + tf.ones({1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND); + op_empty_dim_order_out(sizes_aref, dim_order, out); +} diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 77b18a4814..2e7c34f147 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -174,6 +174,7 @@ def define_common_targets(): codegen_function_header_wrapper("executorch/kernels/test/custom_kernel_example", "custom_kernel_example") _common_op_test("op__to_dim_order_copy_test", ["aten", "portable"]) + _common_op_test("op__empty_dim_order_test", ["aten", "portable"]) _common_op_test("op_abs_test", ["aten", "portable"]) _common_op_test("op_acos_test", ["aten", "portable"]) _common_op_test("op_acosh_test", ["aten", "portable"]) diff --git a/shim/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim/xplat/executorch/kernels/portable/op_registration_util.bzl index 53698e7f21..b88e83b058 100644 --- a/shim/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -1252,6 +1252,12 @@ ATEN_OPS = ( op_target( name = "op_zeros", ), + op_target( + name = "op__empty_dim_order", + deps = [ + ":scalar_utils", + ], + ), op_target( name = "op__to_dim_order_copy", deps = [