Skip to content

[wip] 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

Open
wants to merge 1 commit into
base: gh/lucylq/88/base
Choose a base branch
from
Open

[wip] Add MergedDataMap to method #12088

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
66 changes: 40 additions & 26 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) {
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,23 +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();
}

// Merge NamedDataMaps.
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;
// Merge them.
if (external_data_map && pte_data_map.ok()) {
const std::array<const NamedDataMap*, 2> data_maps = {
external_data_map, pte_data_map.ok() ? pte_data_map.get() : nullptr};

auto merged = MergedDataMap<2>::load(data_maps);
if (!merged.ok()) {
return merged.error();
}
merged_data_map_ = method_allocator->allocateInstance<MergedDataMap<2>>();
if (merged_data_map_ == nullptr) {
return Error::MemoryAllocationFailed;
}
new (merged_data_map_) MergedDataMap<2>(std::move(merged.get()));
} 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
// ~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.
Expand Down Expand Up @@ -1680,6 +1692,8 @@ Method::~Method() {
for (const auto i : c10::irange(n_external_constants_)) {
external_constants_[i].buffer.~FreeableBuffer();
}
// Free the MergedDataMap
merged_data_map_->~MergedDataMap<2>();
// 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_;

MergedDataMap<2>* 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