Fix device placement for memory-planned buffers in Runtime.load_program

.github/workflows/cuda.yml

@@ -26,6 +26,8 @@ on:
       - .ci/scripts/export_model_artifact.sh
       - .ci/scripts/test_model_e2e.sh
       - examples/models/muse-glimmer/**
+      - extension/pybindings/**
+      - runtime/__init__.py
   workflow_dispatch:
 
 concurrency:
@@ -160,13 +162,17 @@ jobs:
   unittest-cuda:
     name: unittest-cuda
     needs: [changed-files, run-decision]
+    # This job runs the pybindings device test, so it also has to fire on the
+    # two pybindings paths the workflow filter lists.
     if: |
       contains(needs.changed-files.outputs.changed-files, 'backends/cuda') ||
       contains(needs.changed-files.outputs.changed-files, 'backends/aoti') ||
       contains(needs.changed-files.outputs.changed-files, '.github/workflows/cuda.yml') ||
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test-cuda-build.sh') ||
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') ||
       contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') ||
+      contains(needs.changed-files.outputs.changed-files, 'extension/pybindings') ||
+      contains(needs.changed-files.outputs.changed-files, 'runtime/__init__.py') ||
       needs.run-decision.outputs.is-full-run == 'true'
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     permissions:
@@ -190,6 +196,29 @@ jobs:
         conda install -y -c conda-forge 'libstdcxx-ng>=12'
         export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH
 
+        # The pybindings loader only asks for device memory on a build that has
+        # a device allocator registered, so this test cannot run in the CPU
+        # jobs. It needs the install above and nothing built below, so run it
+        # here and fail before the long builds.
+        #
+        # unittest exits 0 when a test skips, and this test skips when the CUDA
+        # backend is missing or no GPU is visible. Both are the prerequisites
+        # this job exists to provide, so treat a skip as a failure rather than
+        # let a packaging regression pass silently. `tee` would otherwise hide
+        # the real exit status behind its own, so read it back explicitly.
+        set +e
+        python -m unittest -v executorch.extension.pybindings.test.test_pybindings.PybindingsTest.test_device_planned_method_allocates_on_the_device 2>&1 | tee /tmp/pybindings_device_test.log
+        pybindings_device_status=${PIPESTATUS[0]}
+        set -e
+        if [ "${pybindings_device_status}" -ne 0 ]; then
+          echo "::error::the pybindings device test failed"
+          exit "${pybindings_device_status}"
+        fi
+        if grep -qE "^OK \(skipped=|skipped=[1-9]" /tmp/pybindings_device_test.log; then
+          echo "::error::the pybindings device test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing"
+          exit 1
+        fi
+
         # Build ExecuTorch with CUDA support
         cmake --workflow --preset llm-release-cuda
 

extension/pybindings/pybindings.cpp

@@ -7,6 +7,7 @@
  */
 
 #include <algorithm>
+#include <cinttypes>
 #include <cstdio>
 #include <iostream>
 #include <memory>
@@ -31,6 +32,7 @@
 #include <executorch/extension/threadpool/threadpool.h>
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/data_loader.h>
+#include <executorch/runtime/core/device_memory_buffer.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 #include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/program.h>
@@ -81,13 +83,15 @@ using ::executorch::ET_RUNTIME_NAMESPACE::get_num_registered_backends;
 using ::executorch::ET_RUNTIME_NAMESPACE::get_registered_kernels;
 using ::executorch::ET_RUNTIME_NAMESPACE::Kernel;
 using ::executorch::ET_RUNTIME_NAMESPACE::Method;
+using ::executorch::ET_RUNTIME_NAMESPACE::MethodMeta;
 using ::executorch::ET_RUNTIME_NAMESPACE::Program;
 using ::executorch::extension::BufferDataLoader;
 using ::executorch::extension::MallocMemoryAllocator;
 using ::executorch::extension::MmapDataLoader;
 using ::executorch::extension::ET_BUNDLED_MODULE_NAMESPACE::BundledModule;
 using ::executorch::extension::pybindings::PyDataLoader;
 using ::executorch::runtime::DataLoader;
+using ::executorch::runtime::DeviceMemoryBuffer;
 using ::executorch::runtime::Error;
 using ::executorch::runtime::EValue;
 using ::executorch::runtime::EventTracerDebugLogLevel;
@@ -1077,18 +1081,33 @@ inline std::shared_ptr<ProgramState> load_program(
 /// A wrapper/util class for executorch memory allocations/manager.
 class ProgramMemory {
  public:
-  explicit ProgramMemory(std::vector<std::vector<uint8_t>>&& non_const_buffers)
+  /// `devices` is empty when every buffer is on the host, which keeps
+  /// `MemoryManager::has_device_memory()` false for CPU-only programs.
+  /// Otherwise it holds one entry per buffer, indexed like `sizes`.
+  ///
+  /// Members initialize in declaration order and each one reads the members
+  /// declared before it, so that order is load-bearing. Device buffers come
+  /// first so that a device that is missing or out of memory throws before the
+  /// host arenas are allocated and zero-filled, rather than after.
+  ProgramMemory(
+      std::vector<int64_t>&& sizes,
+      std::vector<runtime::etensor::Device>&& devices)
       : runtime_allocator_(),
-        non_const_buffers_(std::move(non_const_buffers)),
+        planned_sizes_(std::move(sizes)),
+        planned_devices_(std::move(devices)),
+        device_buffers_(allocate_device_buffers()),
+        non_const_buffers_(allocate_host_buffers()),
         non_const_spans_(create_non_const_spans()),
-        non_const_allocator_(
-            {non_const_spans_.data(), non_const_spans_.size()}),
+        non_const_allocator_(create_non_const_allocator()),
         mem_manager_(
             &const_allocator_,
             &non_const_allocator_,
             &runtime_allocator_,
             &temp_allocator_) {}
 
+  explicit ProgramMemory(std::vector<int64_t>&& sizes)
+      : ProgramMemory(std::move(sizes), {}) {}
+
   /// Returns a pointer to the internal memory manager, the Memory instance
   /// must outlive this pointer.
   MemoryManager* mem_manager() {
@@ -1105,6 +1124,16 @@ class ProgramMemory {
 
   MallocMemoryAllocator temp_allocator_{};
 
+  std::vector<int64_t> planned_sizes_;
+
+  std::vector<runtime::etensor::Device> planned_devices_;
+
+  // Backs device-tagged buffers; the entry is empty for a CPU-tagged buffer.
+  // Parallel to non_const_buffers_ so both index by planned buffer id. Empty
+  // for an all-host program.
+  std::vector<DeviceMemoryBuffer> device_buffers_;
+
+  // Backs CPU-tagged buffers; the entry is empty for a device-tagged buffer.
   std::vector<std::vector<uint8_t>> non_const_buffers_;
 
   std::vector<Span<uint8_t>> non_const_spans_;
@@ -1113,16 +1142,125 @@ class ProgramMemory {
 
   MemoryManager mem_manager_;
 
+  bool is_device_buffer(size_t index) const {
+    return index < planned_devices_.size() && !planned_devices_[index].is_cpu();
+  }
+
+  std::vector<std::vector<uint8_t>> allocate_host_buffers() {
+    std::vector<std::vector<uint8_t>> result;
+    result.reserve(planned_sizes_.size());
+    for (size_t i = 0; i < planned_sizes_.size(); ++i) {
+      if (is_device_buffer(i)) {
+        result.emplace_back();
+      } else {
+        result.emplace_back(planned_sizes_[i]);
+      }
+    }
+    return result;
+  }
+
+  std::vector<DeviceMemoryBuffer> allocate_device_buffers() {
+    std::vector<DeviceMemoryBuffer> result;
+    if (planned_devices_.empty()) {
+      return result;
+    }
+    // Both vectors are filled in lockstep today, so this only fires if a
+    // future caller breaks that. HierarchicalAllocator aborts on a mismatch,
+    // so check here instead, where a Python caller can catch it.
+    THROW_IF_ERROR(
+        planned_devices_.size() == planned_sizes_.size()
+            ? Error::Ok
+            : Error::InvalidArgument,
+        "Have %zu planned buffer sizes but %zu device tags",
+        planned_sizes_.size(),
+        planned_devices_.size());
+    result.reserve(planned_sizes_.size());
+    for (size_t i = 0; i < planned_sizes_.size(); ++i) {
+      if (!is_device_buffer(i)) {
+        result.emplace_back();
+        continue;
+      }
+      auto buffer = DeviceMemoryBuffer::create(
+          planned_sizes_[i],
+          planned_devices_[i].type(),
+          planned_devices_[i].index());
+      THROW_IF_ERROR(
+          buffer.error(),
+          "Failed to allocate %" PRId64 " bytes for buffer %zu on device %d:%d",
+          planned_sizes_[i],
+          i,
+          static_cast<int>(planned_devices_[i].type()),
+          static_cast<int>(planned_devices_[i].index()));
+      result.emplace_back(std::move(buffer.get()));
+    }
+    return result;
+  }
+
   std::vector<Span<uint8_t>> create_non_const_spans() {
     std::vector<Span<uint8_t>> result;
-    for (size_t i = 0; i < non_const_buffers_.size(); i++) {
-      result.push_back(
-          {non_const_buffers_[i].data(), non_const_buffers_[i].size()});
+    result.reserve(planned_sizes_.size());
+    for (size_t i = 0; i < planned_sizes_.size(); ++i) {
+      if (is_device_buffer(i)) {
+        result.push_back(device_buffers_[i].as_span());
+      } else {
+        result.push_back(
+            {non_const_buffers_[i].data(), non_const_buffers_[i].size()});
+      }
     }
     return result;
   }
+
+  HierarchicalAllocator create_non_const_allocator() {
+    Span<Span<uint8_t>> buffers(
+        non_const_spans_.data(), non_const_spans_.size());
+    return planned_devices_.empty()
+        ? HierarchicalAllocator(buffers)
+        : HierarchicalAllocator(
+              buffers, {planned_devices_.data(), planned_devices_.size()});
+  }
 };
 
+/// True if any of the method's memory-planned buffers must live off the host.
+bool has_device_buffers(const MethodMeta& method_meta) {
+  for (size_t i = 0; i < method_meta.num_memory_planned_buffers(); ++i) {
+    auto device = method_meta.memory_planned_buffer_device(i);
+    THROW_IF_ERROR(
+        device.error(), "Failed to get device of planned buffer %zu", i);
+    if (!device.get().is_cpu()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+/// Arenas sized and placed for a single method, used when that method's
+/// buffers cannot come from the program-wide host arenas. Returns nullptr when
+/// every buffer is on the host, so one pass over the metadata answers both
+/// whether the method needs its own arenas and how big they are.
+std::shared_ptr<ProgramMemory> make_method_memory(
+    const MethodMeta& method_meta) {
+  const size_t num_buffers = method_meta.num_memory_planned_buffers();
+  std::vector<int64_t> sizes;
+  std::vector<runtime::etensor::Device> devices;
+  sizes.reserve(num_buffers);
+  devices.reserve(num_buffers);
+  bool needs_device_memory = false;
+  for (size_t i = 0; i < num_buffers; ++i) {
+    auto size = method_meta.memory_planned_buffer_size(i);
+    THROW_IF_ERROR(size.error(), "Failed to get size of planned buffer %zu", i);
+    auto device = method_meta.memory_planned_buffer_device(i);
+    THROW_IF_ERROR(
+        device.error(), "Failed to get device of planned buffer %zu", i);
+    needs_device_memory |= !device.get().is_cpu();
+    sizes.push_back(size.get());
+    devices.push_back(device.get());
+  }
+  if (!needs_device_memory) {
+    return nullptr;
+  }
+  return std::make_shared<ProgramMemory>(std::move(sizes), std::move(devices));
+}
+
 struct PyMethod final {
   explicit PyMethod(
       std::shared_ptr<ProgramMemory> memory,
@@ -1423,8 +1561,17 @@ struct PyProgram final {
     for (size_t i = 0; i < state_->program_->num_methods(); ++i) {
       auto name = state_->program_->get_method_name(i).get();
       auto method_meta = state_->program_->method_meta(name).get();
-      for (size_t j = 0; j < method_meta.num_non_const_buffers(); j++) {
-        int64_t buffer_size = method_meta.non_const_buffer_size(j).get();
+      // A device-planned method gets its own arenas in load_method and never
+      // reads these, so letting its sizes in would only grow the host arenas
+      // the other methods share.
+      if (has_device_buffers(method_meta)) {
+        continue;
+      }
+      for (size_t j = 0; j < method_meta.num_memory_planned_buffers(); ++j) {
+        auto size = method_meta.memory_planned_buffer_size(j);
+        THROW_IF_ERROR(
+            size.error(), "Failed to get size of planned buffer %zu", j);
+        int64_t buffer_size = size.get();
         if (non_const_buffer_sizes.find(j) == non_const_buffer_sizes.end()) {
           non_const_buffer_sizes.insert({j, buffer_size});
         } else {
@@ -1434,16 +1581,14 @@ struct PyProgram final {
       }
     }
 
-    // Allocate the arenas. Using vector because we need to remember the size as
-    // well, so vector is easier then unique_ptr.
-    std::vector<std::vector<uint8_t>> non_const_buffers_;
-    for (std::map<size_t, int64_t>::iterator i = non_const_buffer_sizes.begin();
-         i != non_const_buffer_sizes.end();
-         i++) {
-      non_const_buffers_.push_back(std::vector<uint8_t>(i->second));
+    // Allocate the shared host arenas.
+    std::vector<int64_t> planned_sizes;
+    planned_sizes.reserve(non_const_buffer_sizes.size());
+    for (const auto& entry : non_const_buffer_sizes) {
+      planned_sizes.push_back(entry.second);
     }
 
-    memory_ = std::make_shared<ProgramMemory>(std::move(non_const_buffers_));
+    memory_ = std::make_shared<ProgramMemory>(std::move(planned_sizes));
     if (event_tracer_ && debug_buffer_size > 0) {
       // If a debug buffer was requested for the ETDump, allocate it and make
       // sure its lifetime is as long as the event_tracer.
@@ -1508,9 +1653,22 @@ struct PyProgram final {
   }
 
   std::unique_ptr<PyMethod> load_method(const std::string& method_name) {
+    Result<MethodMeta> meta =
+        state_->program_->method_meta(method_name.c_str());
+    THROW_IF_ERROR(
+        meta.error(),
+        "Failed to get method meta for method %s, error: 0x:%" PRIx32,
+        method_name.c_str(),
+        static_cast<uint32_t>(meta.error()));
+    // Device memory is claimed here rather than at program load so that one
+    // accelerator method cannot make the rest of the program unloadable. A
+    // host-only method keeps sharing the program-wide arenas, so its planned
+    // memory is not isolated from the other host-only methods of this program.
+    auto method_memory = make_method_memory(meta.get());
+    auto memory = method_memory ? std::move(method_memory) : memory_;
     Result<Method> res = state_->program_->load_method(
         method_name.c_str(),
-        memory_->mem_manager(),
+        memory->mem_manager(),
         event_tracer_.get(),
         state_->data_map_.get());
     THROW_IF_ERROR(
@@ -1519,7 +1677,9 @@ struct PyProgram final {
         method_name.c_str(),
         static_cast<uint32_t>(res.error()));
     return std::make_unique<PyMethod>(
-        memory_, state_, std::make_unique<Method>(std::move(res.get())));
+        std::move(memory),
+        state_,
+        std::make_unique<Method>(std::move(res.get())));
   }
 
   Span<uint8_t> get_etdump_debug_buffer() {

extension/pybindings/test/BUCK

@@ -26,6 +26,7 @@ fbcode_target(
         "//executorch/exir:pass_manager",
         "//executorch/exir:scalar_type",
         "//executorch/exir/_serialize:lib",
+        "//executorch/exir/backend:partitioner",
         "//executorch/exir/emit:lib",
         "//executorch/exir/passes:lib",
         "//executorch/runtime/core:core",
@@ -39,6 +40,8 @@ fbcode_target(
     preload_deps = ["//executorch/kernels/quantized:aot_lib"],
     deps = [
         ":make_test",
+        "//executorch/exir:schema",
+        "//executorch/exir/backend/test:device_util",
         "//executorch/extension/pybindings:portable_lib",
     ],
 )
@@ -50,6 +53,8 @@ fbcode_target(
     preload_deps = ["//executorch/kernels/quantized:aot_lib"],
     deps = [
         ":make_test",
+        "//executorch/exir:schema",
+        "//executorch/exir/backend/test:device_util",
         "//executorch/extension/pybindings:aten_lib",
         "//executorch/kernels/quantized:aot_lib",
     ],
@@ -61,6 +66,8 @@ fbcode_target(
     srcs = ["test_pybindings.py"],
     deps = [
         ":make_test",
+        "//executorch/exir:schema",
+        "//executorch/exir/backend/test:device_util",
     ],
 )