From 4df19823b6699f30ac5bb77d3786e777aae45e24 Mon Sep 17 00:00:00 2001 From: Lucy Qiu Date: Mon, 27 Jan 2025 10:17:00 -0800 Subject: [PATCH] TensorLayout updates (#7870) Summary: Addressing comments from D67048723 Reviewed By: JacobSzwejbka Differential Revision: D68535453 --- runtime/core/targets.bzl | 12 +++- runtime/core/tensor_layout.cpp | 54 ++++++++++++++++++ runtime/core/tensor_layout.h | 72 ++++++++++++------------ runtime/core/test/targets.bzl | 3 +- runtime/core/test/tensor_layout_test.cpp | 52 ++++++++++++++--- 5 files changed, 148 insertions(+), 45 deletions(-) create mode 100644 runtime/core/tensor_layout.cpp diff --git a/runtime/core/targets.bzl b/runtime/core/targets.bzl index 253e2526107..61351b083bf 100644 --- a/runtime/core/targets.bzl +++ b/runtime/core/targets.bzl @@ -44,7 +44,6 @@ def define_common_targets(): "named_data_map.h", "result.h", "span.h", - "tensor_layout.h", ], visibility = [ "//executorch/...", @@ -133,3 +132,14 @@ def define_common_targets(): "//executorch/...", ], ) + + runtime.cxx_library( + name = "tensor_layout", + srcs = ["tensor_layout.cpp"], + exported_headers = ["tensor_layout.h"], + exported_deps = [ + ":core", + "//executorch/runtime/core/exec_aten:lib", + ], + visibility = ["//executorch/..."], + ) diff --git a/runtime/core/tensor_layout.cpp b/runtime/core/tensor_layout.cpp new file mode 100644 index 00000000000..4dad23dbd14 --- /dev/null +++ b/runtime/core/tensor_layout.cpp @@ -0,0 +1,54 @@ +/* + * 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 executorch { +namespace runtime { + +namespace { +Result calculate_nbytes( + const Span& sizes, + const exec_aten::ScalarType& scalar_type) { + ssize_t n = 1; + for (ssize_t i = 0; i < sizes.size(); i++) { + if (sizes[i] < 0) { + return Error::InvalidArgument; + } + n *= sizes[i]; + } + // Use the full namespace to disambiguate from c10::elementSize. + return n * executorch::runtime::elementSize(scalar_type); +} +} // namespace + +Result TensorLayout::create( + Span sizes, + Span dim_order, + executorch::aten::ScalarType scalar_type) { + auto nbytes = calculate_nbytes(sizes, scalar_type); + if (!nbytes.ok()) { + return nbytes.error(); + } + + if (dim_order.size() != sizes.size()) { + return Error::InvalidArgument; + } + + for (size_t i = 0; i < dim_order.size(); i++) { + if (dim_order[i] >= sizes.size()) { + return Error::InvalidArgument; + } + } + return TensorLayout(sizes, dim_order, scalar_type, nbytes.get()); +} +} // namespace runtime +} // namespace executorch diff --git a/runtime/core/tensor_layout.h b/runtime/core/tensor_layout.h index 84238561412..dcd0a2af839 100644 --- a/runtime/core/tensor_layout.h +++ b/runtime/core/tensor_layout.h @@ -10,55 +10,48 @@ #include #include +#include #include namespace executorch { namespace runtime { -namespace { -size_t calculate_nbytes( - const Span& sizes, - const exec_aten::ScalarType& scalar_type) { - ssize_t n = 1; - for (ssize_t i = 0; i < sizes.size(); i++) { - ET_CHECK(sizes[i] >= 0); - n *= sizes[i]; - } - // Use the full namespace to disambiguate from c10::elementSize. - return n * executorch::runtime::elementSize(scalar_type); -} -} // namespace - /** - * Metadata describing the layout of external tensors (tensors that are not - stored in the PTE file). - * - * The NamedDataMap used to create the TensorLayout must outlive the - TensorLayout. + * Describes the layout of a tensor. */ -class TensorLayout { +class ET_EXPERIMENTAL TensorLayout final { public: - TensorLayout( - executorch::aten::ScalarType scalar_type, - Span sizes, - Span dim_order) - : sizes_(sizes), - dim_order_(dim_order), - scalar_type_(scalar_type), - nbytes_(calculate_nbytes(sizes_, scalar_type_)) {} + TensorLayout() = delete; - TensorLayout(const TensorLayout&) = default; - TensorLayout(TensorLayout&&) = default; - TensorLayout& operator=(const TensorLayout&) = default; - TensorLayout& operator=(TensorLayout&& other) = default; - ~TensorLayout() = default; + /** + * Creates a TensorLayout from the given parameters. + * + * @param[in] sizes The sizes of the tensor. Note: the span passed here must + * outlive the TensorLayout and all copies of it. + * @param[in] dim_order The dim order of the tensor. Note: the span passed + * here must outlive the TensorLayout and all copies of it. + * @param[in] scalar_type The scalar type of the tensor. + * @return A Result containing the TensorLayout on success, or an error. + */ + static executorch::runtime::Result create( + Span sizes, + Span dim_order, + executorch::aten::ScalarType scalar_type); - /// Returns the sizes of the tensor. + /** + * Returns the sizes of the tensor. + * + * NOTE: The TensorLayout must outlive the spans returned here. + */ Span sizes() const { return sizes_; } - /// Returns the dim order of the tensor. + /** + * Returns the dim order of the tensor. + * + * NOTE: The TensorLayout must outlive the spans returned here. + */ Span dim_order() const { return dim_order_; } @@ -74,6 +67,15 @@ class TensorLayout { } private: + TensorLayout( + Span sizes, + Span dim_order, + executorch::aten::ScalarType scalar_type, + size_t nbytes) + : sizes_(sizes), + dim_order_(dim_order), + scalar_type_(scalar_type), + nbytes_(nbytes) {} /// The sizes of the tensor. Span sizes_; diff --git a/runtime/core/test/targets.bzl b/runtime/core/test/targets.bzl index 34dc6e0f3b2..7332aad8a3d 100644 --- a/runtime/core/test/targets.bzl +++ b/runtime/core/test/targets.bzl @@ -19,8 +19,7 @@ def define_common_targets(): name = "tensor_layout_test", srcs = ["tensor_layout_test.cpp"], deps = [ - "//executorch/runtime/core:core", - "//executorch/runtime/core/exec_aten:lib", + "//executorch/runtime/core:tensor_layout", ], ) diff --git a/runtime/core/test/tensor_layout_test.cpp b/runtime/core/test/tensor_layout_test.cpp index e9065f93cce..3e032788eab 100644 --- a/runtime/core/test/tensor_layout_test.cpp +++ b/runtime/core/test/tensor_layout_test.cpp @@ -6,26 +6,31 @@ * LICENSE file in the root directory of this source tree. */ +#include #include +#include #include #include using namespace ::testing; using executorch::aten::ScalarType; +using executorch::runtime::Error; +using executorch::runtime::Result; using executorch::runtime::Span; using executorch::runtime::TensorLayout; TEST(TestTensorLayout, Ctor) { - int32_t sizes[2] = {1, 2}; - uint8_t dim_order[2] = {0, 1}; + std::array sizes = {1, 2}; + std::array dim_order = {0, 1}; + Span sizes_span = {sizes.data(), sizes.size()}; + Span dim_order_span = {dim_order.data(), dim_order.size()}; - Span sizes_span = {sizes, sizes + 2}; - Span dim_order_span = {dim_order, dim_order + 2}; - - TensorLayout layout = - TensorLayout(ScalarType::Float, sizes_span, dim_order_span); + Result layout_res = + TensorLayout::create(sizes_span, dim_order_span, ScalarType::Float); + EXPECT_TRUE(layout_res.ok()); + TensorLayout layout = layout_res.get(); EXPECT_EQ(layout.scalar_type(), ScalarType::Float); EXPECT_EQ(layout.sizes().size(), sizes_span.size()); @@ -38,3 +43,36 @@ TEST(TestTensorLayout, Ctor) { EXPECT_EQ(layout.nbytes(), 8); } + +TEST(TestTensorLayout, Ctor_InvalidDimOrder) { + std::array sizes = {2}; + std::array dim_order = {1}; + Span sizes_span = {sizes.data(), sizes.size()}; + Span dim_order_span = {dim_order.data(), dim_order.size()}; + + Result layout_res = + TensorLayout::create(sizes_span, dim_order_span, ScalarType::Float); + EXPECT_EQ(layout_res.error(), Error::InvalidArgument); +} + +TEST(TestTensorLayout, Ctor_InvalidSizes) { + std::array sizes = {-1}; + std::array dim_order = {0}; + Span sizes_span = {sizes.data(), sizes.size()}; + Span dim_order_span = {dim_order.data(), dim_order.size()}; + + Result layout_res = + TensorLayout::create(sizes_span, dim_order_span, ScalarType::Float); + EXPECT_EQ(layout_res.error(), Error::InvalidArgument); +} + +TEST(TestTensorLayout, Ctor_SizesDimOrderMismatch) { + std::array sizes = {2}; + std::array dim_order = {0, 1}; + Span sizes_span = {sizes.data(), sizes.size()}; + Span dim_order_span = {dim_order.data(), dim_order.size()}; + + Result layout_res = + TensorLayout::create(sizes_span, dim_order_span, ScalarType::Float); + EXPECT_EQ(layout_res.error(), Error::InvalidArgument); +}