Return an error instead of aborting when device memory cannot be allocated (pytorch#22085)

### What is going on

A model can be exported so that some of its memory-planned buffers live
on an
accelerator instead of on the CPU. When `Module` loads such a model, it
asks the
runtime for a block of memory on that device.

That request can fail for two ordinary reasons that say nothing about
the model
file being loaded:

1. this build has no allocator registered for that device type, or
2. the device is out of memory.

Both were checked with `ET_CHECK_MSG`, which terminates the whole
process. So
loading a device-annotated model on a machine without the matching
backend, or
on a busy GPU, killed the process outright. Through the Python bindings
(`_load_for_executorch_from_buffer` and friends build one of these
`Module`
objects) that shows up as an abort and a core dump: no traceback, no
error to
catch, no chance to fall back to the CPU.

Neither condition means the program is broken. They describe the machine
it is
running on, so `load_method` now returns the error to the caller and the
process stays alive.

Before this change, both reasons ended the same way:

```
module.load_method("forward");   // never returns, process aborts
```

After, the two reasons stay apart, because `DeviceMemoryBuffer::create`
already
reports them with different codes:

```
// nothing registered for that device type
module.load_method("forward");   // returns Error::NotFound

// allocator is there, the device has no room
module.load_method("forward");   // returns Error::MemoryAllocationFailed
```

### Also in this change

Two other `MethodMeta` lookups in the same function, one for the buffer
size
and one for the buffer device, also used `ET_CHECK_MSG`. They return the
error
now as well, so all three abort sites in that function are gone.

Separately, the loop that decides whether a model uses device buffers at
all
threw away a failed device query, which then read as "this buffer is on
the
CPU". It returns the error now instead of ignoring it.

That last path is not reachable today. The only case the query rejects
is an
out-of-range buffer index, and the loop bounds already keep the index in
range.
It is fixed for consistency, not because it can fire. Quietly treating a
failed
query as CPU would hand a backend host memory, which is a far harder
failure to
debug than an error return.

### Test coverage

`extension/module/test/module_device_memory_test.cpp` was registered in
the
internal build only, so no open source job ran any of it. This change
adds it
to `extension/module/test/CMakeLists.txt`, along with the model file it
needs,
so it now builds and runs as part of `extension_module_test`.

### Test plan

* New test
`ModuleDeviceMemoryTest.DeviceAllocationFailureIsReportedNotFatal`.
  It makes the test allocator refuse the request, then checks that
`load_method` returns `MemoryAllocationFailed`, that the method is not
left
half loaded, and that a later attempt with the device healthy allocates
  normally.
* Confirmed the test actually catches the bug. With this fix reverted,
the test
binary exits with signal 6 (abort) and prints no result at all. With the
fix,
  all 6 tests in the suite pass.
* Ran the whole `extension_module_test` binary: 57 tests, all pass. The
51
pre-existing tests are unaffected by adding the device test file to the
same
  binary.
* Ran `extension_module_test --gtest_repeat=3`: 57 tests pass on every
iteration, exit 0. The suite registers its allocator only once, so
process
reuse is safe. Put the unconditional registration back and the same
command
  aborts with exit 134.
* `clang-format` reports no changes needed on the touched source file.
  `cmake-format` 0.6.13 and `cmake-lint` report no changes needed on the
  touched `CMakeLists.txt`.

### Not covered

Nothing here runs against real accelerator hardware. The failure is
driven
through the existing mock allocator, which stands in for a device that
is out
of memory. The committed test therefore covers the
`Error::MemoryAllocationFailed` branch only. The `Error::NotFound`
branch, for
a device type with no registered allocator, is not exercised by a
committed
test.

Author: shoumikhin

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;