diff --git a/include/nvexec/detail/memory.cuh b/include/nvexec/detail/memory.cuh index b0ab1a768..498332a43 100644 --- a/include/nvexec/detail/memory.cuh +++ b/include/nvexec/detail/memory.cuh @@ -60,6 +60,8 @@ namespace nv::execution::_strm if (status = STDEXEC_LOG_CUDA_API(::cudaMalloc(reinterpret_cast(&ptr), sizeof(Type))); status == cudaSuccess) { + device_ptr_t result{ptr}; + STDEXEC_TRY { Type h(static_cast(args)...); @@ -67,13 +69,12 @@ namespace nv::execution::_strm ::cudaMemcpy(ptr, &h, sizeof(Type), cudaMemcpyHostToDevice)); if (status == cudaSuccess) { - return device_ptr_t(ptr); + return result; } } STDEXEC_CATCH_ALL { status = cudaErrorUnknown; - STDEXEC_ASSERT_CUDA_API(::cudaFree(ptr)); } } } diff --git a/test/nvexec/CMakeLists.txt b/test/nvexec/CMakeLists.txt index bc4a819b6..d46a551e5 100644 --- a/test/nvexec/CMakeLists.txt +++ b/test/nvexec/CMakeLists.txt @@ -16,6 +16,7 @@ set(nvexec_test_sources continues_on.cpp + device_allocate.cpp bulk.cpp ensure_started.cpp start_detached.cpp diff --git a/test/nvexec/device_allocate.cpp b/test/nvexec/device_allocate.cpp new file mode 100644 index 000000000..6bfb1df54 --- /dev/null +++ b/test/nvexec/device_allocate.cpp @@ -0,0 +1,43 @@ +#include +#include + +static int test_device_allocate_free_calls{}; + +static cudaError_t +test_device_allocate_cudaMemcpy(void*, void const *, std::size_t, cudaMemcpyKind) noexcept +{ + return cudaErrorInvalidValue; +} + +static cudaError_t test_device_allocate_cudaFree(void* ptr) noexcept +{ + ++test_device_allocate_free_calls; + return ::cudaFree(ptr); +} + +#define cudaMemcpy test_device_allocate_cudaMemcpy +#define cudaFree test_device_allocate_cudaFree +#include "nvexec/detail/memory.cuh" +#undef cudaFree +#undef cudaMemcpy + +#include + +namespace +{ + TEST_CASE("device allocation frees storage when cudaMemcpy fails", + "[cuda][stream][memory][device_allocate]") + { + test_device_allocate_free_calls = 0; + cudaError_t status = cudaSuccess; + + { + auto ptr = nvexec::_strm::device_allocate(status, 42); + + REQUIRE(status == cudaErrorInvalidValue); + REQUIRE(ptr == nullptr); + } + + REQUIRE(test_device_allocate_free_calls == 1); + } +} // namespace