From ab9467f5eca49628ba853b4fff5d1c756b02aaca Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 24 Aug 2026 08:54:12 -0700 Subject: [PATCH 1/4] Return an error instead of aborting when device memory cannot be allocated A model can ask for its memory-planned buffers to 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 have nothing to do with the model file: no allocator has been registered for that device type in this build, or the device is out of memory. Both were checked with ET_CHECK_MSG, which terminates the whole process. A Python caller got a SIGABRT and a core dump with no traceback and no chance to fall back to the CPU. These are properties of the machine, not of the program, so load_method now returns the error to the caller and the process survives. The same function also ignored the error from two MethodMeta lookups, and the loop that decides whether a model uses device buffers at all discarded a failed device query, which then read as "this buffer is on the CPU". That last path is not reachable today, because the loop bounds keep the index in range and that is the only case the query rejects, but silently treating a failed query as CPU would hand a backend host memory, so it now reports the error too. Test plan: - Added 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. - Verified the test catches the bug: with the fix reverted the test binary exits with signal 6 (abort) and prints no result. With the fix, all 6 tests in the suite pass. - module_device_memory_test.cpp was only registered in the internal build, so no open source job ran it. Added it to extension/module/test/CMakeLists.txt along with the model file it needs. The full extension_module_test binary now runs 57 tests and all pass. - clang-format 18.1.3, cmake-format 0.6.13 and cmake-lint all report no changes on the touched files. --- extension/module/module.cpp | 32 +++++++++++++------ extension/module/module.h | 3 +- extension/module/test/CMakeLists.txt | 7 +++- .../module/test/module_device_memory_test.cpp | 23 +++++++++++++ runtime/core/test/mock_cuda_allocator.h | 8 +++++ 5 files changed, 62 insertions(+), 11 deletions(-) diff --git a/extension/module/module.cpp b/extension/module/module.cpp index 11fea031603..90d688830e3 100644 --- a/extension/module/module.cpp +++ b/extension/module/module.cpp @@ -368,7 +368,8 @@ Module::make_planned_memory_with_shared_arenas( return planned; } -std::unique_ptr Module::make_planned_memory_with_devices( +runtime::Result> +Module::make_planned_memory_with_devices( const ET_RUNTIME_NAMESPACE::MethodMeta& method_meta) { auto planned = std::make_unique(); const size_t num_buffers = method_meta.num_memory_planned_buffers(); @@ -379,9 +380,11 @@ std::unique_ptr 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::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(device->type())); + static_cast(device->type()), + static_cast(device->index())); planned->planned_spans.emplace_back(dmb->as_span()); planned->device_buffers.push_back(std::move(dmb.get())); } @@ -473,10 +480,15 @@ 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 index is always + // in range here, so this only fires if the metadata itself is broken. 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 +505,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); diff --git a/extension/module/module.h b/extension/module/module.h index 91c7feaad9b..f30cfc39e90 100644 --- a/extension/module/module.h +++ b/extension/module/module.h @@ -730,7 +730,8 @@ class Module { std::unique_ptr make_planned_memory_with_shared_arenas( const std::vector& buffer_sizes, std::vector>& shared_arenas); - std::unique_ptr make_planned_memory_with_devices( + runtime::Result> + make_planned_memory_with_devices( const ET_RUNTIME_NAMESPACE::MethodMeta& method_meta); runtime::Result> get_mem_planned_buffer_sizes( const std::string& method_name); diff --git a/extension/module/test/CMakeLists.txt b/extension/module/test/CMakeLists.txt index 57cb24413ee..40d5fb285ba 100644 --- a/extension/module/test/CMakeLists.txt +++ b/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( diff --git a/extension/module/test/module_device_memory_test.cpp b/extension/module/test/module_device_memory_test.cpp index 0e5111177f1..779b2dbcc32 100644 --- a/extension/module/test/module_device_memory_test.cpp +++ b/extension/module/test/module_device_memory_test.cpp @@ -44,6 +44,7 @@ class ModuleDeviceMemoryTest : public ::testing::Test { } 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 +147,28 @@ 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, or one whose allocator was + // never registered. Either way 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) { diff --git a/runtime/core/test/mock_cuda_allocator.h b/runtime/core/test/mock_cuda_allocator.h index 238d819311f..b9639d1021e 100644 --- a/runtime/core/test/mock_cuda_allocator.h +++ b/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; From a9d024ca5d291d0e156d1acd23b3a7c0a47d0ada Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 24 Aug 2026 10:24:20 -0700 Subject: [PATCH 2/4] Register the mock CUDA allocator only once in the module device test The device allocator registry is a process wide static. Registering a second allocator for the same device type aborts the process on purpose. GoogleTest calls SetUpTestSuite again for every repeat iteration, so running this suite with --gtest_repeat=2 aborted instead of passing. Check the registry first and register only when nothing is there yet. This is what runtime/core/test/device_allocator_test.cpp and runtime/core/test/device_memory_buffer_test.cpp already do. Test plan: ctest -R '^extension_module_test$' 1/1 Test #54: extension_module_test ... Passed extension_module_test --gtest_repeat=3 [ PASSED ] 57 tests. three times, exit 0 With the check removed again, the same repeat command aborts with exit 134, so the test really covers this. --- extension/module/test/module_device_memory_test.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/extension/module/test/module_device_memory_test.cpp b/extension/module/test/module_device_memory_test.cpp index 779b2dbcc32..188fe69a28c 100644 --- a/extension/module/test/module_device_memory_test.cpp +++ b/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,7 +41,11 @@ 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 { From 2adb061d9e4211526b1b3791eb82bf156d1933d3 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 24 Aug 2026 10:24:20 -0700 Subject: [PATCH 3/4] Correct a wrong comment in the module device memory test The comment claimed the forced allocator failure also stands in for a device whose allocator was never registered. It does not. Two different errors are possible here. A missing allocator makes DeviceMemoryBuffer::create return Error::NotFound, while a registered allocator that refuses the request returns Error::MemoryAllocationFailed. This test only reaches the second one. Comment only, no behavior change. Test plan: ctest -R '^extension_module_test$' 1/1 Test #54: extension_module_test ... Passed --- extension/module/test/module_device_memory_test.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/extension/module/test/module_device_memory_test.cpp b/extension/module/test/module_device_memory_test.cpp index 188fe69a28c..3f1a5604ecd 100644 --- a/extension/module/test/module_device_memory_test.cpp +++ b/extension/module/test/module_device_memory_test.cpp @@ -156,9 +156,8 @@ 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, or one whose allocator was - // never registered. Either way the caller gets an error and the process - // survives to handle it. + // 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); From 820c4e2e9aede5177bc9d544990b265ce834c134 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 24 Aug 2026 14:37:35 -0700 Subject: [PATCH 4/4] Fix comment: the range check is unreachable, not a broken-metadata signal The callee's only error return is a range check that this loop's bound makes unreachable, and broken metadata yields a CPU device rather than an error. Describe what the handling is actually for. --- extension/module/module.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/extension/module/module.cpp b/extension/module/module.cpp index 90d688830e3..a9cf8578131 100644 --- a/extension/module/module.cpp +++ b/extension/module/module.cpp @@ -481,8 +481,9 @@ runtime::Error Module::load_method( 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 index is always - // in range here, so this only fires if the metadata itself is broken. + // 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);