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

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.