From 145620366ce6d4c8072f02254fc8eb710407eaf0 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Fri, 21 Aug 2026 10:33:11 +0200 Subject: [PATCH 1/2] Fix schedule_from task cleanup before start --- include/nvexec/stream/schedule_from.cuh | 42 ++++++++++++-------- test/nvexec/continues_on.cpp | 52 +++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 16 deletions(-) diff --git a/include/nvexec/stream/schedule_from.cuh b/include/nvexec/stream/schedule_from.cuh index 8f7edee11..4d45d6734 100644 --- a/include/nvexec/stream/schedule_from.cuh +++ b/include/nvexec/stream/schedule_from.cuh @@ -19,10 +19,9 @@ #pragma once #include "../../stdexec/execution.hpp" +#include #include -#include "../detail/cuda_atomic.cuh" // IWYU pragma: keep - #include "common.cuh" namespace nv::execution @@ -71,21 +70,34 @@ namespace nv::execution opstate& opstate_; }; + struct task_deleter + { + void operator()(task_t* task) const noexcept + { + if (task != nullptr) + { + task->free_(task); + } + } + }; + + using task_ptr_t = std::unique_ptr; + opstate(Sender&& sndr, Receiver&& rcvr, context ctx) : _strm::opstate_base(static_cast(rcvr), ctx) , ctx_(ctx) , storage_(host_allocate(this->status_, ctx.pinned_resource_)) - , task_(host_allocate(this->status_, - ctx.pinned_resource_, - receiver{*this}, - storage_.get(), - this->get_stream(), - ctx.pinned_resource_) - .release()) + , task_(task_ptr_t{host_allocate(this->status_, + ctx.pinned_resource_, + receiver{*this}, + storage_.get(), + this->get_stream(), + ctx.pinned_resource_) + .release()}) , env_(host_allocate(this->status_, ctx_.pinned_resource_, this->make_env())) - , inner_op_{ - connect(static_cast(sndr), - enqueue_receiver_t{env_.get(), storage_.get(), task_, ctx_.hub_->producer()})} + , inner_op_{connect( + static_cast(sndr), + enqueue_receiver_t{env_.get(), storage_.get(), task_.get(), ctx_.hub_->producer()})} { if (this->status_ == cudaSuccess) { @@ -97,8 +109,6 @@ namespace nv::execution void start() & noexcept { - started_.test_and_set(::cuda::std::memory_order::relaxed); - if (status_ != cudaSuccess) { // Couldn't allocate memory for operation state, complete with error @@ -106,14 +116,14 @@ namespace nv::execution return; } + task_.release(); STDEXEC::start(inner_op_); } cudaError_t status_{cudaSuccess}; context ctx_; host_ptr_t storage_; - task_t* task_; - ::cuda::std::atomic_flag started_{}; + task_ptr_t task_; host_ptr_t<__decay_t> env_{}; inner_opstate_t inner_op_; }; diff --git a/test/nvexec/continues_on.cpp b/test/nvexec/continues_on.cpp index 5826c8155..29e250c17 100644 --- a/test/nvexec/continues_on.cpp +++ b/test/nvexec/continues_on.cpp @@ -10,22 +10,56 @@ namespace { class pinned_memory_resource_t : public std::pmr::memory_resource { + std::size_t allocations_{}; + std::size_t deallocations_{}; + void* do_allocate(std::size_t bytes, std::size_t) override { void* storage{}; STDEXEC_TRY_CUDA_API(cudaMallocHost(&storage, bytes)); + ++allocations_; return storage; } void do_deallocate(void* storage, std::size_t, std::size_t) override { STDEXEC_ASSERT_CUDA_API(cudaFreeHost(storage)); + ++deallocations_; } auto do_is_equal(std::pmr::memory_resource const & other) const noexcept -> bool override { return this == &other; } + + public: + auto allocations() const noexcept -> std::size_t + { + return allocations_; + } + + auto deallocations() const noexcept -> std::size_t + { + return deallocations_; + } + }; + + struct noop_receiver + { + using receiver_concept = STDEXEC::receiver_tag; + + auto get_env() const noexcept -> STDEXEC::env<> + { + return {}; + } + + void set_value() noexcept {} + + template + void set_error(Error&&) noexcept + {} + + void set_stopped() noexcept {} }; class destruction_probe_t @@ -68,6 +102,24 @@ namespace STDEXEC::sync_wait(std::move(sndr)); } + TEST_CASE("continues_on frees its task when the operation is not started", + "[cuda][stream][adaptors][continues_on]") + { + pinned_memory_resource_t pinned_memory; + nvexec::stream_context ctx; + auto scheduler = ctx.get_scheduler(); + scheduler.ctx_.pinned_resource_ = &pinned_memory; + + auto sndr = STDEXEC::just() | STDEXEC::continues_on(scheduler); + { + auto op = STDEXEC::connect(std::move(sndr), noop_receiver{}); + (void) op; + } + + REQUIRE(pinned_memory.allocations() > 0); + REQUIRE(pinned_memory.allocations() == pinned_memory.deallocations()); + } + TEST_CASE("continues on after schedule", "[cuda][stream][adaptors][continues_on]") { nvexec::stream_context ctx; From 3de8db6f1e1d17c2906b8bde4078cd9f222c6480 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Fri, 21 Aug 2026 10:45:25 +0200 Subject: [PATCH 2/2] Test schedule_from setup error cleanup --- test/nvexec/continues_on.cpp | 57 ++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/test/nvexec/continues_on.cpp b/test/nvexec/continues_on.cpp index 29e250c17..d2a6a9100 100644 --- a/test/nvexec/continues_on.cpp +++ b/test/nvexec/continues_on.cpp @@ -5,6 +5,7 @@ #include "nvexec/stream_context.cuh" #include +#include namespace { @@ -12,9 +13,16 @@ namespace { std::size_t allocations_{}; std::size_t deallocations_{}; + std::size_t fail_after_{}; + bool fail_enabled_{}; void* do_allocate(std::size_t bytes, std::size_t) override { + if (fail_enabled_ && allocations_ == fail_after_) + { + throw std::bad_alloc(); + } + void* storage{}; STDEXEC_TRY_CUDA_API(cudaMallocHost(&storage, bytes)); ++allocations_; @@ -42,6 +50,12 @@ namespace { return deallocations_; } + + void fail_after(std::size_t allocation) noexcept + { + fail_after_ = allocation; + fail_enabled_ = true; + } }; struct noop_receiver @@ -62,6 +76,27 @@ namespace void set_stopped() noexcept {} }; + struct error_receiver + { + using receiver_concept = STDEXEC::receiver_tag; + + cudaError_t* error_; + + auto get_env() const noexcept -> STDEXEC::env<> + { + return {}; + } + + void set_value() noexcept {} + + void set_error(cudaError_t error) noexcept + { + *error_ = error; + } + + void set_stopped() noexcept {} + }; + class destruction_probe_t { flags_storage_t<>::flags_t flags_; @@ -120,6 +155,28 @@ namespace REQUIRE(pinned_memory.allocations() == pinned_memory.deallocations()); } + TEST_CASE("schedule_from frees its task when setup fails", + "[cuda][stream][adaptors][schedule_from]") + { + pinned_memory_resource_t pinned_memory; + pinned_memory.fail_after(2); + + nvexec::stream_context ctx; + auto scheduler = ctx.get_scheduler(); + scheduler.ctx_.pinned_resource_ = &pinned_memory; + + cudaError_t error = cudaSuccess; + auto sndr = STDEXEC::schedule_from(STDEXEC::schedule(scheduler)); + { + auto op = STDEXEC::connect(std::move(sndr), error_receiver{&error}); + STDEXEC::start(op); + } + + REQUIRE(error == cudaErrorMemoryAllocation); + REQUIRE(pinned_memory.allocations() == 2); + REQUIRE(pinned_memory.allocations() == pinned_memory.deallocations()); + } + TEST_CASE("continues on after schedule", "[cuda][stream][adaptors][continues_on]") { nvexec::stream_context ctx;