From 2a1922a898f2d6224e8683ba1254ee89e31deb91 Mon Sep 17 00:00:00 2001 From: lucylq Date: Fri, 17 Jan 2025 16:44:55 -0800 Subject: [PATCH] [executorch][core] Add TensorLayout to core Introduce TensorLayout class, used to describe external tensors. Currently contains: - scalar_type - sizes - dim_order Differential Revision: [D67048723](https://our.internmc.facebook.com/intern/diff/D67048723/) ghstack-source-id: 262004745 Pull Request resolved: https://github.com/pytorch/executorch/pull/7761 --- runtime/core/targets.bzl | 1 + runtime/core/tensor_layout.h | 91 ++++++++++++++++++++++++ runtime/core/test/targets.bzl | 9 +++ runtime/core/test/tensor_layout_test.cpp | 40 +++++++++++ 4 files changed, 141 insertions(+) create mode 100644 runtime/core/tensor_layout.h create mode 100644 runtime/core/test/tensor_layout_test.cpp diff --git a/runtime/core/targets.bzl b/runtime/core/targets.bzl index 352c962eb2..d64d457953 100644 --- a/runtime/core/targets.bzl +++ b/runtime/core/targets.bzl @@ -43,6 +43,7 @@ def define_common_targets(): "freeable_buffer.h", "result.h", "span.h", + "tensor_layout.h", ], visibility = [ "//executorch/...", diff --git a/runtime/core/tensor_layout.h b/runtime/core/tensor_layout.h new file mode 100644 index 0000000000..8423856141 --- /dev/null +++ b/runtime/core/tensor_layout.h @@ -0,0 +1,91 @@ +/* + * 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. + */ + +#pragma once + +#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. + */ +class TensorLayout { + 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(const TensorLayout&) = default; + TensorLayout(TensorLayout&&) = default; + TensorLayout& operator=(const TensorLayout&) = default; + TensorLayout& operator=(TensorLayout&& other) = default; + ~TensorLayout() = default; + + /// Returns the sizes of the tensor. + Span sizes() const { + return sizes_; + } + + /// Returns the dim order of the tensor. + Span dim_order() const { + return dim_order_; + } + + /// Returns the scalar type of the tensor. + executorch::aten::ScalarType scalar_type() const { + return scalar_type_; + } + + /// Returns the size of the tensor in bytes. + size_t nbytes() const { + return nbytes_; + } + + private: + /// The sizes of the tensor. + Span sizes_; + + /// The dim order of the tensor. + Span dim_order_; + + /// The scalar type of the tensor. + executorch::aten::ScalarType scalar_type_; + + /// The size in bytes of the tensor. + size_t nbytes_; +}; + +} // namespace runtime +} // namespace executorch diff --git a/runtime/core/test/targets.bzl b/runtime/core/test/targets.bzl index d221d49e6e..34dc6e0f3b 100644 --- a/runtime/core/test/targets.bzl +++ b/runtime/core/test/targets.bzl @@ -15,6 +15,15 @@ def define_common_targets(): ], ) + runtime.cxx_test( + name = "tensor_layout_test", + srcs = ["tensor_layout_test.cpp"], + deps = [ + "//executorch/runtime/core:core", + "//executorch/runtime/core/exec_aten:lib", + ], + ) + runtime.cxx_test( name = "error_handling_test", srcs = [ diff --git a/runtime/core/test/tensor_layout_test.cpp b/runtime/core/test/tensor_layout_test.cpp new file mode 100644 index 0000000000..e9065f93cc --- /dev/null +++ b/runtime/core/test/tensor_layout_test.cpp @@ -0,0 +1,40 @@ +/* + * 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 + +using namespace ::testing; +using executorch::aten::ScalarType; +using executorch::runtime::Span; +using executorch::runtime::TensorLayout; + +TEST(TestTensorLayout, Ctor) { + int32_t sizes[2] = {1, 2}; + uint8_t dim_order[2] = {0, 1}; + + 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); + + EXPECT_EQ(layout.scalar_type(), ScalarType::Float); + + EXPECT_EQ(layout.sizes().size(), sizes_span.size()); + EXPECT_EQ(layout.sizes()[0], sizes_span[0]); + EXPECT_EQ(layout.sizes()[1], sizes_span[1]); + + EXPECT_EQ(layout.dim_order().size(), dim_order_span.size()); + EXPECT_EQ(layout.dim_order()[0], dim_order_span[0]); + EXPECT_EQ(layout.dim_order()[1], dim_order_span[1]); + + EXPECT_EQ(layout.nbytes(), 8); +}