Return an error instead of aborting when device memory cannot be allocated

extension/module/module.cpp

@@ -368,7 +368,8 @@ Module::make_planned_memory_with_shared_arenas(
   return planned;
 }
 
-std::unique_ptr<Module::PlannedMemory> Module::make_planned_memory_with_devices(
+runtime::Result<std::unique_ptr<Module::PlannedMemory>>
+Module::make_planned_memory_with_devices(
     const ET_RUNTIME_NAMESPACE::MethodMeta& method_meta) {
   auto planned = std::make_unique<PlannedMemory>();
   const size_t num_buffers = method_meta.num_memory_planned_buffers();
@@ -379,9 +380,11 @@ std::unique_ptr<Module::PlannedMemory> Module::make_planned_memory_with_devices(
 
   for (size_t i = 0; i < num_buffers; ++i) {
     auto size = method_meta.memory_planned_buffer_size(i);
-    ET_CHECK_MSG(size.ok(), "Failed to get buffer size for index %zu", i);
+    ET_CHECK_OK_OR_RETURN_ERROR(
+        size.error(), "Failed to get buffer size for index %zu", i);
     auto device = method_meta.memory_planned_buffer_device(i);
-    ET_CHECK_MSG(device.ok(), "Failed to get buffer device for index %zu", i);
+    ET_CHECK_OK_OR_RETURN_ERROR(
+        device.error(), "Failed to get buffer device for index %zu", i);
     planned->planned_devices.push_back(device.get());
 
     if (device->is_cpu()) {
@@ -391,13 +394,17 @@ std::unique_ptr<Module::PlannedMemory> Module::make_planned_memory_with_devices(
     } else {
       // Allocate device memory via DeviceAllocator and store the RAII buffer.
       planned->planned_buffers.emplace_back(); // empty CPU placeholder
+      // Whether a device allocator exists, and whether the device has room,
+      // are properties of the machine rather than of the program, so report
+      // them to the caller instead of aborting the process.
       auto dmb = runtime::DeviceMemoryBuffer::create(
           size.get(), device->type(), device->index());
-      ET_CHECK_MSG(
-          dmb.ok(),
-          "Failed to allocate device memory for buffer %zu (device_type=%d)",
+      ET_CHECK_OK_OR_RETURN_ERROR(
+          dmb.error(),
+          "Failed to allocate device memory for buffer %zu (device_type=%d, device_index=%d)",
           i,
-          static_cast<int>(device->type()));
+          static_cast<int>(device->type()),
+          static_cast<int>(device->index()));
       planned->planned_spans.emplace_back(dmb->as_span());
       planned->device_buffers.push_back(std::move(dmb.get()));
     }
@@ -473,10 +480,16 @@ runtime::Error Module::load_method(
       ET_CHECK_OK_OR_RETURN_ERROR(meta_res.error());
       auto& meta = meta_res.get();
 
+      // A failed device query must not read as "this buffer is on the host",
+      // which would silently hand a backend host memory. The loop bounds i by
+      // num_memory_planned_buffers(), so the callee's range check cannot fire
+      // today; this keeps the failure handled if that ever stops holding.
       bool has_device_buffers = false;
       for (size_t i = 0; i < meta.num_memory_planned_buffers(); ++i) {
         auto dev = meta.memory_planned_buffer_device(i);
-        if (dev.ok() && !dev->is_cpu()) {
+        ET_CHECK_OK_OR_RETURN_ERROR(
+            dev.error(), "Failed to get buffer device for index %zu", i);
+        if (!dev->is_cpu()) {
           has_device_buffers = true;
           break;
         }
@@ -493,7 +506,9 @@ runtime::Error Module::load_method(
 
         // Device-aware path: allocate CPU and device buffers. The device
         // span is owned by the HierarchicalAllocator inside PlannedMemory.
-        method_holder.planned_memory = make_planned_memory_with_devices(meta);
+        auto planned_res = make_planned_memory_with_devices(meta);
+        ET_CHECK_OK_OR_RETURN_ERROR(planned_res.error());
+        method_holder.planned_memory = std::move(planned_res.get());
         planned_memory = method_holder.planned_memory->planned_memory.get();
       } else if (!share_memory_arenas_) {
         auto sizes_res = get_mem_planned_buffer_sizes(method_name);

extension/module/module.h

@@ -730,7 +730,8 @@ class Module {
   std::unique_ptr<PlannedMemory> make_planned_memory_with_shared_arenas(
       const std::vector<size_t>& buffer_sizes,
       std::vector<std::vector<uint8_t>>& shared_arenas);
-  std::unique_ptr<PlannedMemory> make_planned_memory_with_devices(
+  runtime::Result<std::unique_ptr<PlannedMemory>>
+  make_planned_memory_with_devices(
       const ET_RUNTIME_NAMESPACE::MethodMeta& method_meta);
   runtime::Result<std::vector<size_t>> get_mem_planned_buffer_sizes(
       const std::string& method_name);

extension/module/test/CMakeLists.txt

@@ -17,7 +17,7 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
 
 include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)
 
-set(_test_srcs module_test.cpp)
+set(_test_srcs module_test.cpp module_device_memory_test.cpp)
 
 add_custom_command(
   OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/ModuleAdd.pte"
@@ -26,12 +26,15 @@ add_custom_command(
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.ptd"
          "${CMAKE_CURRENT_BINARY_DIR}/ModuleSharedState.pte"
+         "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddWithDevice.pte"
   COMMAND ${PYTHON_EXECUTABLE} -m test.models.export_program --modules
           "ModuleAdd,ModuleSharedState" --outdir "${CMAKE_CURRENT_BINARY_DIR}"
   COMMAND
     ${PYTHON_EXECUTABLE} -m test.models.export_program --modules
     "ModuleAddMul,ModuleLinear" --external-constants --outdir
     "${CMAKE_CURRENT_BINARY_DIR}"
+  COMMAND ${PYTHON_EXECUTABLE} -m test.models.export_program_with_device_info
+          --outdir "${CMAKE_CURRENT_BINARY_DIR}"
   WORKING_DIRECTORY ${EXECUTORCH_ROOT}
 )
 
@@ -43,6 +46,7 @@ add_custom_target(
           "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
           "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.ptd"
           "${CMAKE_CURRENT_BINARY_DIR}/ModuleSharedState.pte"
+          "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddWithDevice.pte"
 )
 
 set(test_env
@@ -52,6 +56,7 @@ set(test_env
     "ET_MODULE_LINEAR_PROGRAM_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte"
     "ET_MODULE_LINEAR_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.ptd"
     "ET_MODULE_SHARED_STATE=${CMAKE_CURRENT_BINARY_DIR}/ModuleSharedState.pte"
+    "ET_MODULE_ADD_WITH_DEVICE_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleAddWithDevice.pte"
 )
 
 et_cxx_test(

extension/module/test/module_device_memory_test.cpp

@@ -30,6 +30,7 @@
 using executorch::extension::Module;
 using executorch::runtime::DeviceMemoryBuffer;
 using executorch::runtime::Error;
+using executorch::runtime::get_device_allocator;
 using executorch::runtime::register_device_allocator;
 using executorch::runtime::etensor::DeviceType;
 using executorch::runtime::testing::MockCudaAllocator;
@@ -40,10 +41,15 @@ class ModuleDeviceMemoryTest : public ::testing::Test {
  protected:
   static void SetUpTestSuite() {
     executorch::runtime::runtime_init();
-    register_device_allocator(&g_mock_cuda);
+    // The registry is a process-wide static, so a second registration for the
+    // same device type aborts. Repeat runs re-enter this function.
+    if (get_device_allocator(DeviceType::CUDA) == nullptr) {
+      register_device_allocator(&g_mock_cuda);
+    }
   }
 
   void SetUp() override {
+    g_mock_cuda.fail_allocations_ = false;
     g_mock_cuda.allocate_count_ = 0;
     g_mock_cuda.deallocate_count_ = 0;
     g_mock_cuda.last_allocate_size_ = 0;
@@ -146,6 +152,27 @@ TEST_F(ModuleDeviceMemoryTest, DeviceModelWithSharedArenasReturnsNotSupported) {
   EXPECT_EQ(err, Error::NotSupported);
 }
 
+TEST_F(ModuleDeviceMemoryTest, DeviceAllocationFailureIsReportedNotFatal) {
+  const char* path = std::getenv("ET_MODULE_ADD_WITH_DEVICE_PATH");
+  ASSERT_NE(path, nullptr) << "ET_MODULE_ADD_WITH_DEVICE_PATH not set";
+
+  // Stand in for a device that is out of memory. The caller gets an error and
+  // the process survives to handle it.
+  g_mock_cuda.fail_allocations_ = true;
+
+  Module module(path);
+  EXPECT_EQ(module.load_method("forward"), Error::MemoryAllocationFailed);
+  EXPECT_FALSE(module.is_method_loaded("forward"));
+  EXPECT_EQ(g_mock_cuda.allocate_count_, 0);
+  EXPECT_EQ(g_mock_cuda.deallocate_count_, 0);
+
+  // The failure must not leave the Module half-built: with the device back,
+  // the same call reaches the allocator again.
+  g_mock_cuda.fail_allocations_ = false;
+  (void)module.load_method("forward");
+  EXPECT_EQ(g_mock_cuda.allocate_count_, 1);
+}
+
 TEST_F(
     ModuleDeviceMemoryTest,
     LoadMethodAllocatesDeviceMemoryAndDeallocatesOnDestroy) {

runtime/core/test/mock_cuda_allocator.h

@@ -34,6 +34,9 @@ class MockCudaAllocator : public DeviceAllocator {
     // malloc returns memory aligned to alignof(max_align_t), which satisfies
     // kDefaultAlignment; the mock only exercises the default alignment.
     (void)alignment;
+    if (fail_allocations_) {
+      return Error::MemoryAllocationFailed;
+    }
     void* ptr = std::malloc(nbytes);
     if (!ptr) {
       return Error::MemoryAllocationFailed;
@@ -98,6 +101,7 @@ class MockCudaAllocator : public DeviceAllocator {
   }
 
   void reset() {
+    fail_allocations_ = false;
     allocate_count_ = 0;
     deallocate_count_ = 0;
     h2d_count_ = 0;
@@ -117,6 +121,10 @@ class MockCudaAllocator : public DeviceAllocator {
     last_d2h_index_ = -1;
   }
 
+  /// When set, allocate() reports MemoryAllocationFailed without allocating,
+  /// so callers can be tested against a device that is out of memory.
+  bool fail_allocations_ = false;
+
   // Allocation tracking
   int allocate_count_ = 0;
   int deallocate_count_ = 0;