Skip to content

Commit c7458ee

Browse files
shoumikhinshoumikhin
committed
Fix device placement for memory-planned buffers in Runtime.load_program
A .pte file records where each memory-planned buffer has to live, on the host or on an accelerator, in the plan's non_const_buffer_device field. ProgramMemory never read that field and allocated every planned buffer as host memory, so a program asking for device memory received a host pointer. With the CUDA backend this failed at execute time with "not backed by CUDA device memory", while the same .pte loaded through _load_for_executorch worked. Allocate each planned buffer on the device it is tagged for, mirroring make_planned_memory_with_devices in extension/module/module.cpp. Buffers shared across methods must agree on their device, so a disagreement is now rejected with a clear error rather than resolved silently. non_const_buffer_device is optional and MethodMeta reports CPU when it is absent, so CPU-only and older programs keep the previous host allocation path.
1 parent cff6f4d commit c7458ee

1 file changed

Lines changed: 110 additions & 16 deletions

File tree

‎extension/pybindings/pybindings.cpp‎

Lines changed: 110 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
*/
88

99
#include <algorithm>
10+
#include <cinttypes>
1011
#include <cstdio>
1112
#include <iostream>
1213
#include <memory>
@@ -31,6 +32,7 @@
3132
#include <executorch/extension/threadpool/threadpool.h>
3233
#include <executorch/runtime/backend/interface.h>
3334
#include <executorch/runtime/core/data_loader.h>
35+
#include <executorch/runtime/core/device_memory_buffer.h>
3436
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
3537
#include <executorch/runtime/executor/method.h>
3638
#include <executorch/runtime/executor/program.h>
@@ -88,6 +90,7 @@ using ::executorch::extension::MmapDataLoader;
8890
using ::executorch::extension::ET_BUNDLED_MODULE_NAMESPACE::BundledModule;
8991
using ::executorch::extension::pybindings::PyDataLoader;
9092
using ::executorch::runtime::DataLoader;
93+
using ::executorch::runtime::DeviceMemoryBuffer;
9194
using ::executorch::runtime::Error;
9295
using ::executorch::runtime::EValue;
9396
using ::executorch::runtime::EventTracerDebugLogLevel;
@@ -1077,12 +1080,22 @@ inline std::shared_ptr<ProgramState> load_program(
10771080
/// A wrapper/util class for executorch memory allocations/manager.
10781081
class ProgramMemory {
10791082
public:
1080-
explicit ProgramMemory(std::vector<std::vector<uint8_t>>&& non_const_buffers)
1083+
/// @param sizes Byte size of each memory-planned buffer.
1084+
/// @param devices Device each planned buffer must live on, same indexing as
1085+
/// `sizes`. A buffer tagged non-CPU is allocated through the device
1086+
/// allocator; handing the backend a host pointer for it is undefined.
1087+
ProgramMemory(
1088+
std::vector<int64_t>&& sizes,
1089+
std::vector<executorch::aten::Device>&& devices)
10811090
: runtime_allocator_(),
1082-
non_const_buffers_(std::move(non_const_buffers)),
1091+
planned_sizes_(std::move(sizes)),
1092+
planned_devices_(std::move(devices)),
1093+
non_const_buffers_(allocate_host_buffers()),
1094+
device_buffers_(allocate_device_buffers()),
10831095
non_const_spans_(create_non_const_spans()),
10841096
non_const_allocator_(
1085-
{non_const_spans_.data(), non_const_spans_.size()}),
1097+
{non_const_spans_.data(), non_const_spans_.size()},
1098+
{planned_devices_.data(), planned_devices_.size()}),
10861099
mem_manager_(
10871100
&const_allocator_,
10881101
&non_const_allocator_,
@@ -1105,19 +1118,71 @@ class ProgramMemory {
11051118

11061119
MallocMemoryAllocator temp_allocator_{};
11071120

1121+
std::vector<int64_t> planned_sizes_;
1122+
1123+
std::vector<executorch::aten::Device> planned_devices_;
1124+
1125+
// Backs CPU-tagged buffers; the entry is empty for a device-tagged buffer.
11081126
std::vector<std::vector<uint8_t>> non_const_buffers_;
11091127

1128+
// Backs device-tagged buffers; the entry is empty for a CPU-tagged buffer.
1129+
// Parallel to non_const_buffers_ so both index by planned buffer id.
1130+
std::vector<DeviceMemoryBuffer> device_buffers_;
1131+
11101132
std::vector<Span<uint8_t>> non_const_spans_;
11111133

11121134
HierarchicalAllocator non_const_allocator_;
11131135

11141136
MemoryManager mem_manager_;
11151137

1138+
std::vector<std::vector<uint8_t>> allocate_host_buffers() {
1139+
std::vector<std::vector<uint8_t>> result;
1140+
result.reserve(planned_sizes_.size());
1141+
for (size_t i = 0; i < planned_sizes_.size(); i++) {
1142+
if (planned_devices_[i].is_cpu()) {
1143+
result.emplace_back(planned_sizes_[i]);
1144+
} else {
1145+
result.emplace_back();
1146+
}
1147+
}
1148+
return result;
1149+
}
1150+
1151+
std::vector<DeviceMemoryBuffer> allocate_device_buffers() {
1152+
std::vector<DeviceMemoryBuffer> result;
1153+
result.reserve(planned_sizes_.size());
1154+
for (size_t i = 0; i < planned_sizes_.size(); i++) {
1155+
if (planned_devices_[i].is_cpu()) {
1156+
result.emplace_back();
1157+
continue;
1158+
}
1159+
auto buffer = DeviceMemoryBuffer::create(
1160+
planned_sizes_[i],
1161+
planned_devices_[i].type(),
1162+
planned_devices_[i].index());
1163+
THROW_IF_ERROR(
1164+
buffer.error(),
1165+
"Failed to allocate %" PRId64
1166+
" bytes for planned buffer %zu on device type %d index %d",
1167+
planned_sizes_[i],
1168+
i,
1169+
static_cast<int>(planned_devices_[i].type()),
1170+
static_cast<int>(planned_devices_[i].index()));
1171+
result.emplace_back(std::move(buffer.get()));
1172+
}
1173+
return result;
1174+
}
1175+
11161176
std::vector<Span<uint8_t>> create_non_const_spans() {
11171177
std::vector<Span<uint8_t>> result;
1118-
for (size_t i = 0; i < non_const_buffers_.size(); i++) {
1119-
result.push_back(
1120-
{non_const_buffers_[i].data(), non_const_buffers_[i].size()});
1178+
result.reserve(planned_sizes_.size());
1179+
for (size_t i = 0; i < planned_sizes_.size(); i++) {
1180+
if (planned_devices_[i].is_cpu()) {
1181+
result.push_back(
1182+
{non_const_buffers_[i].data(), non_const_buffers_[i].size()});
1183+
} else {
1184+
result.push_back(device_buffers_[i].as_span());
1185+
}
11211186
}
11221187
return result;
11231188
}
@@ -1420,30 +1485,59 @@ struct PyProgram final {
14201485
// method in the program. Map will be easier to use than a list because we
14211486
// dont know how many non_const arenas there will be
14221487
std::map<size_t, int64_t> non_const_buffer_sizes;
1488+
// Device each arena must live on. Arenas are shared across methods, so all
1489+
// methods have to agree; disagreement is rejected below rather than
1490+
// silently resolved, because one arena cannot be both host and device.
1491+
std::map<size_t, executorch::aten::Device> non_const_buffer_devices;
14231492
for (size_t i = 0; i < state_->program_->num_methods(); ++i) {
14241493
auto name = state_->program_->get_method_name(i).get();
14251494
auto method_meta = state_->program_->method_meta(name).get();
1426-
for (size_t j = 0; j < method_meta.num_non_const_buffers(); j++) {
1427-
int64_t buffer_size = method_meta.non_const_buffer_size(j).get();
1495+
for (size_t j = 0; j < method_meta.num_memory_planned_buffers(); j++) {
1496+
int64_t buffer_size = method_meta.memory_planned_buffer_size(j).get();
1497+
auto device_result = method_meta.memory_planned_buffer_device(j);
1498+
THROW_IF_ERROR(
1499+
device_result.error(),
1500+
"Failed to get device of planned buffer %zu of method %s",
1501+
j,
1502+
name);
1503+
auto device = device_result.get();
14281504
if (non_const_buffer_sizes.find(j) == non_const_buffer_sizes.end()) {
14291505
non_const_buffer_sizes.insert({j, buffer_size});
1506+
non_const_buffer_devices.insert({j, device});
14301507
} else {
14311508
non_const_buffer_sizes[j] =
14321509
std::max(non_const_buffer_sizes[j], buffer_size);
1510+
const auto& seen = non_const_buffer_devices.at(j);
1511+
THROW_IF_ERROR(
1512+
(seen.type() == device.type() && seen.index() == device.index())
1513+
? Error::Ok
1514+
: Error::InvalidProgram,
1515+
"Planned buffer %zu is on device type %d index %d for method %s "
1516+
"but device type %d index %d for an earlier method; a shared "
1517+
"arena cannot satisfy both",
1518+
j,
1519+
static_cast<int>(device.type()),
1520+
static_cast<int>(device.index()),
1521+
name,
1522+
static_cast<int>(seen.type()),
1523+
static_cast<int>(seen.index()));
14331524
}
14341525
}
14351526
}
14361527

1437-
// Allocate the arenas. Using vector because we need to remember the size as
1438-
// well, so vector is easier then unique_ptr.
1439-
std::vector<std::vector<uint8_t>> non_const_buffers_;
1440-
for (std::map<size_t, int64_t>::iterator i = non_const_buffer_sizes.begin();
1441-
i != non_const_buffer_sizes.end();
1442-
i++) {
1443-
non_const_buffers_.push_back(std::vector<uint8_t>(i->second));
1528+
// Hand the sizes and devices to ProgramMemory, which allocates each arena
1529+
// on its tagged device.
1530+
std::vector<int64_t> planned_sizes;
1531+
std::vector<executorch::aten::Device> planned_devices;
1532+
planned_sizes.reserve(non_const_buffer_sizes.size());
1533+
planned_devices.reserve(non_const_buffer_sizes.size());
1534+
for (const auto& entry : non_const_buffer_sizes) {
1535+
planned_sizes.push_back(entry.second);
1536+
planned_devices.push_back(non_const_buffer_devices.at(entry.first));
14441537
}
14451538

1446-
memory_ = std::make_shared<ProgramMemory>(std::move(non_const_buffers_));
1539+
memory_ = std::make_shared<ProgramMemory>(
1540+
std::move(planned_sizes), std::move(planned_devices));
14471541
if (event_tracer_ && debug_buffer_size > 0) {
14481542
// If a debug buffer was requested for the ETDump, allocate it and make
14491543
// sure its lifetime is as long as the event_tracer.

0 commit comments

Comments
 (0)