Skip to content

Add MergedDataMap to method #12088

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jul 9, 2025
Merged

Add MergedDataMap to method #12088

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
19c16cb
Add MergedDataMap to method
lucylq Jun 28, 2025
a37a615
Update on "[wip] Add MergedDataMap to method"
lucylq Jul 2, 2025
9d83166
Update on "[wip] Add MergedDataMap to method"
lucylq Jul 2, 2025
f5da980
Update on "[wip] Add MergedDataMap to method"
lucylq Jul 8, 2025
6e9429a
Update on "[wip] Add MergedDataMap to method"
lucylq Jul 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 43 additions & 25 deletions runtime/executor/method.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#include <executorch/runtime/core/named_data_map.h>
#include <executorch/runtime/core/span.h>
#include <executorch/runtime/executor/memory_manager.h>
#include <executorch/runtime/executor/merged_data_map.h>
#include <executorch/runtime/executor/platform_memory_allocator.h>
#include <executorch/runtime/executor/program.h>
#include <executorch/runtime/executor/tensor_parser.h>
Expand Down Expand Up @@ -328,9 +329,9 @@ Result<size_t> Method::get_num_external_constants() {
return n_external_constants;
}

Error Method::parse_external_constants(const NamedDataMap* named_data_map) {
Error Method::parse_external_constants(const NamedDataMap* external_data_map) {
ET_CHECK_OR_RETURN_ERROR(
named_data_map != nullptr, InvalidState, "named_data_map is null");
external_data_map != nullptr, InvalidState, "external_data_map is null");
auto flatbuffer_values = serialization_plan_->values();
size_t n_value = flatbuffer_values->size();

Expand Down Expand Up @@ -372,7 +373,7 @@ Error Method::parse_external_constants(const NamedDataMap* named_data_map) {
continue;
}
Result<const TensorLayout> tensor_layout =
named_data_map->get_tensor_layout(key);
external_data_map->get_tensor_layout(key);
if (!tensor_layout.ok()) {
ET_LOG(Info, "Failed to get metadata for key %s", key);
return tensor_layout.error();
Expand All @@ -387,7 +388,7 @@ Error Method::parse_external_constants(const NamedDataMap* named_data_map) {
external_constants_[n_external_constants_].key = key;

// Save the buffer.
Result<FreeableBuffer> buffer = named_data_map->get_data(key);
Result<FreeableBuffer> buffer = external_data_map->get_data(key);
ET_CHECK_OR_RETURN_ERROR(
buffer.ok(),
InvalidExternalData,
Expand All @@ -400,7 +401,7 @@ Error Method::parse_external_constants(const NamedDataMap* named_data_map) {
return Error::Ok;
}

Error Method::parse_values(const NamedDataMap* named_data_map) {
Error Method::parse_values(const NamedDataMap* external_data_map) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel like named was the better name since it possibly contains data embedded in the pte

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oops, didn't see this. This data map comes from the PTD file (external), so it shouldn't contain items from the PTE --> that's the pte_data_map

auto flatbuffer_values = serialization_plan_->values();
ET_CHECK_OR_RETURN_ERROR(
flatbuffer_values != nullptr, InvalidProgram, "Missing values");
Expand Down Expand Up @@ -428,7 +429,7 @@ Error Method::parse_values(const NamedDataMap* named_data_map) {
if (external_constants_ == nullptr) {
return Error::MemoryAllocationFailed;
}
Error err = parse_external_constants(named_data_map);
Error err = parse_external_constants(external_data_map);
if (err != Error::Ok) {
return err;
}
Expand Down Expand Up @@ -541,7 +542,7 @@ Error Method::parse_values(const NamedDataMap* named_data_map) {
program_,
memory_manager_,
static_cast<const executorch_flatbuffer::Tensor*>(val),
named_data_map,
external_data_map,
Span<NamedData>(external_constants_, n_external_constants_));
if (!t.ok()) {
ET_LOG(
Expand Down Expand Up @@ -741,7 +742,7 @@ Result<Method> Method::load(
const Program* program,
MemoryManager* memory_manager,
EventTracer* event_tracer,
const NamedDataMap* named_data_map) {
const NamedDataMap* external_data_map) {
MemoryAllocator* temp_allocator = memory_manager->temp_allocator();
if (temp_allocator == nullptr) {
PlatformMemoryAllocator* platform_allocator =
Expand All @@ -755,7 +756,7 @@ Result<Method> Method::load(
}
Method method(program, memory_manager, event_tracer, temp_allocator);
ET_LOG(Debug, "Loading method: %s.", s_plan->name()->c_str());
Error err = method.init(s_plan, named_data_map);
Error err = method.init(s_plan, external_data_map);
if (err != Error::Ok) {
return err;
} else {
Expand All @@ -766,7 +767,7 @@ Result<Method> Method::load(

Error Method::init(
executorch_flatbuffer::ExecutionPlan* s_plan,
const NamedDataMap* named_data_map) {
const NamedDataMap* external_data_map) {
EXECUTORCH_SCOPE_PROF("Method::init");
internal::EventTracerProfileMethodScope event_tracer_profile_scope =
internal::EventTracerProfileMethodScope(event_tracer_, "Method::init");
Expand All @@ -783,7 +784,7 @@ Error Method::init(

{
// Parse the elements of the values_ array.
Error err = parse_values(named_data_map);
Error err = parse_values(external_data_map);
if (err != Error::Ok) {
return err;
}
Expand All @@ -800,21 +801,34 @@ Error Method::init(
return Error::MemoryAllocationFailed;
}

// Get NamedDataMap, if it exists.
const NamedDataMap* pte_data_map = nullptr;
Result<const NamedDataMap*> pte_data_map_res =
program_->get_named_data_map();
if (pte_data_map_res.ok()) {
pte_data_map = pte_data_map_res.get();
}

// Get PTE data map, if it exists.
auto pte_data_map = program_->get_named_data_map();
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;
pte_data_map.ok() || pte_data_map.error() == Error::NotFound,
InvalidProgram,
"Failed to get named data map from program: 0x%" PRIx32,
static_cast<uint32_t>(pte_data_map.error()));

const NamedDataMap* named_data_map = nullptr;
if (external_data_map && pte_data_map.ok()) {
// Merge external_data_map and pte_data_map if both are present.
auto merged =
internal::MergedDataMap::load(external_data_map, pte_data_map.get());
if (!merged.ok()) {
return merged.error();
}
// Allocate memory for the merged data map.
merged_data_map_ =
method_allocator->allocateInstance<internal::MergedDataMap>();
if (merged_data_map_ == nullptr) {
return Error::MemoryAllocationFailed;
}
new (merged_data_map_) internal::MergedDataMap(std::move(merged.get()));
named_data_map = merged_data_map_;
} else if (external_data_map) {
named_data_map = external_data_map;
} else if (pte_data_map.ok()) {
named_data_map = pte_data_map.get();
}

// n_delegate_ counts the number of successfully-initialized delegates for
Expand Down Expand Up @@ -1680,6 +1694,10 @@ Method::~Method() {
for (const auto i : c10::irange(n_external_constants_)) {
external_constants_[i].buffer.~FreeableBuffer();
}
// Free the MergedDataMap.
Copy link
Contributor

@JacobSzwejbka JacobSzwejbka Jul 9, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess this is fine. It doesnt do anything though. The MergedMap is trivially destructable

if (merged_data_map_ != nullptr) {
merged_data_map_->~MergedDataMap();
}
// All other fields are trivially destructible.
}
} // namespace ET_RUNTIME_NAMESPACE
Expand Down
6 changes: 6 additions & 0 deletions runtime/executor/method.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#include <executorch/runtime/core/named_data_map.h>
#include <executorch/runtime/core/span.h>
#include <executorch/runtime/executor/memory_manager.h>
#include <executorch/runtime/executor/merged_data_map.h>
#include <executorch/runtime/executor/method_meta.h>
#include <executorch/runtime/platform/compiler.h>

Expand Down Expand Up @@ -76,6 +77,7 @@ class Method final {
delegates_(rhs.delegates_),
n_chains_(rhs.n_chains_),
chains_(rhs.chains_),
merged_data_map_(std::move(rhs.merged_data_map_)),
external_constants_(rhs.external_constants_),
n_external_constants_(rhs.n_external_constants_),
init_state_(rhs.init_state_) {
Expand All @@ -85,6 +87,8 @@ class Method final {
rhs.values_ = nullptr;
rhs.n_delegate_ = 0;
rhs.delegates_ = nullptr;

rhs.merged_data_map_ = nullptr;
rhs.n_external_constants_ = 0;
rhs.external_constants_ = nullptr;

Expand Down Expand Up @@ -314,6 +318,7 @@ class Method final {
delegates_(nullptr),
n_chains_(0),
chains_(nullptr),
merged_data_map_(nullptr),
external_constants_(nullptr),
n_external_constants_(0),
init_state_(InitializationState::Uninitialized) {}
Expand Down Expand Up @@ -364,6 +369,7 @@ class Method final {
size_t n_chains_;
Chain* chains_;

internal::MergedDataMap* merged_data_map_;
NamedData* external_constants_;
size_t n_external_constants_ = 0;

Expand Down
1 change: 1 addition & 0 deletions runtime/executor/targets.bzl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,7 @@ def define_common_targets():
exported_deps = [
":memory_manager",
":pte_data_map" + aten_suffix,
":merged_data_map" + aten_suffix,
"//executorch/runtime/backend:interface" + aten_suffix,
"//executorch/runtime/core:core",
"//executorch/runtime/core:named_data_map" + aten_suffix,
Expand Down
1 change: 1 addition & 0 deletions runtime/executor/test/targets.bzl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,7 @@ def define_common_targets(is_fbcode = False):
],
deps = [
":managed_memory_manager",
"//executorch/runtime/executor:merged_data_map",
"//executorch/runtime/executor:program",
"//executorch/extension/data_loader:file_data_loader",
"//executorch/extension/flat_tensor:flat_tensor_data_map",
Expand Down
Loading