Free device allocation when copy fails

include/nvexec/detail/memory.cuh

@@ -60,20 +60,21 @@ namespace nv::execution::_strm
       if (status = STDEXEC_LOG_CUDA_API(::cudaMalloc(reinterpret_cast<void**>(&ptr), sizeof(Type)));
           status == cudaSuccess)
       {
+        device_ptr_t<Type> result{ptr};
+
         STDEXEC_TRY
         {
           Type h(static_cast<Args&&>(args)...);
           status = STDEXEC_LOG_CUDA_API(
             ::cudaMemcpy(ptr, &h, sizeof(Type), cudaMemcpyHostToDevice));
           if (status == cudaSuccess)
           {
-            return device_ptr_t<Type>(ptr);
+            return result;
           }
         }
         STDEXEC_CATCH_ALL
         {
           status = cudaErrorUnknown;
-          STDEXEC_ASSERT_CUDA_API(::cudaFree(ptr));
         }
       }
     }

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

test/nvexec/device_allocate.cpp

@@ -0,0 +1,43 @@
+#include <cstddef>
+#include <cuda_runtime_api.h>
+
+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 <test_common/catch2.hpp>
+
+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<int>(status, 42);
+
+      REQUIRE(status == cudaErrorInvalidValue);
+      REQUIRE(ptr == nullptr);
+    }
+
+    REQUIRE(test_device_allocate_free_calls == 1);
+  }
+}  // namespace