From cfd05d6907a7032fa2d8e728f78608b3e57d3526 Mon Sep 17 00:00:00 2001 From: lucylq Date: Thu, 13 Feb 2025 14:42:44 -0800 Subject: [PATCH] [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-poisoned] --- exir/emit/_emit_program.py | 1 + exir/emit/_emitter.py | 4 +++ exir/schema.py | 1 + runtime/executor/method.cpp | 51 ++++++++----------------------------- runtime/executor/method.h | 5 ---- schema/program.fbs | 2 ++ 6 files changed, 18 insertions(+), 46 deletions(-) diff --git a/exir/emit/_emit_program.py b/exir/emit/_emit_program.py index f9571143a1b..2c28379c147 100644 --- a/exir/emit/_emit_program.py +++ b/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) diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py index 0cbc63bde21..1b30b090653 100644 --- a/exir/emit/_emitter.py +++ b/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, ) diff --git a/exir/schema.py b/exir/schema.py index 8e1434a2fe4..a200e97b974 100644 --- a/exir/schema.py +++ b/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 diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index 7eb15820bde..d27907185b4 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -289,42 +289,6 @@ Result parse_cond_value(const EValue& cond_value) { } // namespace -Result 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( - 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 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_ = diff --git a/runtime/executor/method.h b/runtime/executor/method.h index bd05dd54fe6..c42fc4c3563 100644 --- a/runtime/executor/method.h +++ b/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 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 diff --git a/schema/program.fbs b/schema/program.fbs index 7ab2175f8ac..86aa070337b 100644 --- a/schema/program.fbs +++ b/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.