TensorLayout updates (#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,45 @@
+/*
+ * 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 {
+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
+
+/**
+ * NOTE: the spans (sizes, dim_order) passed here must outlive the
+ * TensorLayout and all copies of it.
+ */
+TensorLayout::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_)) {}
+
+} // namespace runtime
+} // namespace executorch

runtime/core/tensor_layout.h

@@ -15,38 +15,19 @@
 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.
+ * NOTE: the underlying data structure/s backing the TensorLayout must outlive
+ * the TensorLayout and all copies of it.
  */
-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)
-      : sizes_(sizes),
-        dim_order_(dim_order),
-        scalar_type_(scalar_type),
-        nbytes_(calculate_nbytes(sizes_, scalar_type_)) {}
-
+      Span<const uint8_t> dim_order);
   TensorLayout(const TensorLayout&) = default;
   TensorLayout(TensorLayout&&) = default;
   TensorLayout& operator=(const TensorLayout&) = default;
@@ -74,10 +55,18 @@ class TensorLayout {
   }
 
  private:
-  /// The sizes of the tensor.
+  /**
+   * The sizes of the tensor.
+   *
+   * NOTE: The TensorLayout must outlive the spans returned here.
+   */
   Span<const int32_t> sizes_;
 
-  /// The dim order of the tensor.
+  /**
+   * The dim order of the tensor.
+   *
+   * NOTE: The TensorLayout must outlive the spans returned here.
+   */
   Span<const uint8_t> dim_order_;
 
   /// The scalar type of the tensor.

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