From 5894ea0ad9f5ace52e630def98bca0059af74835 Mon Sep 17 00:00:00 2001 From: lucylq Date: Wed, 30 Apr 2025 13:15:34 -0700 Subject: [PATCH 1/2] Support named_data in flat_tensor Pull Request resolved: https://github.com/pytorch/executorch/pull/10527 Currently flat_tensor ndm only accounts for tensors in get_data, get_num_keys, get_key functions. Add support to return named_data values as well. TODO:consolidate tensors and named_data into one structure in the flatbuffer. This will simplify all the serialization and runtime code. ghstack-source-id: 280786211 Differential Revision: [D73679683](https://our.internmc.facebook.com/intern/diff/D73679683/) --- .../flat_tensor/flat_tensor_data_map.cpp | 83 +++++++++++++++++-- 1 file changed, 78 insertions(+), 5 deletions(-) diff --git a/extension/flat_tensor/flat_tensor_data_map.cpp b/extension/flat_tensor/flat_tensor_data_map.cpp index 8aa0af13928..c5590cb61b1 100644 --- a/extension/flat_tensor/flat_tensor_data_map.cpp +++ b/extension/flat_tensor/flat_tensor_data_map.cpp @@ -44,6 +44,28 @@ bool is_aligned(const void* data) { return addr % kMinimumAlignment == 0; } +Result get_named_data( + const char* key, + const flatbuffers::Vector< + flatbuffers::Offset>* named_data) { + // Linear search by name. + if (named_data == nullptr) { + return Error::NotFound; + } + for (int i = 0; i < named_data->size(); i++) { + if (std::strcmp(named_data->Get(i)->key()->c_str(), key) == 0) { + const auto* metadata = named_data->Get(i); + ET_CHECK_OR_RETURN_ERROR( + metadata->segment_index() >= 0, + InvalidExternalData, + "Invalid segment_index %d; malformed PTD file.", + metadata->segment_index()); + return metadata; + } + } + return Error::NotFound; +} + Result get_flat_tensor_metadata( const char* key, const flatbuffers::Vector< @@ -109,6 +131,39 @@ ET_NODISCARD Result FlatTensorDataMap::get_metadata( ET_NODISCARD Result FlatTensorDataMap::get_data( const char* key) const { + // TODO(lfq): consolidate named_data and tensors. + // Check named data. + Result named_data = + get_named_data(key, flat_tensor_->named_data()); + if (named_data.ok()) { + size_t segment_index = named_data.get()->segment_index(); + ET_CHECK_OR_RETURN_ERROR( + segment_index < flat_tensor_->segments()->size(), + InvalidExternalData, + "Invalid segment_index %zu; malformed PTD file.", + segment_index); + + size_t segment_offset = + flat_tensor_->segments()->Get(segment_index)->offset(); + size_t segment_size = flat_tensor_->segments()->Get(segment_index)->size(); + ET_CHECK_OR_RETURN_ERROR( + segment_offset < + header_.segment_base_offset + header_.segment_data_size, + InvalidExternalData, + "Invalid segment offset %zu is larger than the segment_base_offset + segment_data_size %" PRIu64 + "; malformed PTD file.", + segment_offset, + header_.segment_base_offset + header_.segment_data_size); + return loader_->load( + /*offset=*/header_.segment_base_offset + segment_offset, + segment_size, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External)); + } + if (named_data.error() != Error::NotFound) { + return named_data.error(); + } + + // Check tensors, if named data is not found. Result metadata = get_flat_tensor_metadata(key, flat_tensor_->tensors()); if (!metadata.ok()) { @@ -179,16 +234,34 @@ ET_NODISCARD Error FlatTensorDataMap::load_data_into( } ET_NODISCARD Result FlatTensorDataMap::get_num_keys() const { - return flat_tensor_->tensors()->size(); + // TODO(lfq): consolidate named_data and tensors. + if (flat_tensor_->named_data() == nullptr) { + return flat_tensor_->tensors()->size(); + } + return flat_tensor_->named_data()->size() + flat_tensor_->tensors()->size(); } ET_NODISCARD Result FlatTensorDataMap::get_key( size_t index) const { - if (index < 0 || index >= flat_tensor_->tensors()->size()) { - return Error::InvalidArgument; - } + // TODO(lfq): consolidate named_data and tensors. + // For now, iterate over named_data and then flat_tensor. + size_t num_keys = get_num_keys().get(); + ET_CHECK_OR_RETURN_ERROR( + index >= 0 && index < num_keys, + InvalidArgument, + "Index %zu out of range of size %zu", + index, + num_keys); - return flat_tensor_->tensors()->Get(index)->fully_qualified_name()->c_str(); + if (flat_tensor_->named_data() != nullptr && + index < flat_tensor_->named_data()->size()) { + return flat_tensor_->named_data()->Get(index)->key()->c_str(); + } else { + if (flat_tensor_->named_data() != nullptr) { + index = index - flat_tensor_->named_data()->size(); + } + return flat_tensor_->tensors()->Get(index)->fully_qualified_name()->c_str(); + } } /* static */ Result FlatTensorDataMap::load( From 333b7294d2bc4024c59bc0583e4e7c8aa2b46932 Mon Sep 17 00:00:00 2001 From: lucylq Date: Wed, 30 Apr 2025 13:15:36 -0700 Subject: [PATCH 2/2] Pass one NDM to backend init Pull Request resolved: https://github.com/pytorch/executorch/pull/10528 Take external NDM if it exists, otherwise internal. Note: by default, xnnpack uses the named_data_map. Constants are not stored with the delegated blob anymore. ghstack-source-id: 281215638 Differential Revision: [D73679710](https://our.internmc.facebook.com/intern/diff/D73679710/) --- runtime/executor/method.cpp | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index f09af8ac2e7..b7cd7fa8d12 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -805,6 +805,15 @@ Error Method::init( pte_data_map = pte_data_map_res.get(); } + ET_CHECK_OR_RETURN_ERROR( + !(pte_data_map && named_data_map), + NotSupported, + "NamedDataMap merge not supported; both pte_data_map and named_data_map are non-empty. If you see this error please file an issue at https://github.com/pytorch/executorch/issues"); + + if (!named_data_map || named_data_map->get_num_keys().get() == 0) { + named_data_map = pte_data_map; + } + // n_delegate_ counts the number of successfully-initialized delegates 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. @@ -816,7 +825,7 @@ Error Method::init( method_allocator, /*event_tracer=*/event_tracer_, /*method_name=*/serialization_plan_->name()->c_str(), - /*named_data_map=*/pte_data_map); + /*named_data_map=*/named_data_map); Error err = BackendDelegate::Init( delegate, program_, backend_init_context, &delegates_[i]); if (err != Error::Ok) {