TensorLayout updates (pytorch#7870)

Summary:

Addressing comments from D67048723

Differential Revision: D68535453

Author: lucylq

runtime/core/targets.bzl

@@ -43,7 +43,6 @@ def define_common_targets():
             "freeable_buffer.h",
             "result.h",
             "span.h",
-            "tensor_layout.h",
         ],
         visibility = [
             "//executorch/...",
@@ -132,3 +131,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/..."],
+    )

runtime/core/tensor_layout.cpp

@@ -0,0 +1,44 @@
+/*
+ * 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/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/span.h>
+#include <executorch/runtime/core/tensor_layout.h>
+
+namespace executorch {
+namespace runtime {
+
+namespace {
+Result<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++) {
+    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> TensorLayout::create(
+    Span<const int32_t> sizes,
+    Span<const uint8_t> dim_order,
+    executorch::aten::ScalarType scalar_type) {
+  auto nbytes = calculate_nbytes(sizes, scalar_type);
+  if (!nbytes.ok()) {
+    return nbytes.error();
+  }
+  return TensorLayout(sizes, dim_order, scalar_type, nbytes.get());
+}
+} // namespace runtime
+} // namespace executorch

runtime/core/tensor_layout.h

@@ -10,55 +10,55 @@
 
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/result.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.
+ * Describes the layout of a tensor.
  */
-class TensorLayout {
+class ET_EXPERIMENTAL TensorLayout final {
  public:
   TensorLayout(
-      executorch::aten::ScalarType scalar_type,
       Span<const int32_t> sizes,
-      Span<const uint8_t> dim_order)
+      Span<const uint8_t> dim_order,
+      executorch::aten::ScalarType scalar_type,
+      size_t nbytes)
       : sizes_(sizes),
         dim_order_(dim_order),
         scalar_type_(scalar_type),
-        nbytes_(calculate_nbytes(sizes_, scalar_type_)) {}
+        nbytes_(nbytes) {}
 
-  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<TensorLayout> create(
+      Span<const int32_t> sizes,
+      Span<const uint8_t> 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<const int32_t> 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<const uint8_t> dim_order() const {
     return dim_order_;
   }

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",
         ],
     )
 

runtime/core/test/tensor_layout_test.cpp

@@ -7,12 +7,14 @@
  */
 
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/result.h>
 #include <executorch/runtime/core/tensor_layout.h>
 
 #include <gtest/gtest.h>
 
 using namespace ::testing;
 using executorch::aten::ScalarType;
+using executorch::runtime::Result;
 using executorch::runtime::Span;
 using executorch::runtime::TensorLayout;
 
@@ -23,9 +25,12 @@ TEST(TestTensorLayout, Ctor) {
   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);
+  Result<TensorLayout> 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());