[executorch][flat tensor] Store number of external tensors in flatbuffer

Store the number of external tensors in each ExecutionPlan.

This saves time during loading, as we can directly allocate memory for external tensors without doing a pass to count how many external tensors require allocation.

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

ghstack-source-id: 266371204
Pull Request resolved: #8483

Author: lucylq

exir/emit/_emit_program.py

@@ -163,6 +163,7 @@ def emit_program(
             operator_cache={},
             delegate_cache={},
             emit_stacktrace=emit_stacktrace,
+            num_external_constants=0,
         )
 
         gm = _remove_non_user_outputs(exported_program)

exir/emit/_emitter.py

@@ -143,6 +143,7 @@ class _EmitterState:
     """
 
     values: List[EValue]
+    num_external_constants: int
     operators: List[Operator]
     delegates: List[BackendDelegate]
     operator_cache: Dict[Tuple[str, str], int]
@@ -407,6 +408,7 @@ def _save_new_const_tensor(
             self.program_state.external_constant_map[constant_tag][
                 spec.extra_tensor_info.fully_qualified_name  # pyre-ignore Undefined attribute [16]: `Optional` has no attribute `fully_qualified_name`.
             ] = buffer_idx
+            self.emitter_state.num_external_constants += 1
         # Tensor is mutable with initial state. Place into mutable segment
         elif allocation_info:
             buffer_idx = len(self.program_state.mutable_buffer)
@@ -1358,6 +1360,7 @@ def _emit_prim_getters(self, prim_getters: Dict[str, Any]) -> List[ExecutionPlan
                     delegates=[],
                     non_const_buffer_sizes=[0],
                     container_meta_type=ContainerMetadata("", spec),
+                    num_external_constants=0,
                 )
             )
         return plans
@@ -1739,4 +1742,5 @@ def plan(self) -> ExecutionPlan:
                 self.module.meta["non_const_buffer_sizes"],
             ),
             container_meta_type=self.container_meta_type,
+            num_external_constants=self.emitter_state.num_external_constants,
         )

exir/schema.py

@@ -276,6 +276,7 @@ class ExecutionPlan:
     # Runtime should use the len(constant_buffer) as the ground truch of
     # constant memory buffer size, and ignore non_const_buffer_sizes[0].
     non_const_buffer_sizes: List[int]
+    num_external_constants: int
 
 
 @dataclass

runtime/executor/method.cpp

@@ -289,42 +289,6 @@ 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;
-}
-
 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) {
@@ -342,6 +306,15 @@ Error Method::parse_external_constants(const NamedDataMap* named_data_map) {
   int index = 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) {
@@ -412,11 +385,7 @@ Error Method::parse_values(const NamedDataMap* named_data_map) {
   }
 
   // Check if there are any 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;
+  num_external_constants_ = serialization_plan_->num_external_constants();
   if (num_external_constants_ > 0) {
     // Allocate space for external tensors.
     external_constants_ =

runtime/executor/method.h

@@ -358,11 +358,6 @@ class Method final {
   NamedData* external_constants_;
   size_t num_external_constants_ = 0;
 
-  /**
-   * Counts the number of external constants for this method.
-   */
-  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

schema/program.fbs

@@ -386,6 +386,8 @@ table ExecutionPlan {
   // constants memory buffer size, and ignore non_const_buffer_sizes[0].
   non_const_buffer_sizes: [int64];
 
+  // Number of external constants used in this ExecutionPlan.
+  num_external_constants: uint32;
 }
 
 // Constant tensor data stored directly in the flatbuffer.