[flat_tensor] Persist FreeableBuffers of external constants in method

## Problem
Currently, the FlatTensorDataMap persists tensors, and returns a FreeableBuffer with an empty free function.

The NamedDataMap should not persist data, as most cases (eg. delegate) will want it to be freed.

Ownership should be on the caller; `get_data` returns a FreeableBuffer that 'owns' the data. The FreeableBuffer in turn is owned by the caller.

NOTE: this doesn't support the case where we want to share plain tensors between methods/pte files at runtime. A custom NDM could support that use-case.

## This diff:
1. Introduces a 'NamedData' struct to method.h. This holds a key and a FreeeableBuffer.
2. Iterate over all the flatbuffer tensors to count the constants tagged with EXTERNAL.
NOTE: this will increase load time for all users. Potentially allocate chunks of 16 and use a linked list to store external constants, or store this number in PTE file.
3. Allocate space for num_external_constants using the method allocator.
4. Iterate over all flatbuffer tensors and use the named_data_map to resolve EXTERNAL tensors into the array of NamedData.
5. Pass the resolved external constants to tensor_parser, along with NDM (used for mutable external tensors).
6. Resolved external tensors are stored inside method. They are freed when the method is destructed.


Some notes:
https://docs.google.com/document/d/1_PBi4JgODuClUPD4PCUWrKNjyUH54zOUHGUJ3QHDNes/edit?tab=t.0#heading=h.blsvwraxss7g

Differential Revision: [D69477027](https://our.internmc.facebook.com/intern/diff/D69477027/)

ghstack-source-id: 266173959
Pull Request resolved: #8437

Author: lucylq

extension/flat_tensor/targets.bzl

@@ -16,6 +16,9 @@ def define_common_targets():
             "//executorch/runtime/core/exec_aten:lib",
             "//executorch/runtime/core/exec_aten/util:tensor_util",
         ],
+        exported_deps = [
+            "//executorch/extension/flat_tensor/serialize:flat_tensor_header",
+        ],
         visibility = [
             "//executorch/...",
         ],

runtime/executor/method.cpp

@@ -289,6 +289,120 @@ Result<bool> parse_cond_value(const EValue& cond_value) {
 
 } // namespace
 
+Result<size_t> Method::get_num_external_constants() {
+  auto flatbuffer_values = serialization_plan_->values();
+  size_t n_value = flatbuffer_values->size();
+
+  size_t num_external_constants = 0;
+  for (size_t i = 0; i < n_value; ++i) {
+    auto serialization_value = flatbuffer_values->Get(i);
+    // Ensure that the `val_as_X()` calls will return non-null pointers.
+    ET_CHECK_OR_RETURN_ERROR(
+        serialization_value != nullptr &&
+            (serialization_value->val_type() ==
+                 executorch_flatbuffer::KernelTypes::Null ||
+             serialization_value->val() != nullptr),
+        InvalidProgram,
+        "Null value at index %" ET_PRIsize_t,
+        i);
+    // Ignore non-tensor types.
+    if (serialization_value->val_type() !=
+        executorch_flatbuffer::KernelTypes::Tensor) {
+      continue;
+    }
+    const auto s_tensor = static_cast<const executorch_flatbuffer::Tensor*>(
+        serialization_value->val());
+
+    // An external constant is tagged with EXTERNAL and has no
+    // allocation_info.
+    if (s_tensor->extra_tensor_info() != nullptr &&
+        s_tensor->extra_tensor_info()->location() ==
+            executorch_flatbuffer::TensorDataLocation::EXTERNAL &&
+        s_tensor->allocation_info() == nullptr) {
+      num_external_constants++;
+    }
+  }
+  return num_external_constants;
+}
+
+// Check if key exists in external_constants_.
+// Helper function for parse_external_constants.
+bool key_exists(const char* key, NamedData* external_constants, int num_keys) {
+  for (int i = 0; i < num_keys; i++) {
+    if (strcmp(key, external_constants[i].key) == 0) {
+      return true;
+    }
+  }
+  return false;
+}
+
+Error Method::parse_external_constants(const NamedDataMap* named_data_map) {
+  auto flatbuffer_values = serialization_plan_->values();
+  size_t n_value = flatbuffer_values->size();
+
+  // Stores the number of unique external tensors that have been resolved.
+  int index = 0;
+  for (size_t i = 0; i < n_value; ++i) {
+    auto serialization_value = flatbuffer_values->Get(i);
+    // Ignore non-tensor types.
+    if (serialization_value->val_type() !=
+        executorch_flatbuffer::KernelTypes::Tensor) {
+      continue;
+    }
+    const auto s_tensor = static_cast<const executorch_flatbuffer::Tensor*>(
+        serialization_value->val());
+    // Note: tensors with allocation_info are mutable, and resolved in
+    // parse_values.
+    if (s_tensor->extra_tensor_info() != nullptr &&
+        s_tensor->extra_tensor_info()->location() ==
+            executorch_flatbuffer::TensorDataLocation::EXTERNAL &&
+        s_tensor->allocation_info() == nullptr) {
+      ET_CHECK_OR_RETURN_ERROR(
+          s_tensor->extra_tensor_info()->fully_qualified_name() != nullptr,
+          InvalidExternalData,
+          "Fully qualified name of external tensor is null");
+
+      const char* key =
+          s_tensor->extra_tensor_info()->fully_qualified_name()->c_str();
+
+      // Check if this tensor has already been resolved.
+      if (!key_exists(key, external_constants_, index)) {
+        Result<const TensorLayout> tensor_layout =
+            named_data_map->get_metadata(key);
+        if (!tensor_layout.ok()) {
+          return tensor_layout.error();
+        }
+        // Check external tensor compatibility.
+        Error err = deserialization::validateExternalTensor(
+            s_tensor, tensor_layout.get());
+        if (err != Error::Ok) {
+          return err;
+        }
+        // Save the key.
+        external_constants_[index].key = key;
+
+        // Save the buffer.
+        Result<FreeableBuffer> buffer = named_data_map->get_data(key);
+        ET_CHECK_OR_RETURN_ERROR(
+            buffer.ok(),
+            InvalidExternalData,
+            "Buffer retrieved from get_data is not valid");
+        external_constants_[index].buffer =
+            memory_manager_->method_allocator()
+                ->allocateInstance<FreeableBuffer>();
+        if (external_constants_[index].buffer == nullptr) {
+          return Error::MemoryAllocationFailed;
+        }
+        new (external_constants_[index].buffer)
+            FreeableBuffer(std::move(buffer.get()));
+
+        index++;
+      }
+    }
+  }
+  return Error::Ok;
+}
+
 Error Method::parse_values(const NamedDataMap* named_data_map) {
   auto flatbuffer_values = serialization_plan_->values();
   ET_CHECK_OR_RETURN_ERROR(
@@ -299,23 +413,33 @@ Error Method::parse_values(const NamedDataMap* named_data_map) {
     return Error::MemoryAllocationFailed;
   }
 
+  // Check if there are external constants.
+  Result<size_t> num_external_constants = get_num_external_constants();
+  if (!num_external_constants.ok()) {
+    return num_external_constants.error();
+  }
+  num_external_constants_ = *num_external_constants;
+  if (num_external_constants_ > 0) {
+    // Allocate space for external tensors.
+    external_constants_ =
+        memory_manager_->method_allocator()->allocateList<NamedData>(
+            num_external_constants_);
+    if (external_constants_ == nullptr) {
+      return Error::MemoryAllocationFailed;
+    }
+    auto err = parse_external_constants(named_data_map);
+    if (err != Error::Ok) {
+      return err;
+    }
+  }
+
   // n_value_ counts the number of successfully-initialized values for ~Method()
   // to clean up, and is incremented at the bottom of the loop. This makes it
   // safe for errors to return without updating any state.
   n_value_ = 0;
 
   for (size_t i = 0; i < n_value; ++i) {
     auto serialization_value = flatbuffer_values->Get(i);
-    // Ensure that the `val_as_X()` calls will return non-null pointers.
-    ET_CHECK_OR_RETURN_ERROR(
-        serialization_value != nullptr &&
-            (serialization_value->val_type() ==
-                 executorch_flatbuffer::KernelTypes::Null ||
-             serialization_value->val() != nullptr),
-        InvalidProgram,
-        "Null value at index %" ET_PRIsize_t,
-        i);
-
     const auto val = serialization_value->val();
 
     switch (serialization_value->val_type()) {
@@ -416,7 +540,8 @@ Error Method::parse_values(const NamedDataMap* named_data_map) {
             program_,
             memory_manager_,
             static_cast<const executorch_flatbuffer::Tensor*>(val),
-            named_data_map);
+            named_data_map,
+            Span<NamedData>(external_constants_, num_external_constants_));
         if (!t.ok()) {
           ET_LOG(
               Error,
@@ -1496,6 +1621,11 @@ Method::~Method() {
       delegates_[i].~BackendDelegate();
     }
   }
+  // Free resources associated with external constants.
+  for (int i = 0; i < num_external_constants_; i++) {
+    external_constants_[i].buffer->Free();
+    external_constants_[i].buffer = nullptr;
+  }
   // All other fields are trivially destructible.
 }
 } // namespace runtime

runtime/executor/method.h

@@ -43,6 +43,13 @@ using OpFunction = void (*)(KernelRuntimeContext&, EValue**);
 /// argument list for a single instruction
 using InstructionArgs = Span<EValue*>;
 
+/// Data structure to hold key and data buffer for external data used
+/// in a method.
+struct NamedData {
+  const char* key;
+  FreeableBuffer* buffer;
+};
+
 /**
  * An executable method of an executorch program. Maps to a python method like
  * `forward()` on the original nn.Module.
@@ -66,7 +73,9 @@ class Method final {
         delegates_(rhs.delegates_),
         n_chains_(rhs.n_chains_),
         chains_(rhs.chains_),
-        init_state_(rhs.init_state_) {
+        init_state_(rhs.init_state_),
+        external_constants_(rhs.external_constants_),
+        num_external_constants_(rhs.num_external_constants_) {
     // Required: clear out fields that the dtor looks at, so that we don't free
     // anything twice.
     rhs.n_value_ = 0;
@@ -84,6 +93,8 @@ class Method final {
     rhs.event_tracer_ = nullptr;
     rhs.n_chains_ = 0;
     rhs.chains_ = nullptr;
+    rhs.external_constants_ = nullptr;
+    rhs.num_external_constants_ = 0;
   }
 
   /**
@@ -288,7 +299,9 @@ class Method final {
         delegates_(nullptr),
         n_chains_(0),
         chains_(nullptr),
-        init_state_(InitializationState::Uninitialized) {}
+        init_state_(InitializationState::Uninitialized),
+        external_constants_(nullptr),
+        num_external_constants_(0) {}
 
   /// Static factory used by Program.
   ET_NODISCARD static Result<Method> load(
@@ -338,6 +351,23 @@ class Method final {
 
   InitializationState init_state_;
 
+  NamedData* external_constants_;
+  size_t num_external_constants_ = 0;
+
+  /**
+   * Counts the number of external constants in flatbuffer.
+   */
+  ET_NODISCARD Result<size_t> get_num_external_constants();
+
+  /**
+   * Parses the flatbuffer for constant tensors tagged as EXTERNAL.
+   * Retrieves the external constants using the named_data_map and places them
+   * into `external_constants_`.
+   * FreeableBuffers returned by the named_data_map are owned by the
+   * method and are freed on method destruction.
+   */
+  ET_NODISCARD Error
+  parse_external_constants(const NamedDataMap* named_data_map);
   /**
    * Parses the elements of the values_ array. On error, n_value_ will be set to
    * the number of successfully-initialized entries so that ~Method doesn't try

runtime/executor/tensor_parser.h

@@ -14,6 +14,7 @@
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/executor/memory_manager.h>
+#include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/schema/program_generated.h>
 
@@ -25,14 +26,21 @@ ET_NODISCARD Result<executorch::aten::Tensor> parseTensor(
     const Program* program,
     MemoryManager* memory_manager,
     const executorch_flatbuffer::Tensor* s_tensor,
-    const NamedDataMap* named_data_map = nullptr);
+    const NamedDataMap* named_data_map = nullptr,
+    Span<NamedData> external_constants = {});
 
 ET_NODISCARD Result<BoxedEvalueList<executorch::aten::Tensor>> parseTensorList(
     const flatbuffers::Vector<int32_t>* tensor_indices,
     EValue* values,
     size_t values_len,
     MemoryManager* memory_manager);
 
+// Checks that the sizes, dim_order and scalar_type match between tensors
+// stored in the PTE and externally.
+ET_NODISCARD Error validateExternalTensor(
+    const executorch_flatbuffer::Tensor* s_tensor,
+    const TensorLayout& tensor_layout);
+
 // Deserializes a List of optional type. The code here is the same between all
 // list of optionals: list of optional Tensor, list of optional float etc, so we
 // just use a template to avoid boilerplate.
@@ -105,7 +113,9 @@ parseListOptionalType(
  * @param[in] nbytes The amount of memory to get from the allocator.
  * @param[in] allocator The source of memory for non-constant tensors.
  * @param[in] named_data_map An optional map of {name, blob} used to resolve
- *     data that is external to the PTE, if any.
+ *     data that is mutable and external to the PTE, if any.
+ * @param[in] external_constants An optional span of {name, buffer} used to
+ *     resolve data that is constant and external to the PTE, if any.
  *
  * @returns On success, the data pointer to use for the tensor. On failure, a
  *     non-Ok Error.
@@ -115,7 +125,8 @@ ET_NODISCARD Result<void*> getTensorDataPtr(
     const Program* program,
     size_t nbytes,
     HierarchicalAllocator* allocator,
-    const NamedDataMap* named_data_map = nullptr);
+    const NamedDataMap* named_data_map = nullptr,
+    Span<NamedData> external_constants = {});
 
 } // namespace deserialization
 } // namespace runtime

runtime/executor/tensor_parser_aten.cpp

@@ -12,6 +12,7 @@
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/core/named_data_map.h>
 #include <executorch/runtime/executor/memory_manager.h>
+#include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/platform/profiler.h>
 #include <executorch/schema/program_generated.h>
@@ -33,7 +34,8 @@ Result<at::Tensor> parseTensor(
     const Program* program,
     MemoryManager* memory_manager,
     const executorch_flatbuffer::Tensor* s_tensor,
-    const NamedDataMap* named_data_map) {
+    const NamedDataMap* named_data_map,
+    Span<NamedData> external_constants) {
   EXECUTORCH_SCOPE_PROF("TensorParser::parseTensor");
 
   ET_CHECK_OR_RETURN_ERROR(
@@ -108,7 +110,8 @@ Result<at::Tensor> parseTensor(
         program,
         tensor.nbytes(),
         memory_manager->planned_memory(),
-        named_data_map);
+        named_data_map,
+        external_constants);
     if (!data_ptr.ok()) {
       ET_LOG(
           Error,