From ec403fa33a75902eb2bfa0e1758d54c1f3645945 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Fri, 10 Jul 2026 08:32:39 +0200 Subject: [PATCH 1/6] [STF] Add exec places from externally-owned CUDA contexts Allow an exec_place to wrap a CUcontext created outside CUDASTF, so SM partitioning can be owned by one library (e.g. cuda.core in Python) and shared across frameworks: the same green context can back a CUDASTF exec place and, for instance, a Warp device mapped from the same handle. - cudax: new exec_place_cuda_ctx_impl generalizing the green-context place; exec_place::cuda_context(ctx, devid = -1) factory (devid derived via cuCtxGetDevice when omitted). exec_place::green_ctx now builds the same impl, converting the CUgreenCtx once at construction instead of on every activate(). Identity (hash/cmp) is keyed on the converted CUcontext, which is stable per green context, so places built from a green_ctx_view and from the converted context compare equal. - C API: stf_exec_place_cuda_context(CUcontext, int dev_id). - Header unittests for equality, device-ordinal derivation, the activate/deactivate round trip, and cross-route equality with the green-context place; C API test running a task on a place built from the primary context. Places are non-owning: the caller must keep the context alive while the place is in use. Co-Authored-By: Claude Fable 5 --- .../stf/include/cccl/c/experimental/stf/stf.h | 9 + c/experimental/stf/src/stf.cu | 11 + c/experimental/stf/test/test_places.cpp | 57 ++++ .../__places/exec/cuda_context.cuh | 253 ++++++++++++++++++ .../__places/exec/green_context.cuh | 139 +++------- .../cuda/experimental/__places/places.cuh | 16 +- cudax/include/cuda/experimental/places.cuh | 1 + cudax/include/cuda/experimental/stf.cuh | 1 + cudax/test/places/CMakeLists.txt | 1 + 9 files changed, 385 insertions(+), 103 deletions(-) create mode 100644 cudax/include/cuda/experimental/__places/exec/cuda_context.cuh diff --git a/c/experimental/stf/include/cccl/c/experimental/stf/stf.h b/c/experimental/stf/include/cccl/c/experimental/stf/stf.h index 489cc64a089..fc6fb17a0ee 100644 --- a/c/experimental/stf/include/cccl/c/experimental/stf/stf.h +++ b/c/experimental/stf/include/cccl/c/experimental/stf/stf.h @@ -161,6 +161,15 @@ stf_exec_place_handle stf_exec_place_device(int dev_id); //! \brief Create execution place for the current CUDA device. stf_exec_place_handle stf_exec_place_current_device(void); +//! \brief Create an execution place from an externally-owned CUDA driver context \p ctx. +//! +//! The place is non-owning: the caller must keep \p ctx alive while the place is in +//! use. This is the natural entry point for contexts created by other libraries, e.g. +//! green contexts converted with cuCtxFromGreenCtx (such as the ones produced by +//! cuda.core in Python). \p dev_id is the device ordinal of the context, or -1 to +//! derive it from the context. Returns NULL on failure. +stf_exec_place_handle stf_exec_place_cuda_context(CUcontext ctx, int dev_id); + //! \brief Create a green-context helper for \p dev_id with \p sm_count SMs per green context. //! Requires CUDA 12.4+. Returns NULL on failure. stf_green_context_helper_handle stf_green_context_helper_create(int sm_count, int dev_id); diff --git a/c/experimental/stf/src/stf.cu b/c/experimental/stf/src/stf.cu index bb724347134..02092c96f55 100644 --- a/c/experimental/stf/src/stf.cu +++ b/c/experimental/stf/src/stf.cu @@ -219,6 +219,17 @@ stf_exec_place_handle stf_exec_place_current_device(void) })); } +stf_exec_place_handle stf_exec_place_cuda_context(CUcontext ctx, int dev_id) +{ + if (ctx == nullptr) + { + return nullptr; + } + return to_opaque(stf_try_allocate([ctx, dev_id] { + return new exec_place(exec_place::cuda_context(ctx, dev_id)); + })); +} + stf_green_context_helper_handle stf_green_context_helper_create(int sm_count, int dev_id) { #if _CCCL_CTK_AT_LEAST(12, 4) diff --git a/c/experimental/stf/test/test_places.cpp b/c/experimental/stf/test/test_places.cpp index 06bc074bb49..b88c76692bd 100644 --- a/c/experimental/stf/test/test_places.cpp +++ b/c/experimental/stf/test/test_places.cpp @@ -41,6 +41,63 @@ static void blocked_mapper_1d(stf_pos4* result, stf_pos4 data_coords, stf_dim4 d result->t = 0; } +C2H_TEST("exec place from an externally-owned CUDA context", "[task][places][cuda_context]") +{ + const size_t N = 1024; + + // Wrap the primary context of device 0 as an exec place + CUdevice dev = 0; + REQUIRE(cuDeviceGet(&dev, 0) == CUDA_SUCCESS); + CUcontext primary_ctx = nullptr; + REQUIRE(cuDevicePrimaryCtxRetain(&primary_ctx, dev) == CUDA_SUCCESS); + + // Null context is rejected + REQUIRE(stf_exec_place_cuda_context(nullptr, 0) == nullptr); + + // dev_id < 0 is derived from the context + stf_exec_place_handle place_derived = stf_exec_place_cuda_context(primary_ctx, -1); + REQUIRE(place_derived != nullptr); + REQUIRE(stf_exec_place_is_device(place_derived) != 0); + stf_exec_place_destroy(place_derived); + + stf_exec_place_handle place = stf_exec_place_cuda_context(primary_ctx, 0); + REQUIRE(place != nullptr); + REQUIRE(stf_exec_place_is_device(place) != 0); + REQUIRE(stf_exec_place_is_host(place) == 0); + + // Run a task on the place and fill the buffer through its stream + stf_ctx_handle ctx = stf_ctx_create(); + REQUIRE(ctx != nullptr); + + std::vector X(N, 1.0f); + stf_logical_data_handle lX = stf_logical_data(ctx, X.data(), N * sizeof(float)); + REQUIRE(lX != nullptr); + + stf_task_handle t = stf_task_create(ctx); + REQUIRE(t != nullptr); + stf_task_set_exec_place(t, place); + stf_task_add_dep(t, lX, STF_RW); + stf_task_start(t); + CUstream stream = stf_task_get_custream(t); + REQUIRE(stream != nullptr); + float* dX = static_cast(stf_task_get(t, 0)); + REQUIRE(dX != nullptr); + REQUIRE(cudaMemsetAsync(dX, 0, N * sizeof(float), stream) == cudaSuccess); + stf_task_end(t); + stf_task_destroy(t); + + stf_logical_data_destroy(lX); + stf_ctx_finalize(ctx); + + for (size_t i = 0; i < N; i++) + { + REQUIRE(X[i] == 0.0f); + } + + stf_exec_place_destroy(place); + REQUIRE(cuDevicePrimaryCtxRelease(dev) == CUDA_SUCCESS); +} + C2H_TEST("empty stf tasks", "[task]") { size_t N = 1000000; diff --git a/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh new file mode 100644 index 00000000000..bcb84037bce --- /dev/null +++ b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh @@ -0,0 +1,253 @@ +//===----------------------------------------------------------------------===// +// +// Part of CUDASTF in CUDA C++ Core Libraries, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +/** + * @file + * @brief Execution place wrapping an externally-owned CUDA driver context + * + * This makes it possible to use a CUcontext created outside CUDASTF (e.g. a + * green context created through cuda.core in Python, or any other library) + * as an execution place. The place is non-owning: the caller must keep the + * context alive while the place (and the streams lazily created in its pool) + * is in use. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +namespace cuda::experimental::places +{ +/** + * @brief Implementation for execution places backed by an externally-owned CUcontext + * + * The context is used as-is: `activate()` saves the current context and makes + * this one current, `deactivate()` restores the saved context. Streams are + * created lazily in the place's own pool while the context is current, so they + * inherit whatever resources the context carries (e.g. the SM partition of a + * green context converted with `cuCtxFromGreenCtx`). + * + * Identity (`hash`/`cmp`) is keyed on the CUcontext handle, which uniquely + * identifies the underlying resources: `cuCtxFromGreenCtx` returns the same + * CUcontext for a given CUgreenCtx on every call. + */ +class exec_place_cuda_ctx_impl : public exec_place::impl +{ +public: + /** + * @brief Construct an execution place from an externally-owned CUDA context + * + * @param ctx The CUDA driver context. Non-owning: the caller keeps it alive. + * @param devid The device ordinal the context belongs to, or -1 to derive it + * from the context (via cuCtxGetDevice). + * @param pool_size Number of streams in the place's stream pool. + */ + exec_place_cuda_ctx_impl(CUcontext ctx, int devid = -1, size_t pool_size = exec_place::impl::pool_size) + : exec_place_cuda_ctx_impl( + ctx, resolve_devid(ctx, devid), stream_pool(pool_size), data_place::device(resolve_devid(ctx, devid))) + {} + + /** + * @brief Full-control constructor, also used by the green-context place + * + * @param ctx The CUDA driver context (already resolved, e.g. from cuCtxFromGreenCtx) + * @param devid The device ordinal (must be valid) + * @param pool A stream pool to use for this place (shared handle) + * @param affine The affine data place for this place + */ + exec_place_cuda_ctx_impl(CUcontext ctx, int devid, stream_pool pool, data_place affine) + : exec_place::impl(mv(affine)) + , devid_(devid) + , driver_context_(ctx) + , pool_(mv(pool)) + { + _CCCL_ASSERT(ctx != nullptr, "cuda_ctx exec_place requires a valid CUcontext"); + _CCCL_ASSERT(devid_ >= 0, "cuda_ctx exec_place requires a valid device ordinal"); + } + + ::std::shared_ptr get_place(size_t idx) override + { + _CCCL_ASSERT(idx == 0, "Index out of bounds for cuda_ctx exec_place"); + return shared_from_this(); + } + + exec_place activate(size_t idx) const override + { + _CCCL_ASSERT(idx == 0, "Index out of bounds for cuda_ctx exec_place"); + + // Save the current context wrapped as a place so deactivate can restore it + CUcontext current_ctx = cuda_try(); + exec_place result = exec_place(::std::make_shared(saved_tag{}, current_ctx)); + + cuda_try(driver_context_); + + return result; + } + + void deactivate(const exec_place& prev, size_t idx = 0) const override + { + _CCCL_ASSERT(idx == 0, "Index out of bounds for cuda_ctx exec_place"); + + auto prev_impl = ::std::static_pointer_cast(prev.get_impl()); + CUcontext saved_ctx = prev_impl->driver_context_; + + cuda_try(saved_ctx); + } + + bool is_device() const override + { + return true; + } + + ::std::string to_string() const override + { + return "cuda_ctx(ctx=" + ::std::to_string(reinterpret_cast<::std::uintptr_t>(driver_context_)) + + " dev=" + ::std::to_string(devid_) + ")"; + } + + stream_pool& get_stream_pool(bool, exec_place_resources&, const exec_place&) const override + { + // This place carries its own pool and bypasses the registry. The user is + // responsible for keeping the underlying CUcontext alive while the pool + // is in use. + return pool_; + } + + int cmp(const exec_place::impl& rhs) const override + { + if (typeid(*this) != typeid(rhs)) + { + return typeid(*this).before(typeid(rhs)) ? -1 : 1; + } + const auto& other = static_cast(rhs); + return (other.driver_context_ < driver_context_) - (driver_context_ < other.driver_context_); + } + + size_t hash() const override + { + return ::std::hash()(driver_context_); + } + +protected: + // Tag type for the internal "saved context" wrapper used by activate/deactivate. + // The type itself is protected so only this class (and derived classes) can + // name it; the constructor below must be public for make_shared. + struct saved_tag + {}; + +public: + // Wrap an existing context with no pool: only used to carry the saved + // context through the activate/deactivate round trip. + exec_place_cuda_ctx_impl(saved_tag, CUcontext saved_context) + : driver_context_(saved_context) + {} + +protected: + static int resolve_devid(CUcontext ctx, int devid) + { + if (devid >= 0) + { + return devid; + } + + // Derive the device ordinal from the context + cuda_try(ctx); + CUdevice dev = cuda_try(); + [[maybe_unused]] const CUcontext popped = cuda_try(); + return static_cast(dev); + } + + int devid_ = -1; + CUcontext driver_context_ = {}; + mutable stream_pool pool_; +}; + +inline exec_place exec_place::cuda_context(CUcontext ctx, int devid, size_t pool_size) +{ + return exec_place(::std::make_shared(ctx, devid, pool_size)); +} + +#ifdef UNITTESTED_FILE +namespace +{ +//! RAII holder for the primary context of device 0, used by the unittests below. +struct primary_ctx_guard +{ + primary_ctx_guard() + { + cuda_try(0); + dev = cuda_try(0); + ctx = cuda_try(dev); + } + + ~primary_ctx_guard() + { + cuda_try(cuDevicePrimaryCtxRelease(dev)); + } + + CUdevice dev = -1; + CUcontext ctx = nullptr; +}; +} // namespace + +UNITTEST("cuda_context exec_place equality") +{ + primary_ctx_guard guard; + + auto p0a = exec_place::cuda_context(guard.ctx, 0); + auto p0b = exec_place::cuda_context(guard.ctx, 0); + + // Same context should be equal + EXPECT(p0a == p0b); + EXPECT(!(p0a != p0b)); + + // A cuda_context place is not a regular device place + auto dev0 = exec_place::device(0); + EXPECT(p0a != dev0); + EXPECT(!(p0a == dev0)); +}; + +UNITTEST("cuda_context exec_place derives the device ordinal") +{ + primary_ctx_guard guard; + + // devid intentionally omitted: derived from the context via cuCtxGetDevice + auto p = exec_place::cuda_context(guard.ctx); + EXPECT(p.is_device()); + EXPECT(p.affine_data_place() == data_place::device(0)); + EXPECT(p == exec_place::cuda_context(guard.ctx, 0)); +}; + +UNITTEST("cuda_context exec_place activate/deactivate round trip") +{ + primary_ctx_guard guard; + + auto p = exec_place::cuda_context(guard.ctx, 0); + + CUcontext before = cuda_try(); + { + exec_place_scope scope(p); + EXPECT(cuda_try() == guard.ctx); + } + EXPECT(cuda_try() == before); +}; +#endif // UNITTESTED_FILE +} // end namespace cuda::experimental::places diff --git a/cudax/include/cuda/experimental/__places/exec/green_context.cuh b/cudax/include/cuda/experimental/__places/exec/green_context.cuh index 2944209604c..ecb4b3d728b 100644 --- a/cudax/include/cuda/experimental/__places/exec/green_context.cuh +++ b/cudax/include/cuda/experimental/__places/exec/green_context.cuh @@ -26,6 +26,7 @@ #endif // no system header #include +#include #include #include #include @@ -255,111 +256,24 @@ private: ::std::vector ctxs; }; -/** - * @brief Implementation for green context execution places +/* + * Green-context execution places are implemented with the generic + * externally-owned-context place (`exec_place_cuda_ctx_impl`): the only + * functional use of the CUgreenCtx is deriving its CUcontext, and + * cuCtxFromGreenCtx returns the same CUcontext on every call, so the + * conversion is done once here and the CUcontext is the canonical identity + * of the place. As a consequence, a place built from a green_ctx_view and a + * place built with exec_place::cuda_context() from the converted context + * compare equal, which is the desired semantics. + * + * The user is responsible for keeping the underlying CUgreenCtx alive while + * the place (and its stream pool) is in use. */ -class exec_place_green_ctx_impl : public exec_place::impl -{ -public: - /** - * @brief Construct a green context execution place - * - * @param gc_view The green context view - * @param use_green_ctx_data_place If true, use a green context data place as the - * affine data place. If false (default), use a regular device data place instead. - */ - exec_place_green_ctx_impl(green_ctx_view gc_view, bool use_green_ctx_data_place = false) - : exec_place::impl( - use_green_ctx_data_place ? make_green_ctx_data_place(gc_view) : data_place::device(gc_view.devid)) - , devid_(gc_view.devid) - , g_ctx_(gc_view.g_ctx) - , pool_(mv(gc_view.pool)) - {} - - // This is used to implement deactivate and wrap an existing context - exec_place_green_ctx_impl(CUcontext saved_context) - : driver_context_(saved_context) - {} - - ::std::shared_ptr get_place(size_t idx) override - { - _CCCL_ASSERT(idx == 0, "Index out of bounds for green_ctx exec_place"); - return shared_from_this(); - } - - exec_place activate(size_t idx) const override - { - _CCCL_ASSERT(idx == 0, "Index out of bounds for green_ctx exec_place"); - - // Save the current context and transform it into a fake green context place - CUcontext current_ctx = cuda_try(); - exec_place result = exec_place(::std::make_shared(current_ctx)); - - // Convert the green context to a primary context - driver_context_ = cuda_try(g_ctx_); - cuda_try(cuCtxSetCurrent(driver_context_)); - - return result; - } - - void deactivate(const exec_place& prev, size_t idx = 0) const override - { - _CCCL_ASSERT(idx == 0, "Index out of bounds for green_ctx exec_place"); - - auto prev_impl = ::std::static_pointer_cast(prev.get_impl()); - CUcontext saved_ctx = prev_impl->driver_context_; - -# ifdef DEBUG - CUcontext current_ctx = cuda_try(); - assert(get_cuda_context_id(current_ctx) == get_cuda_context_id(driver_context_)); -# endif - - cuda_try(cuCtxSetCurrent(saved_ctx)); - } - - bool is_device() const override - { - return true; - } - - ::std::string to_string() const override - { - return "green_ctx(id=" + ::std::to_string(get_cuda_context_id(g_ctx_)) + " dev=" + ::std::to_string(devid_) + ")"; - } - - stream_pool& get_stream_pool(bool, exec_place_resources&, const exec_place&) const override - { - // Green-context places carry their own pool (constructed from the - // green_ctx_view) and bypass the registry. The user is responsible for - // keeping the underlying CUgreenCtx alive while the pool is in use. - return pool_; - } - - int cmp(const exec_place::impl& rhs) const override - { - if (typeid(*this) != typeid(rhs)) - { - return typeid(*this).before(typeid(rhs)) ? -1 : 1; - } - const auto& other = static_cast(rhs); - return (other.g_ctx_ < g_ctx_) - (g_ctx_ < other.g_ctx_); - } - - size_t hash() const override - { - return ::std::hash()(g_ctx_); - } - -private: - int devid_ = -1; - CUgreenCtx g_ctx_ = {}; - mutable CUcontext driver_context_ = {}; - mutable stream_pool pool_; -}; - inline exec_place exec_place::green_ctx(const green_ctx_view& gc_view, bool use_green_ctx_data_place) { - return exec_place(::std::make_shared(gc_view, use_green_ctx_data_place)); + CUcontext ctx = cuda_try(gc_view.g_ctx); + data_place affine = use_green_ctx_data_place ? make_green_ctx_data_place(gc_view) : data_place::device(gc_view.devid); + return exec_place(::std::make_shared(ctx, gc_view.devid, gc_view.pool, mv(affine))); } inline ::std::shared_ptr green_ctx_data_place_impl::get_affine_exec_impl() const @@ -405,6 +319,27 @@ UNITTEST("green context exec_place equality") EXPECT(!(p0a == dev0)); }; +UNITTEST("green_ctx place equals the cuda_context place wrapping the same partition") +{ + green_context_helper gc_helper(8, 0); // 8 SMs per green context + + if (gc_helper.get_count() < 1) + { + return; + } + + auto view = gc_helper.get_view(0); + + // Identity is keyed on the converted CUcontext, which is a stable accessor + // (cuCtxFromGreenCtx returns the same handle on every call), so both + // construction routes must yield equal places. + auto p_green = exec_place::green_ctx(view); + auto p_ctx = exec_place::cuda_context(cuda_try(view.g_ctx), view.devid); + + EXPECT(p_green == p_ctx); + EXPECT(!(p_green != p_ctx)); +}; + UNITTEST("green context data_place equality") { green_context_helper gc_helper(8, 0); diff --git a/cudax/include/cuda/experimental/__places/places.cuh b/cudax/include/cuda/experimental/__places/places.cuh index 61df652fce2..6c8600f35f6 100644 --- a/cudax/include/cuda/experimental/__places/places.cuh +++ b/cudax/include/cuda/experimental/__places/places.cuh @@ -518,7 +518,7 @@ public: * impls). * * Self-contained implementations (`exec_place_cuda_stream_impl`, - * `exec_place_green_ctx_impl`) override this method and ignore the + * `exec_place_cuda_ctx_impl`) override this method and ignore the * registry, returning their embedded pool instead. * * The grid implementation forwards `res` to its first sub-place. @@ -783,6 +783,20 @@ public: static exec_place cuda_stream(cudaStream_t stream); static exec_place cuda_stream(const augmented_stream& dstream); + /** + * @brief Create an execution place from an externally-owned CUDA driver context + * + * The place is non-owning: the caller must keep the context alive while the + * place is in use. This is the natural entry point for contexts created by + * other libraries (e.g. green contexts converted with cuCtxFromGreenCtx, such + * as the ones produced by cuda.core in Python). + * + * @param ctx The CUDA driver context + * @param devid The device ordinal of the context, or -1 to derive it from the context + * @param pool_size Number of streams in the place's stream pool + */ + static exec_place cuda_context(CUcontext ctx, int devid = -1, size_t pool_size = impl::pool_size); + /** * @brief Returns the currently active device. * diff --git a/cudax/include/cuda/experimental/places.cuh b/cudax/include/cuda/experimental/places.cuh index 6da44fb7bfe..65a31bc92bc 100644 --- a/cudax/include/cuda/experimental/places.cuh +++ b/cudax/include/cuda/experimental/places.cuh @@ -20,6 +20,7 @@ #pragma once +#include #include #include #include diff --git a/cudax/include/cuda/experimental/stf.cuh b/cudax/include/cuda/experimental/stf.cuh index 80db4950c82..f7d05a3f67b 100644 --- a/cudax/include/cuda/experimental/stf.cuh +++ b/cudax/include/cuda/experimental/stf.cuh @@ -28,6 +28,7 @@ #include #include // #include +#include #include #include #include diff --git a/cudax/test/places/CMakeLists.txt b/cudax/test/places/CMakeLists.txt index 2619fbf5a19..da948c6bbc8 100644 --- a/cudax/test/places/CMakeLists.txt +++ b/cudax/test/places/CMakeLists.txt @@ -12,6 +12,7 @@ set(places_fail_tests exec_place_scope_data_place_fail.cu) set( places_unittested_headers cuda/experimental/__places/places.cuh + cuda/experimental/__places/exec/cuda_context.cuh cuda/experimental/__places/exec/green_context.cuh cuda/experimental/__places/partitions/blocked_partition.cuh cuda/experimental/__places/partitions/cyclic_shape.cuh From 94d24615a7668005c1ecc4cd1dc1859eb8250bd9 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Fri, 10 Jul 2026 09:23:30 +0200 Subject: [PATCH 2/6] [STF] Reject a null CUcontext in exec_place::cuda_context A null context with an explicit dev_id would construct a place that only misbehaves at activation time (cuCtxSetCurrent(nullptr) succeeds and unbinds the current context). Reject it eagerly instead: - exec_place::cuda_context throws std::invalid_argument (works in release builds, matching the logic_error throws elsewhere in places). - The C API asserts in debug builds; in release the exception is mapped to a null handle with a stderr trace by stf_try_allocate, like every other failure on this surface. - Header unittest for the rejection; the C API test already covers the null-handle return. Co-Authored-By: Claude Fable 5 --- .../stf/include/cccl/c/experimental/stf/stf.h | 3 ++- c/experimental/stf/src/stf.cu | 7 +++---- .../__places/exec/cuda_context.cuh | 21 +++++++++++++++++++ 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/c/experimental/stf/include/cccl/c/experimental/stf/stf.h b/c/experimental/stf/include/cccl/c/experimental/stf/stf.h index fc6fb17a0ee..241e6bc2255 100644 --- a/c/experimental/stf/include/cccl/c/experimental/stf/stf.h +++ b/c/experimental/stf/include/cccl/c/experimental/stf/stf.h @@ -167,7 +167,8 @@ stf_exec_place_handle stf_exec_place_current_device(void); //! use. This is the natural entry point for contexts created by other libraries, e.g. //! green contexts converted with cuCtxFromGreenCtx (such as the ones produced by //! cuda.core in Python). \p dev_id is the device ordinal of the context, or -1 to -//! derive it from the context. Returns NULL on failure. +//! derive it from the context. \p ctx must not be NULL. Returns NULL on failure +//! (invalid context, allocation failure), with a diagnostic printed to stderr. stf_exec_place_handle stf_exec_place_cuda_context(CUcontext ctx, int dev_id); //! \brief Create a green-context helper for \p dev_id with \p sm_count SMs per green context. diff --git a/c/experimental/stf/src/stf.cu b/c/experimental/stf/src/stf.cu index 02092c96f55..3649196f422 100644 --- a/c/experimental/stf/src/stf.cu +++ b/c/experimental/stf/src/stf.cu @@ -221,10 +221,9 @@ stf_exec_place_handle stf_exec_place_current_device(void) stf_exec_place_handle stf_exec_place_cuda_context(CUcontext ctx, int dev_id) { - if (ctx == nullptr) - { - return nullptr; - } + _CCCL_ASSERT(ctx != nullptr, "CUcontext must not be null"); + // A null context in release builds throws in exec_place::cuda_context and is + // mapped to a null handle (with a stderr trace) by stf_try_allocate. return to_opaque(stf_try_allocate([ctx, dev_id] { return new exec_place(exec_place::cuda_context(ctx, dev_id)); })); diff --git a/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh index bcb84037bce..6d4f3694c44 100644 --- a/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh +++ b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh @@ -182,6 +182,13 @@ protected: inline exec_place exec_place::cuda_context(CUcontext ctx, int devid, size_t pool_size) { + if (ctx == nullptr) + { + // Reject eagerly: with an explicit devid a null context would construct a + // place that only fails (or silently unbinds the current context) at + // activation time. + throw ::std::invalid_argument("exec_place::cuda_context requires a valid CUcontext"); + } return exec_place(::std::make_shared(ctx, devid, pool_size)); } @@ -236,6 +243,20 @@ UNITTEST("cuda_context exec_place derives the device ordinal") EXPECT(p == exec_place::cuda_context(guard.ctx, 0)); }; +UNITTEST("cuda_context exec_place rejects a null context") +{ + bool thrown = false; + try + { + auto p = exec_place::cuda_context(nullptr, 0); + } + catch (const ::std::invalid_argument&) + { + thrown = true; + } + EXPECT(thrown); +}; + UNITTEST("cuda_context exec_place activate/deactivate round trip") { primary_ctx_guard guard; From 9e48d8390dab7a13eee32e2bedac28b2b5a4f5a8 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Wed, 15 Jul 2026 16:08:03 +0200 Subject: [PATCH 3/6] [STF] Restore CUDA context with RAII Use the existing context guard so device lookup restores the caller's context when querying the device fails. --- .../cuda/experimental/__places/exec/cuda_context.cuh | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh index 6d4f3694c44..6467315130e 100644 --- a/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh +++ b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh @@ -22,6 +22,7 @@ #pragma once #include +#include #if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) # pragma GCC system_header @@ -169,10 +170,8 @@ protected: } // Derive the device ordinal from the context - cuda_try(ctx); - CUdevice dev = cuda_try(); - [[maybe_unused]] const CUcontext popped = cuda_try(); - return static_cast(dev); + ::cuda::__ensure_current_context guard{ctx}; + return static_cast(cuda_try()); } int devid_ = -1; From 55a49cc4200aa8511cd7a5afdfeb684f0c2ba9e5 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Wed, 15 Jul 2026 16:21:55 +0200 Subject: [PATCH 4/6] [STF] Validate external CUDA context places Keep context identity and device affinity consistent, preserve C API error translation, and use portable context ordering. --- c/experimental/stf/src/stf.cu | 3 -- c/experimental/stf/test/test_places.cpp | 42 +++++++++---------- .../__places/exec/cuda_context.cuh | 39 +++++++++++++---- 3 files changed, 51 insertions(+), 33 deletions(-) diff --git a/c/experimental/stf/src/stf.cu b/c/experimental/stf/src/stf.cu index 3649196f422..4bbcf89c9b0 100644 --- a/c/experimental/stf/src/stf.cu +++ b/c/experimental/stf/src/stf.cu @@ -221,9 +221,6 @@ stf_exec_place_handle stf_exec_place_current_device(void) stf_exec_place_handle stf_exec_place_cuda_context(CUcontext ctx, int dev_id) { - _CCCL_ASSERT(ctx != nullptr, "CUcontext must not be null"); - // A null context in release builds throws in exec_place::cuda_context and is - // mapped to a null handle (with a stderr trace) by stf_try_allocate. return to_opaque(stf_try_allocate([ctx, dev_id] { return new exec_place(exec_place::cuda_context(ctx, dev_id)); })); diff --git a/c/experimental/stf/test/test_places.cpp b/c/experimental/stf/test/test_places.cpp index b88c76692bd..fb0de7dfab9 100644 --- a/c/experimental/stf/test/test_places.cpp +++ b/c/experimental/stf/test/test_places.cpp @@ -43,7 +43,7 @@ static void blocked_mapper_1d(stf_pos4* result, stf_pos4 data_coords, stf_dim4 d C2H_TEST("exec place from an externally-owned CUDA context", "[task][places][cuda_context]") { - const size_t N = 1024; + constexpr size_t element_count{1024}; // Wrap the primary context of device 0 as an exec place CUdevice dev = 0; @@ -55,43 +55,43 @@ C2H_TEST("exec place from an externally-owned CUDA context", "[task][places][cud REQUIRE(stf_exec_place_cuda_context(nullptr, 0) == nullptr); // dev_id < 0 is derived from the context - stf_exec_place_handle place_derived = stf_exec_place_cuda_context(primary_ctx, -1); + const stf_exec_place_handle place_derived = stf_exec_place_cuda_context(primary_ctx, -1); REQUIRE(place_derived != nullptr); REQUIRE(stf_exec_place_is_device(place_derived) != 0); stf_exec_place_destroy(place_derived); - stf_exec_place_handle place = stf_exec_place_cuda_context(primary_ctx, 0); + const stf_exec_place_handle place = stf_exec_place_cuda_context(primary_ctx, 0); REQUIRE(place != nullptr); REQUIRE(stf_exec_place_is_device(place) != 0); REQUIRE(stf_exec_place_is_host(place) == 0); // Run a task on the place and fill the buffer through its stream - stf_ctx_handle ctx = stf_ctx_create(); + const stf_ctx_handle ctx = stf_ctx_create(); REQUIRE(ctx != nullptr); - std::vector X(N, 1.0f); - stf_logical_data_handle lX = stf_logical_data(ctx, X.data(), N * sizeof(float)); - REQUIRE(lX != nullptr); + std::vector x(element_count, 1.0f); + const stf_logical_data_handle logical_x = stf_logical_data(ctx, x.data(), element_count * sizeof(float)); + REQUIRE(logical_x != nullptr); - stf_task_handle t = stf_task_create(ctx); - REQUIRE(t != nullptr); - stf_task_set_exec_place(t, place); - stf_task_add_dep(t, lX, STF_RW); - stf_task_start(t); - CUstream stream = stf_task_get_custream(t); + const stf_task_handle task = stf_task_create(ctx); + REQUIRE(task != nullptr); + stf_task_set_exec_place(task, place); + stf_task_add_dep(task, logical_x, STF_RW); + stf_task_start(task); + const CUstream stream = stf_task_get_custream(task); REQUIRE(stream != nullptr); - float* dX = static_cast(stf_task_get(t, 0)); - REQUIRE(dX != nullptr); - REQUIRE(cudaMemsetAsync(dX, 0, N * sizeof(float), stream) == cudaSuccess); - stf_task_end(t); - stf_task_destroy(t); + float* const device_x = static_cast(stf_task_get(task, 0)); + REQUIRE(device_x != nullptr); + REQUIRE(cudaMemsetAsync(device_x, 0, element_count * sizeof(float), stream) == cudaSuccess); + stf_task_end(task); + stf_task_destroy(task); - stf_logical_data_destroy(lX); + stf_logical_data_destroy(logical_x); stf_ctx_finalize(ctx); - for (size_t i = 0; i < N; i++) + for (size_t i = 0; i < element_count; i++) { - REQUIRE(X[i] == 0.0f); + REQUIRE(x[i] == 0.0f); } stf_exec_place_destroy(place); diff --git a/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh index 6467315130e..8c1849a6cbe 100644 --- a/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh +++ b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh @@ -35,6 +35,8 @@ #include #include +#include + namespace cuda::experimental::places { /** @@ -62,8 +64,7 @@ public: * @param pool_size Number of streams in the place's stream pool. */ exec_place_cuda_ctx_impl(CUcontext ctx, int devid = -1, size_t pool_size = exec_place::impl::pool_size) - : exec_place_cuda_ctx_impl( - ctx, resolve_devid(ctx, devid), stream_pool(pool_size), data_place::device(resolve_devid(ctx, devid))) + : exec_place_cuda_ctx_impl(ctx, resolve_devid(ctx, devid), stream_pool(pool_size)) {} /** @@ -139,7 +140,8 @@ public: return typeid(*this).before(typeid(rhs)) ? -1 : 1; } const auto& other = static_cast(rhs); - return (other.driver_context_ < driver_context_) - (driver_context_ < other.driver_context_); + return ::std::less{}(other.driver_context_, driver_context_) + - ::std::less{}(driver_context_, other.driver_context_); } size_t hash() const override @@ -162,16 +164,19 @@ public: {} protected: + exec_place_cuda_ctx_impl(CUcontext ctx, int devid, stream_pool pool) + : exec_place_cuda_ctx_impl(ctx, devid, mv(pool), data_place::device(devid)) + {} + static int resolve_devid(CUcontext ctx, int devid) { - if (devid >= 0) + ::cuda::__ensure_current_context guard{ctx}; + const int context_devid = static_cast(cuda_try()); + if (devid >= 0 && devid != context_devid) { - return devid; + throw ::std::invalid_argument("CUcontext device ordinal does not match devid"); } - - // Derive the device ordinal from the context - ::cuda::__ensure_current_context guard{ctx}; - return static_cast(cuda_try()); + return context_devid; } int devid_ = -1; @@ -256,6 +261,22 @@ UNITTEST("cuda_context exec_place rejects a null context") EXPECT(thrown); }; +UNITTEST("cuda_context exec_place rejects a mismatched device ordinal") +{ + primary_ctx_guard guard; + + bool thrown = false; + try + { + auto p = exec_place::cuda_context(guard.ctx, 1); + } + catch (const ::std::invalid_argument&) + { + thrown = true; + } + EXPECT(thrown); +}; + UNITTEST("cuda_context exec_place activate/deactivate round trip") { primary_ctx_guard guard; From 8c3845b2b532233913d4f876f7d5f778eb1c9213 Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Wed, 15 Jul 2026 17:16:20 +0200 Subject: [PATCH 5/6] [STF] Document external CUDA context places List the context factory with the other execution-place factories and clarify its non-owning lifetime contract. --- docs/cudax/places.rst | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/cudax/places.rst b/docs/cudax/places.rst index 63a5b94734e..ea04856b3a2 100644 --- a/docs/cudax/places.rst +++ b/docs/cudax/places.rst @@ -34,6 +34,9 @@ The following factory methods create the most common execution places: - ``exec_place::device(id)`` -- a specific CUDA device - ``exec_place::host()`` -- the host CPU - ``exec_place::current_device()`` -- the CUDA device that is currently active +- ``exec_place::cuda_context(ctx, devid)`` -- an externally-owned CUDA driver + context; the device ordinal is derived from the context when ``devid`` is + omitted When an execution place is activated, it sets the appropriate CUDA context (e.g. calls ``cudaSetDevice``). Each execution place also has an *affine* @@ -41,6 +44,9 @@ data place: the memory location naturally associated with it. For a device execution place the affine data place is the device's global memory; for the host it is pinned host memory (RAM). +A CUDA-context execution place is non-owning. The caller must keep the +``CUcontext`` alive while the place and any streams obtained from it are in use. + .. _places-data-places: Data places From 592057b5670de6ef02b96956f0845cbe6d4775bc Mon Sep 17 00:00:00 2001 From: Cedric AUGONNET Date: Wed, 15 Jul 2026 17:52:01 +0200 Subject: [PATCH 6/6] [STF] Follow system-header include ordering Keep only the configuration header before the system-header pragma block, matching cudax header conventions. --- cudax/include/cuda/experimental/__places/exec/cuda_context.cuh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh index 8c1849a6cbe..26315e74272 100644 --- a/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh +++ b/cudax/include/cuda/experimental/__places/exec/cuda_context.cuh @@ -22,7 +22,6 @@ #pragma once #include -#include #if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) # pragma GCC system_header @@ -32,6 +31,8 @@ # pragma system_header #endif // no system header +#include + #include #include