diff --git a/extension/module/module.cpp b/extension/module/module.cpp index 11fea031603..a9cf8578131 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,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); 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..3f1a5604ecd 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,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) { 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;