Merge pull request #2190 from fallintoplace/fix/nvexec-continues-on-error-cleanup

Destroy nvexec continues_on storage on CUDA errors

Author: ericniebler

include/nvexec/stream/continues_on.cuh

@@ -83,11 +83,20 @@ namespace nv::execution::_strm
             storage->template emplace<tuple_t>(Tag(), static_cast<Args&&>(args)...);
           }
 
+          auto complete_error = [storage, &opstate = opstate_](cudaError_t status) noexcept
+          {
+            if constexpr (!construct_on_device)
+            {
+              storage->~storage_t();
+            }
+            opstate.propagate_completion_signal(STDEXEC::set_error, std::move(status));
+          };
+
           int dev_id{};
           if (cudaError_t status = STDEXEC_LOG_CUDA_API(cudaGetDevice(&dev_id));
               status != cudaSuccess)
           {
-            opstate_.propagate_completion_signal(STDEXEC::set_error, std::move(status));
+            complete_error(std::move(status));
             return;
           }
 
@@ -98,7 +107,7 @@ namespace nv::execution::_strm
                                        dev_id));
               status != cudaSuccess)
           {
-            opstate_.propagate_completion_signal(STDEXEC::set_error, std::move(status));
+            complete_error(std::move(status));
             return;
           }
 
@@ -110,7 +119,7 @@ namespace nv::execution::_strm
                   cudaMemPrefetchAsync(storage, sizeof(storage_t), dev_id, stream));
                 status != cudaSuccess)
             {
-              opstate_.propagate_completion_signal(STDEXEC::set_error, std::move(status));
+              complete_error(std::move(status));
               return;
             }
           }
@@ -123,7 +132,7 @@ namespace nv::execution::_strm
             if (cudaError_t status = STDEXEC_LOG_CUDA_API(cudaPeekAtLastError());
                 status != cudaSuccess)
             {
-              opstate_.propagate_completion_signal(STDEXEC::set_error, std::move(status));
+              complete_error(std::move(status));
               return;
             }
           }

test/nvexec/continues_on.cpp

@@ -1,10 +1,64 @@
 #include <stdexec/execution.hpp>
 #include <test_common/catch2.hpp>
 
+#include "common.cuh"
 #include "nvexec/stream_context.cuh"
 
+#include <memory_resource>
+
 namespace
 {
+  class pinned_memory_resource_t : public std::pmr::memory_resource
+  {
+    void* do_allocate(std::size_t bytes, std::size_t) override
+    {
+      void* storage{};
+      STDEXEC_TRY_CUDA_API(cudaMallocHost(&storage, bytes));
+      return storage;
+    }
+
+    void do_deallocate(void* storage, std::size_t, std::size_t) override
+    {
+      STDEXEC_ASSERT_CUDA_API(cudaFreeHost(storage));
+    }
+
+    auto do_is_equal(std::pmr::memory_resource const & other) const noexcept -> bool override
+    {
+      return this == &other;
+    }
+  };
+
+  class destruction_probe_t
+  {
+    flags_storage_t<>::flags_t flags_;
+    bool                       owns_{true};
+
+   public:
+    destruction_probe_t()                                               = delete;
+    destruction_probe_t(destruction_probe_t const &)                    = delete;
+    auto operator=(destruction_probe_t const &) -> destruction_probe_t& = delete;
+    auto operator=(destruction_probe_t&&) -> destruction_probe_t&       = delete;
+
+    __host__ __device__ explicit destruction_probe_t(flags_storage_t<>::flags_t flags)
+      : flags_(flags)
+    {}
+
+    __host__ __device__ destruction_probe_t(destruction_probe_t&& other)
+      : flags_(other.flags_)
+      , owns_(other.owns_)
+    {
+      other.owns_ = false;
+    }
+
+    __host__ __device__ ~destruction_probe_t()
+    {
+      if (owns_)
+      {
+        flags_.set();
+      }
+    }
+  };
+
   TEST_CASE("continues on after just", "[cuda][stream][adaptors][continues_on]")
   {
     nvexec::stream_context ctx;
@@ -43,4 +97,32 @@ namespace
 
     REQUIRE(result.has_value());
   }
+
+  TEST_CASE("continues_on destroys host-constructed storage after a CUDA error",
+            "[cuda][stream][adaptors][continues_on]")
+  {
+    int device{};
+    STDEXEC_TRY_CUDA_API(cudaGetDevice(&device));
+
+    int concurrent_managed_access{};
+    STDEXEC_TRY_CUDA_API(cudaDeviceGetAttribute(&concurrent_managed_access,
+                                                cudaDevAttrConcurrentManagedAccess,
+                                                device));
+    if (!concurrent_managed_access)
+    {
+      SKIP("device does not support concurrent managed access");
+    }
+
+    pinned_memory_resource_t pinned_memory;
+    nvexec::stream_context   ctx;
+    auto                     scheduler = ctx.get_scheduler();
+    scheduler.ctx_.managed_resource_   = &pinned_memory;
+
+    flags_storage_t<> destructions{};
+    auto              sndr = STDEXEC::just(destruction_probe_t{destructions.get()})
+              | STDEXEC::continues_on(scheduler);
+
+    REQUIRE_THROWS(STDEXEC::sync_wait(std::move(sndr)));
+    REQUIRE(destructions.all_set_once());
+  }
 }  // namespace