diff --git a/include/rmm/detail/error.hpp b/include/rmm/detail/error.hpp index 1f550f75e..89cbcd916 100644 --- a/include/rmm/detail/error.hpp +++ b/include/rmm/detail/error.hpp @@ -60,6 +60,16 @@ class bad_alloc : public std::bad_alloc { std::string _what; }; +/** + * @brief Exception thrown when RMM runs out of memory + * + * This error should only be thrown when we know for sure a resource is out of memory. + */ +class out_of_memory : public bad_alloc { + public: + using bad_alloc::bad_alloc; +}; + /** * @brief Exception thrown when attempting to access outside of a defined range * @@ -167,6 +177,31 @@ class out_of_range : public std::out_of_range { } while (0) #define RMM_CUDA_TRY_1(_call) RMM_CUDA_TRY_2(_call, rmm::cuda_error) +/** + * @brief Error checking macro for CUDA memory allocation calls. + * + * Invokes a CUDA memory allocation function call. If the call does not return + * `cudaSuccess`, invokes cudaGetLastError() to clear the error and throws an + * exception detailing the CUDA error that occurred + * + * Defaults to throwing `rmm::bad_alloc`, but when `cudaErrorMemoryAllocation` is returned, + * `rmm::out_of_memory` is thrown instead. + */ +#define RMM_CUDA_TRY_ALLOC(_call) \ + do { \ + cudaError_t const error = (_call); \ + if (cudaSuccess != error) { \ + cudaGetLastError(); \ + auto const msg = std::string{"CUDA error at: "} + __FILE__ + ":" + RMM_STRINGIFY(__LINE__) + \ + ": " + cudaGetErrorName(error) + " " + cudaGetErrorString(error); \ + if (cudaErrorMemoryAllocation == error) { \ + throw rmm::out_of_memory{msg}; \ + } else { \ + throw rmm::bad_alloc{msg}; \ + } \ + } \ + } while (0) + /** * @brief Error checking macro similar to `assert` for CUDA runtime API calls * diff --git a/include/rmm/mr/device/arena_memory_resource.hpp b/include/rmm/mr/device/arena_memory_resource.hpp index ce8737225..c9954a507 100644 --- a/include/rmm/mr/device/arena_memory_resource.hpp +++ b/include/rmm/mr/device/arena_memory_resource.hpp @@ -154,7 +154,7 @@ class arena_memory_resource final : public device_memory_resource { pointer = arena.allocate(bytes); if (pointer == nullptr) { if (dump_log_on_failure_) { dump_memory_log(bytes); } - RMM_FAIL("Maximum pool size exceeded", rmm::bad_alloc); + RMM_FAIL("Maximum pool size exceeded", rmm::out_of_memory); } } diff --git a/include/rmm/mr/device/cuda_async_memory_resource.hpp b/include/rmm/mr/device/cuda_async_memory_resource.hpp index a2d8c8567..949a4d940 100644 --- a/include/rmm/mr/device/cuda_async_memory_resource.hpp +++ b/include/rmm/mr/device/cuda_async_memory_resource.hpp @@ -148,8 +148,7 @@ class cuda_async_memory_resource final : public device_memory_resource { void* ptr{nullptr}; #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT if (bytes > 0) { - RMM_CUDA_TRY(cudaMallocFromPoolAsync(&ptr, bytes, pool_handle(), stream.value()), - rmm::bad_alloc); + RMM_CUDA_TRY_ALLOC(cudaMallocFromPoolAsync(&ptr, bytes, pool_handle(), stream.value())); } #else (void)bytes; diff --git a/include/rmm/mr/device/cuda_memory_resource.hpp b/include/rmm/mr/device/cuda_memory_resource.hpp index b5b3d87df..800d0b62c 100644 --- a/include/rmm/mr/device/cuda_memory_resource.hpp +++ b/include/rmm/mr/device/cuda_memory_resource.hpp @@ -67,7 +67,7 @@ class cuda_memory_resource final : public device_memory_resource { void* do_allocate(std::size_t bytes, cuda_stream_view) override { void* ptr{nullptr}; - RMM_CUDA_TRY(cudaMalloc(&ptr, bytes), rmm::bad_alloc); + RMM_CUDA_TRY_ALLOC(cudaMalloc(&ptr, bytes)); return ptr; } diff --git a/include/rmm/mr/device/detail/stream_ordered_memory_resource.hpp b/include/rmm/mr/device/detail/stream_ordered_memory_resource.hpp index 2a726377d..43da17d0f 100644 --- a/include/rmm/mr/device/detail/stream_ordered_memory_resource.hpp +++ b/include/rmm/mr/device/detail/stream_ordered_memory_resource.hpp @@ -208,7 +208,7 @@ class stream_ordered_memory_resource : public crtp, public device_ size = rmm::detail::align_up(size, rmm::detail::CUDA_ALLOCATION_ALIGNMENT); RMM_EXPECTS(size <= this->underlying().get_maximum_allocation_size(), - rmm::bad_alloc, + rmm::out_of_memory, "Maximum allocation size exceeded"); auto const block = this->underlying().get_block(size, stream_event); diff --git a/include/rmm/mr/device/limiting_resource_adaptor.hpp b/include/rmm/mr/device/limiting_resource_adaptor.hpp index 05b45ce31..cdcee0857 100644 --- a/include/rmm/mr/device/limiting_resource_adaptor.hpp +++ b/include/rmm/mr/device/limiting_resource_adaptor.hpp @@ -140,7 +140,7 @@ class limiting_resource_adaptor final : public device_memory_resource { } allocated_bytes_ -= proposed_size; - RMM_FAIL("Exceeded memory limit", rmm::bad_alloc); + RMM_FAIL("Exceeded memory limit", rmm::out_of_memory); } /** diff --git a/include/rmm/mr/device/managed_memory_resource.hpp b/include/rmm/mr/device/managed_memory_resource.hpp index a2be418c7..1bc917e0e 100644 --- a/include/rmm/mr/device/managed_memory_resource.hpp +++ b/include/rmm/mr/device/managed_memory_resource.hpp @@ -71,7 +71,7 @@ class managed_memory_resource final : public device_memory_resource { if (bytes == 0) { return nullptr; } void* ptr{nullptr}; - RMM_CUDA_TRY(cudaMallocManaged(&ptr, bytes), rmm::bad_alloc); + RMM_CUDA_TRY_ALLOC(cudaMallocManaged(&ptr, bytes)); return ptr; } diff --git a/include/rmm/mr/device/pool_memory_resource.hpp b/include/rmm/mr/device/pool_memory_resource.hpp index 526852355..5a4d6fc1e 100644 --- a/include/rmm/mr/device/pool_memory_resource.hpp +++ b/include/rmm/mr/device/pool_memory_resource.hpp @@ -180,7 +180,7 @@ class pool_memory_resource final RMM_LOG_ERROR("[A][Stream {}][Upstream {}B][FAILURE maximum pool size exceeded]", fmt::ptr(stream.value()), min_size); - RMM_FAIL("Maximum pool size exceeded", rmm::bad_alloc); + RMM_FAIL("Maximum pool size exceeded", rmm::out_of_memory); } /** diff --git a/tests/mr/device/limiting_mr_tests.cpp b/tests/mr/device/limiting_mr_tests.cpp index 3bc643abc..80c49104f 100644 --- a/tests/mr/device/limiting_mr_tests.cpp +++ b/tests/mr/device/limiting_mr_tests.cpp @@ -38,7 +38,7 @@ TEST(LimitingTest, TooBig) { auto const max_size{5_MiB}; Limiting_adaptor mr{rmm::mr::get_current_device_resource(), max_size}; - EXPECT_THROW(mr.allocate(max_size + 1), rmm::bad_alloc); + EXPECT_THROW(mr.allocate(max_size + 1), rmm::out_of_memory); } TEST(LimitingTest, UnderLimitDueToFrees) @@ -83,7 +83,7 @@ TEST(LimitingTest, OverLimit) EXPECT_EQ(mr.get_allocated_bytes(), allocated_bytes); EXPECT_EQ(mr.get_allocation_limit() - mr.get_allocated_bytes(), max_size - allocated_bytes); auto const size2{3_MiB}; - EXPECT_THROW(mr.allocate(size2), rmm::bad_alloc); + EXPECT_THROW(mr.allocate(size2), rmm::out_of_memory); EXPECT_EQ(mr.get_allocated_bytes(), allocated_bytes); EXPECT_EQ(mr.get_allocation_limit() - mr.get_allocated_bytes(), max_size - allocated_bytes); mr.deallocate(ptr1, 4_MiB); diff --git a/tests/mr/device/mr_multithreaded_tests.cpp b/tests/mr/device/mr_multithreaded_tests.cpp index 838035d9f..38c34d93f 100644 --- a/tests/mr/device/mr_multithreaded_tests.cpp +++ b/tests/mr/device/mr_multithreaded_tests.cpp @@ -37,6 +37,9 @@ struct mr_test_mt : public mr_test { INSTANTIATE_TEST_CASE_P(MultiThreadResourceTests, mr_test_mt, ::testing::Values(mr_factory{"CUDA", &make_cuda}, +#ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT + mr_factory{"CUDA_Async", &make_cuda_async}, +#endif mr_factory{"Managed", &make_managed}, mr_factory{"Pool", &make_pool}, mr_factory{"Arena", &make_arena}, diff --git a/tests/mr/device/mr_test.hpp b/tests/mr/device/mr_test.hpp index 4bef2b54e..7635901c8 100644 --- a/tests/mr/device/mr_test.hpp +++ b/tests/mr/device/mr_test.hpp @@ -128,7 +128,7 @@ inline void test_various_allocations(rmm::mr::device_memory_resource* mr, cuda_s // should fail to allocate too much { void* ptr{nullptr}; - EXPECT_THROW(ptr = mr->allocate(1_PiB, stream), rmm::bad_alloc); + EXPECT_THROW(ptr = mr->allocate(1_PiB, stream), rmm::out_of_memory); EXPECT_EQ(nullptr, ptr); } } diff --git a/tests/mr/device/pool_mr_tests.cpp b/tests/mr/device/pool_mr_tests.cpp index 9f2020785..ca8888a52 100644 --- a/tests/mr/device/pool_mr_tests.cpp +++ b/tests/mr/device/pool_mr_tests.cpp @@ -86,7 +86,7 @@ TEST(PoolTest, ForceGrowth) EXPECT_NO_THROW(mr.allocate(1000)); EXPECT_NO_THROW(mr.allocate(4000)); EXPECT_NO_THROW(mr.allocate(500)); - EXPECT_THROW(mr.allocate(2000), rmm::bad_alloc); // too much + EXPECT_THROW(mr.allocate(2000), rmm::out_of_memory); // too much } TEST(PoolTest, DeletedStream) diff --git a/tests/mr/device/thrust_allocator_tests.cu b/tests/mr/device/thrust_allocator_tests.cu index eabdfe143..c80f3d9e2 100644 --- a/tests/mr/device/thrust_allocator_tests.cu +++ b/tests/mr/device/thrust_allocator_tests.cu @@ -37,8 +37,12 @@ TEST_P(allocator_test, first) INSTANTIATE_TEST_CASE_P(ThrustAllocatorTests, allocator_test, ::testing::Values(mr_factory{"CUDA", &make_cuda}, +#ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT + mr_factory{"CUDA_Async", &make_cuda_async}, +#endif mr_factory{"Managed", &make_managed}, mr_factory{"Pool", &make_pool}, + mr_factory{"Arena", &make_arena}, mr_factory{"Binning", &make_binning}), [](auto const& info) { return info.param.name; });