diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index d435678ca2b..0857bc1c976 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -33,6 +33,7 @@ namespace executorch { namespace runtime { +using deserialization::NamedData; using internal::PlatformMemoryAllocator; /** @@ -289,6 +290,113 @@ 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 n_external_constants = 0; + for (size_t i = 0; i < n_value; ++i) { + auto serialization_value = flatbuffer_values->Get(i); + // Ensure values are non-null. + // Note that as a side-effect of this check, we're guaranteed that all + // values are non-null, so later loops can skip that check. + 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) { + n_external_constants++; + } + } + return n_external_constants; +} + +Error Method::parse_external_constants(const NamedDataMap* named_data_map) { + auto flatbuffer_values = serialization_plan_->values(); + size_t n_value = flatbuffer_values->size(); + + // n_external_constants_ counts the number of successfully-initialized + // external constants 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_external_constants_ = 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( + serialization_value->val()); + // Constant tensors are resolved here; tensors with allocation_info are + // mutable and are 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) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + s_tensor->extra_tensor_info()->fully_qualified_name() != nullptr, + InvalidExternalData, + "Fully qualified name of external tensor is null at index %zu", + i); + + const char* key = + s_tensor->extra_tensor_info()->fully_qualified_name()->c_str(); + + // Check if this tensor has already been resolved. + if (get_data_by_key( + key, Span(external_constants_, n_external_constants_)) != + nullptr) { + continue; + } + Result tensor_layout = + named_data_map->get_metadata(key); + if (!tensor_layout.ok()) { + return tensor_layout.error(); + } + // Check external tensor compatibility. + Error err = + deserialization::validateTensorLayout(s_tensor, tensor_layout.get()); + if (err != Error::Ok) { + return err; + } + // Save the key. + external_constants_[n_external_constants_].key = key; + + // Save the buffer. + Result buffer = named_data_map->get_data(key); + ET_CHECK_OR_RETURN_ERROR( + buffer.ok(), + InvalidExternalData, + "Buffer retrieved from get_data is not valid"); + new (&external_constants_[n_external_constants_].buffer) + FreeableBuffer(std::move(buffer.get())); + + n_external_constants_ += 1; + } + return Error::Ok; +} + Error Method::parse_values(const NamedDataMap* named_data_map) { auto flatbuffer_values = serialization_plan_->values(); ET_CHECK_OR_RETURN_ERROR( @@ -299,6 +407,30 @@ Error Method::parse_values(const NamedDataMap* named_data_map) { return Error::MemoryAllocationFailed; } + // Count the number of tensors marked as EXTERNAL for this method. The actual + // number of external constants may be smaller, eg. if multiple tensors point + // to the same underlying data buffer. + // This function also ensures that all flatbuffer_values entries + // are non-null, so `val_as_X()` calls below are guaranteed to return + // non-null pointers. + Result max_external_constants = get_num_external_constants(); + if (!max_external_constants.ok()) { + return max_external_constants.error(); + } + if (max_external_constants.get() > 0) { + // Allocate space for external tensors. + external_constants_ = + memory_manager_->method_allocator()->allocateList( + max_external_constants.get()); + if (external_constants_ == nullptr) { + return Error::MemoryAllocationFailed; + } + Error 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. @@ -306,16 +438,6 @@ Error Method::parse_values(const NamedDataMap* named_data_map) { 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 +538,8 @@ Error Method::parse_values(const NamedDataMap* named_data_map) { program_, memory_manager_, static_cast(val), - named_data_map); + named_data_map, + Span(external_constants_, n_external_constants_)); if (!t.ok()) { ET_LOG( Error, @@ -1496,6 +1619,10 @@ Method::~Method() { delegates_[i].~BackendDelegate(); } } + // Free resources associated with external constants. + for (int i = 0; i < n_external_constants_; i++) { + external_constants_[i].buffer.~FreeableBuffer(); + } // All other fields are trivially destructible. } } // namespace runtime diff --git a/runtime/executor/method.h b/runtime/executor/method.h index dff4e818f9f..4108db8810e 100644 --- a/runtime/executor/method.h +++ b/runtime/executor/method.h @@ -31,6 +31,12 @@ struct EValue; namespace executorch { namespace runtime { +// Forward declare NamedData. This is a public header and must not include +// internal data types. +namespace deserialization { +struct NamedData; +} // namespace deserialization + // Forward declare Program to avoid a circular reference. class Program; @@ -42,6 +48,7 @@ using OpFunction = void (*)(KernelRuntimeContext&, EValue**); /// A list of pointers into the master values table that together compose the /// argument list for a single instruction using InstructionArgs = Span; +using deserialization::NamedData; /** * An executable method of an executorch program. Maps to a python method like @@ -66,6 +73,8 @@ class Method final { delegates_(rhs.delegates_), n_chains_(rhs.n_chains_), chains_(rhs.chains_), + external_constants_(rhs.external_constants_), + n_external_constants_(rhs.n_external_constants_), init_state_(rhs.init_state_) { // Required: clear out fields that the dtor looks at, so that we don't free // anything twice. @@ -73,6 +82,8 @@ class Method final { rhs.values_ = nullptr; rhs.n_delegate_ = 0; rhs.delegates_ = nullptr; + rhs.n_external_constants_ = 0; + rhs.external_constants_ = nullptr; // Helpful: Try to ensure that any other interactions with the old object // result in failures. @@ -288,6 +299,8 @@ class Method final { delegates_(nullptr), n_chains_(0), chains_(nullptr), + external_constants_(nullptr), + n_external_constants_(0), init_state_(InitializationState::Uninitialized) {} /// Static factory used by Program. @@ -336,8 +349,31 @@ class Method final { size_t n_chains_; Chain* chains_; + NamedData* external_constants_; + size_t n_external_constants_ = 0; + InitializationState init_state_; + /** + * Counts the number of tensors marked as EXTERNAL in the flatbuffer + * 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 + * into `external_constants_`. Updates `n_external_constants_` to count the + * number of successfully-initialized external constants. + * FreeableBuffers returned by the named_data_map are owned by the + * method and are freed on method destruction. + * + * @param[in] named_data_map, to retrieve external constants from. + * @returns Error::Ok on success, non-Ok on failure. + */ + 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 diff --git a/runtime/executor/tensor_parser.h b/runtime/executor/tensor_parser.h index 2ffb473544d..cfd711713ac 100644 --- a/runtime/executor/tensor_parser.h +++ b/runtime/executor/tensor_parser.h @@ -21,11 +21,21 @@ namespace executorch { namespace runtime { namespace deserialization { +/// Data structure to hold key and data buffer for external data used +/// in a method. +struct NamedData { + const char* key; + FreeableBuffer buffer; +}; + +NamedData* get_data_by_key(const char* key, Span entries); + ET_NODISCARD Result 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 external_constants = {}); ET_NODISCARD Result> parseTensorList( const flatbuffers::Vector* tensor_indices, @@ -33,6 +43,12 @@ ET_NODISCARD Result> parseTensorList( 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 validateTensorLayout( + const executorch_flatbuffer::Tensor* s_tensor, + const TensorLayout& expected_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 +121,11 @@ 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 containing tensor fqn to + * corresponding tensor data. Used to resolve data that is constant and + * external to the PTE, if any. Referencing data from external_constants is + * safe, as it has the same lifetime as the method. * * @returns On success, the data pointer to use for the tensor. On failure, a * non-Ok Error. @@ -115,7 +135,8 @@ ET_NODISCARD Result getTensorDataPtr( const Program* program, size_t nbytes, HierarchicalAllocator* allocator, - const NamedDataMap* named_data_map = nullptr); + const NamedDataMap* named_data_map = nullptr, + Span external_constants = {}); } // namespace deserialization } // namespace runtime diff --git a/runtime/executor/tensor_parser_aten.cpp b/runtime/executor/tensor_parser_aten.cpp index ab9af3d0399..d1a2f712853 100644 --- a/runtime/executor/tensor_parser_aten.cpp +++ b/runtime/executor/tensor_parser_aten.cpp @@ -33,7 +33,8 @@ Result parseTensor( const Program* program, MemoryManager* memory_manager, const executorch_flatbuffer::Tensor* s_tensor, - const NamedDataMap* named_data_map) { + const NamedDataMap* named_data_map, + Span external_constants) { EXECUTORCH_SCOPE_PROF("TensorParser::parseTensor"); ET_CHECK_OR_RETURN_ERROR( @@ -108,7 +109,8 @@ Result parseTensor( program, tensor.nbytes(), memory_manager->planned_memory(), - named_data_map); + named_data_map, + external_constants); if (!data_ptr.ok()) { ET_LOG( Error, diff --git a/runtime/executor/tensor_parser_exec_aten.cpp b/runtime/executor/tensor_parser_exec_aten.cpp index 83310ff680c..a1ac245acca 100644 --- a/runtime/executor/tensor_parser_exec_aten.cpp +++ b/runtime/executor/tensor_parser_exec_aten.cpp @@ -111,12 +111,60 @@ ET_NODISCARD Result> parseTensorList( evalp_list, tensor_list, tensor_indices->size()); } +ET_NODISCARD Error validateTensorLayout( + const executorch_flatbuffer::Tensor* s_tensor, + const TensorLayout& expected_layout) { + ET_CHECK_OR_RETURN_ERROR( + static_cast(s_tensor->scalar_type()) == + expected_layout.scalar_type(), + InvalidExternalData, + "Scalar type mismatch. Expected %hhd, got %hhd.", + static_cast(s_tensor->scalar_type()), + static_cast(expected_layout.scalar_type())); + int dim = s_tensor->sizes()->size(); + ET_CHECK_OR_RETURN_ERROR( + dim == expected_layout.sizes().size(), + InvalidExternalData, + "Dim mismatch. Expected %d, got %zu.", + dim, + expected_layout.sizes().size()); + for (int i = 0; i < dim; i++) { + ET_CHECK_OR_RETURN_ERROR( + s_tensor->sizes()->Get(i) == expected_layout.sizes()[i], + InvalidExternalData, + "Sizes mismatch. Expected %d, got %d for size at index %d.", + s_tensor->sizes()->Get(i), + expected_layout.sizes()[i], + i); + ET_CHECK_OR_RETURN_ERROR( + s_tensor->dim_order()->Get(i) == expected_layout.dim_order()[i], + InvalidExternalData, + "Dim order mismatch. Expected %d, got %d for dim at index %d.", + s_tensor->dim_order()->Get(i), + expected_layout.dim_order()[i], + i); + } + return Error::Ok; +} + +// Check if key exists in entries. If it does, return a pointer to the entry +// otherwise return a nullptr. +NamedData* get_data_by_key(const char* key, Span entries) { + for (int i = 0; i < entries.size(); i++) { + if (strcmp(key, entries[i].key) == 0) { + return &entries[i]; + } + } + return nullptr; +} + ET_NODISCARD Result getTensorDataPtr( const executorch_flatbuffer::Tensor* s_tensor, const Program* program, size_t nbytes, HierarchicalAllocator* allocator, - const NamedDataMap* named_data_map) { + const NamedDataMap* named_data_map, + Span external_constants) { auto data_buffer_idx = s_tensor->data_buffer_idx(); const executorch_flatbuffer::AllocationDetails* allocation_info = s_tensor->allocation_info(); @@ -146,76 +194,38 @@ ET_NODISCARD Result getTensorDataPtr( s_tensor->extra_tensor_info()->fully_qualified_name() != nullptr, InvalidExternalData, "Fully qualified name of external tensor is null"); - // Look up tensor in named data map. - Result tensor_layout_res = named_data_map->get_metadata( - s_tensor->extra_tensor_info()->fully_qualified_name()->c_str()); - if (!tensor_layout_res.ok()) { - return tensor_layout_res.error(); - } - const TensorLayout& tensor_layout = tensor_layout_res.get(); - - // Compatibility checking. - ET_CHECK_OR_RETURN_ERROR( - static_cast(s_tensor->scalar_type()) == - tensor_layout.scalar_type(), - InvalidExternalData, - "Scalar type mismatch. Expected %hhd, got %hhd.", - static_cast(s_tensor->scalar_type()), - static_cast(tensor_layout.scalar_type())); - ET_CHECK_OR_RETURN_ERROR( - nbytes == tensor_layout.nbytes(), - InvalidExternalData, - "Nbytes mismatch. Expected %zu, got %zu.", - nbytes, - tensor_layout.nbytes()); - int dim = s_tensor->sizes()->size(); - ET_CHECK_OR_RETURN_ERROR( - dim == tensor_layout.sizes().size(), - InvalidExternalData, - "Dim mismatch. Expected %d, got %zu.", - dim, - tensor_layout.sizes().size()); - for (int i = 0; i < dim; i++) { - ET_CHECK_OR_RETURN_ERROR( - s_tensor->sizes()->Get(i) == tensor_layout.sizes()[i], - InvalidExternalData, - "Sizes mismatch. Expected %d, got %d for size at index %d.", - s_tensor->sizes()->Get(i), - tensor_layout.sizes()[i], - i); - ET_CHECK_OR_RETURN_ERROR( - s_tensor->dim_order()->Get(i) == tensor_layout.dim_order()[i], - InvalidExternalData, - "Dim order mismatch. Expected %d, got %d for dim at index %d.", - s_tensor->dim_order()->Get(i), - tensor_layout.dim_order()[i], - i); - } + const char* fqn = + s_tensor->extra_tensor_info()->fully_qualified_name()->c_str(); // Constant value. if (allocation_info == nullptr) { - Result data_res = named_data_map->get_data( - s_tensor->extra_tensor_info()->fully_qualified_name()->c_str()); - if (!data_res.ok()) { - return data_res.error(); + NamedData* data = get_data_by_key(fqn, external_constants); + if (data != nullptr) { + return const_cast(data->buffer.data()); + } + // Should never reach here; these tensors are resolved in + // Method::parse_external_constants. Any errors should be caught there. + return Error::Internal; + } else { + // Mutable value. + // Look up tensor in named data map. + Result tensor_layout_res = + named_data_map->get_metadata(fqn); + if (!tensor_layout_res.ok()) { + return tensor_layout_res.error(); + } + const TensorLayout& tensor_layout = tensor_layout_res.get(); + Error err = validateTensorLayout(s_tensor, tensor_layout); + if (err != Error::Ok) { + return err; } - // The const_cast is 'ok' here because program and runtime should - // guarantee that this data is never modified. Temporary until runtime - // takes ownership of FreeableBuffers in TODO(T214294528). - return const_cast(data_res.get().data()); - } - - // Mutable value. - else { // Call load_into. auto planned_ptr = getMemPlannedPtr(allocation_info, nbytes, allocator); if (!planned_ptr.ok()) { return planned_ptr.error(); } - auto size = named_data_map->load_data_into( - s_tensor->extra_tensor_info()->fully_qualified_name()->c_str(), - planned_ptr.get(), - nbytes); + auto size = + named_data_map->load_data_into(fqn, planned_ptr.get(), nbytes); if (size.error() != Error::Ok) { return size.error(); } diff --git a/runtime/executor/tensor_parser_portable.cpp b/runtime/executor/tensor_parser_portable.cpp index a53295470fc..3a29c86700c 100644 --- a/runtime/executor/tensor_parser_portable.cpp +++ b/runtime/executor/tensor_parser_portable.cpp @@ -21,6 +21,7 @@ namespace executorch { namespace runtime { namespace deserialization { +using executorch::runtime::Span; using torch::executor::ScalarType; using torch::executor::Tensor; using torch::executor::TensorImpl; @@ -29,7 +30,8 @@ Result parseTensor( const Program* program, MemoryManager* memory_manager, const executorch_flatbuffer::Tensor* s_tensor, - const NamedDataMap* named_data_map) { + const NamedDataMap* named_data_map, + Span external_constants) { EXECUTORCH_SCOPE_PROF("TensorParser::parseTensor"); auto method_allocator = memory_manager->method_allocator(); @@ -149,7 +151,8 @@ Result parseTensor( program, tensor_impl->nbytes(), memory_manager->planned_memory(), - named_data_map); + named_data_map, + external_constants); if (!data_ptr.ok()) { ET_LOG( Error,