[executorch][core] Add TensorLayout to core

Differential Revision: D67048723

Pull Request resolved: #7761

runtime/core/targets.bzl

@@ -34,6 +34,7 @@ def define_common_targets():
             "freeable_buffer.h",
             "result.h",
             "span.h",
+            "tensor_layout.h",
         ],
         visibility = [
             "//executorch/...",

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 <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/span.h>
+
+namespace executorch {
+namespace runtime {
+
+namespace {
+size_t calculate_nbytes(
+    const Span<const int32_t>& 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<const int32_t> sizes,
+      Span<const uint8_t> 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<const int32_t> sizes() const {
+    return sizes_;
+  }
+
+  /// Returns the dim order of the tensor.
+  Span<const uint8_t> 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<const int32_t> sizes_;
+
+  /// The dim order of the tensor.
+  Span<const uint8_t> 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

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 = [

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 <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/tensor_layout.h>
+
+#include <gtest/gtest.h>
+
+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<const int32_t> sizes_span = {sizes, sizes + 2};
+  Span<const uint8_t> 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);
+}