diff --git a/include/nvexec/stream/continues_on.cuh b/include/nvexec/stream/continues_on.cuh index 57466b896..657f49821 100644 --- a/include/nvexec/stream/continues_on.cuh +++ b/include/nvexec/stream/continues_on.cuh @@ -83,11 +83,20 @@ namespace nv::execution::_strm storage->template emplace(Tag(), static_cast(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; } } diff --git a/test/nvexec/continues_on.cpp b/test/nvexec/continues_on.cpp index a05dd0a81..5826c8155 100644 --- a/test/nvexec/continues_on.cpp +++ b/test/nvexec/continues_on.cpp @@ -1,10 +1,64 @@ #include #include +#include "common.cuh" #include "nvexec/stream_context.cuh" +#include + 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