[STF] C API and Python bindings for placement evaluation and structured partitions#9892
[STF] C API and Python bindings for placement evaluation and structured partitions#9892caugonnet wants to merge 657 commits into
Conversation
Exposes the opaque C handle as an integer so Python code can pass exec_place handles through ctypes to external C libraries (e.g. torq's multi-GPU sort). Made-with: Cursor
Expose green_context_helper and green_ctx place factories through the STF C and Python APIs so downstream code can create green-context places directly. Include focused coverage and the DeviceArray weakref fix needed to exercise the new bindings end-to-end. Made-with: Cursor
Stream pools used to live in process-global `static` members of
`exec_place::impl` singletons, which made cached `cudaStream_t`
handles outlive any individual STF context and the primary CUDA
context itself. After an external `cudaDeviceReset()` (Numba
`cuda.close()`, PyTorch teardown, etc.) the handles became stale
and the next STF call hit `cudaErrorContextIsDestroyed`.
Introduce `exec_place_resources`: a standalone, mutex-guarded
registry of per-place {compute_pool, data_pool} keyed by the
opaque `exec_place::impl*` pointer. It can be created and used
without an STF context. `async_resources_handle` now owns one
`exec_place_resources` so each STF context gets its own pools,
released when the handle is destroyed.
`exec_place::impl::get_stream_pool(bool, exec_place_resources&)`
becomes virtual: the default override does the registry lookup,
the cuda-stream impl returns its own fixed `dummy_pool_` and
ignores the registry, the grid impl forwards to its first
sub-place, and green-context impls keep their own pool.
Made-with: Cursor
Every internal `getStream` / `getDataStream` / `get_stream_pool` call now goes through `ctx.async_resources().get_place_resources()` (or an explicit registry) instead of the retired no-arg pooled API. This makes stream-pool lookups go through the per-handle registry introduced in the previous commit, so the lifetime of every stream returned by `pick_stream(...)` is exactly tied to the handle that owns it. The no-argument `exec_place::pick_stream() / getStream() / stream_pool_size() / pick_all_streams()` overloads are removed: callers must pass an `exec_place_resources&` (or an `async_resources_handle&` that provides one). Made-with: Cursor
… pools
- cudax/test/places/stream_pool.cu: pass an explicit
`exec_place_resources` to every `pick_stream(...)` and add new
`test_two_handles_isolation` and `test_reset_survives_with_fresh_registry`
tests that exercise the per-handle ownership and survival across
`cudaDeviceReset()`.
- cudax/test/stf/cpp/test_pick_stream{,_green_context}.cu: new
end-to-end tests covering `ctx.async_resources()`-borrowed pools
for plain device, host, and green-context places, registered in
cudax/test/stf/CMakeLists.txt.
- docs/cudax/places.rst: update examples to show the explicit
`pick_stream(resources)` / `pick_stream(ctx.async_resources())`
API and the standalone `exec_place_resources` use case.
Made-with: Cursor
…lace_resources
C API additions in c/experimental/stf:
- New `stf_exec_place_resources_handle` opaque type with
`stf_exec_place_resources_create() / _destroy()` for standalone
use, plus `stf_ctx_get_place_resources()` to borrow the pool
registry that lives inside an STF context.
- `stf_exec_place_pick_stream(res, place, for_computation)` now
takes an explicit registry handle.
- New stackable-context surface: `stf_stackable_ctx_create /
_finalize / _fence`, `stf_stackable_push_graph / pop`,
`stf_stackable_push_while / pop_while`, `stf_stackable_push_repeat
/ pop_repeat`, `stf_stackable_while_cond_scalar`,
`stf_stackable_logical_data{,_with_place,_empty}`,
`stf_stackable_token`, `stf_stackable_task_create /
_add_dep{,_with_dplace}`, and `stf_stackable_host_launch_*`.
- New tests: stackable host_launch test in test_host_launch.cu;
test_places.cpp split into standalone vs context-borrowed
pick_stream coverage and updated allocator path.
Python bindings (`_stf_bindings_impl.pyx`, `__init__.py`):
- `cdef class exec_place_resources` with `_borrow_from(ctx)` and
owned/borrowed lifetime tracking; `context.place_resources`
returns a borrowed view; `exec_place.pick_stream` requires a
resources argument.
- `cdef class stackable_context` exposing nested graph, while-loop
and repeat scopes, stackable logical data / tokens, stackable
tasks and stackable host launches; helpers `_AliveFlag` and
`_PrimaryContextPin` to make Python lifetimes safe across the
framework boundary.
Internal C++ stackable support: record nested-body graph stats
and add a `CUDASTF_DUMP_GRAPHS` debug-dot dump for nested bodies
in `stackable_ctx_impl.cuh`.
Tests (Python):
- `tests/stf/test_lifecycle.py` (new): regression suite for
context / logical-data / token / task lifetimes, including
`test_stackable_repeat_after_device_reset` which exercises the
per-handle `exec_place_resources` registry across a real
primary-context teardown.
- `tests/stf/test_place_support.py`: refactor `pick_stream` tests
into standalone, context-borrowed, requires-resources, and
two-handles-isolated variants.
- `tests/stf/test_context.py`: minor updates to match the new
context lifecycle and explicit-resources API.
Made-with: Cursor
End-to-end tests exercising the stackable-context Python API: - test_burger_stackable.py / test_burger_stackable_fast.py: 5-level nested stackable Burger solver (ctx.repeat / graph_scope / while_loop) using PyTorch tasks and Numba kernels respectively. - test_cg_stackable.py: conjugate-gradient solver using nested while_loop scopes. - test_jacobi_stackable_numba.py / test_jacobi_stackable_pytorch.py: Jacobi solvers covering both Numba and PyTorch task backends. - test_nested_stackable.py: targeted coverage of nested graph_scope / push / pop interactions, repeat counts, and while-loop scalar conditions. - test_stackable_graph_scope.py: graph_scope-only smoke + edge cases. Made-with: Cursor
Rewrite the tiled Cholesky and POTRI Python examples so all heavy BLAS/LAPACK kernels go through direct nvmath-python bindings (cuBLAS + cuSOLVER) instead of CuPy's high-level linalg helpers, with every cuSOLVER/cuBLAS scratch buffer declared as an STF ``logical_data_empty(...).write()`` task argument. This mirrors the C++ linear algebra examples under ``cudax/examples/stf/linear_algebra`` and eliminates hidden allocations and workspace churn through CuPy's memory pool on the numerical path. Both examples now reach machine-precision residuals under ``--check`` (order 1e-16 vs. 1e-8 previously). The only remaining CuPy usage is CUDA-runtime glue (pinned memory, device/stream/event APIs, timing). POTRI additionally uses a tiny module-level ``cp.RawKernel`` for the triangular copy-with-zero-fill needed by ``DLAAUM`` because ``cusolverDnDlacpy`` is not exposed by nvmath-python; it is compiled once and does not allocate on invocation. Also drops stale helpers that this refactor made obsolete (``CAIWrapper``, ``get_cupy_arrays``, the ``cupyx.scipy`` import). Made-with: Cursor
Expose stackable_ctx::logical_data_no_export through the C API as stf_stackable_logical_data_no_export_empty, and surface it in the Cython bindings via a new ``no_export`` keyword on ``stackable_context.logical_data_empty``. Such data is created at the current head context and is not exported to parent scopes, which is the right semantics for per-iteration temporaries inside while/repeat bodies. Made-with: Cursor
Expose stackable_ctx::logical_data::push through the C API as stf_stackable_logical_data_push(ld, mode, dplace), and surface it on the Python stackable_logical_data class via a new ``push(mode, dplace)`` method. Also re-export ``AccessMode`` from ``cuda.stf`` so callers can pass ``AccessMode.READ`` symbolically. By default, the first access from a nested scope auto-pushes a piece of data with a conservative (read-write) mode, which serialises sibling scopes that only need to read it. Calling ``ld.push(AccessMode.READ)`` inside each sibling scope lets those scopes execute concurrently without having to globally mark the data read-only via ``set_read_only()``. Made-with: Cursor
Brings the Python STF bindings to parity with the C++ idiom `stream_ctx ctx(stream, handle)` / `graph_ctx ctx(stream, handle)` used by legacy_to_stf.cu's `lib_call_with_handle`. Without this, each Python `stf.context(use_graph=True)` had to rebuild its instantiated CUDA graph from scratch, so the graph backend was strictly a loss on short DAGs. C facade (`c/experimental/stf/...`): - New opaque `stf_async_resources_handle` + `stf_async_resources_create` / `stf_async_resources_destroy` wrapping the C++ `async_resources_handle`. - New `stf_backend_kind` enum, `stf_ctx_options` struct (zero-init = default), and unified factory `stf_ctx_create_ex(const stf_ctx_options*)` covering every combination of backend / caller stream / shared resources handle. `opts == NULL` is equivalent to `stf_ctx_create()`. Existing `stf_ctx_create[_graph]()` entry points are left untouched. Cython + Python (`python/cuda_cccl_experimental/cuda/stf/...`): - New `cuda.stf.async_resources` class managing the shared handle's lifetime (RAII via `__dealloc__`). - `cuda.stf.context(...)` gained `stream=` and `handle=` kwargs. When either is provided the binding dispatches through `stf_ctx_create_ex`; otherwise it keeps calling the legacy no-arg factories so pre-existing call sites behave bit-identically. A keep-alive reference to the `async_resources` prevents the handle from being GC'd underneath a live context. Example (`tests/stf/test_legacy_to_stf.py`, new file): - Python port of `cudax/test/stf/local_stf/legacy_to_stf.cu` mirroring its four variants (legacy single-stream, STF with logical_data, STF with tokens, STF with tokens on the graph backend), plus a fifth variant exercising the new `stream=`/`handle=` kwargs. - pytest case `test_lib_call_token_shared_handle` back-to-back reuses one `async_resources` across two graph contexts on a caller-owned stream. - Benchmark sweep adds rows for the (+stream,+handle) variants. On small problem sizes where the graph backend used to lose (2.0x vs legacy), sharing a resources handle brings it to 0.75x -- i.e. 25% faster than the sequential baseline and ~2.7x faster than the cache-miss graph path. Made-with: Cursor
Expose a two-phase pop on stackable_ctx so the CUDA graph built in a
nested scope can be reused:
auto h = ctx.pop_prologue();
for (int k = 0; k < N; ++k) h.launch();
ctx.pop_epilogue();
The graph_ctx_node finalization path is split into four idempotent
phases (prepare_graph / ensure_instantiated / ensure_prereqs_synced /
launch_once + finalize_after_launch) so that the existing single-shot
pop() stays unchanged while the split pop can drive them
independently. pop_prologue() only runs phase 1
(cudaGraphInstantiate is deferred until launch() or exec() is
actually called), so callers that just want to embed the nested
cudaGraph_t into another graph via cudaGraphAddChildGraphNode pay no
instantiation cost.
The returned launchable_graph_handle exposes:
* launch() - dispatch once on the support stream (lazy dep-A sync
+ lazy instantiation on first call).
* exec() - cudaGraphExec_t, forces lazy instantiation + dep-A
sync on first call so cudaGraphLaunch(exec(), stream())
is always safe.
* stream() - support stream. Purely observational, no side effects.
* graph() - underlying cudaGraph_t. Does NOT force instantiation
but triggers the dep-A sync so stream() becomes a
ready event source for cross-stream ordering.
Misuse of the handle (use after pop_epilogue, stale handle from a
different pop, double pop_prologue) is detected via a shared/weak
token pair and aborts with a clear diagnostic.
Also add a launchable_graph_scope RAII helper mirroring
graph_scope_guard but driving pop_prologue/pop_epilogue, and a
matching UNITTEST suite covering: repeated launches accumulating,
zero-launch reusability, handle invalidation after pop_epilogue,
while-graph_scope interplay, and manual cudaGraphLaunch via
exec()/stream().
Made-with: Cursor
Expose the two-phase stackable_ctx::pop_prologue / pop_epilogue API
added in the previous commit through the C ABI and the Cython
bindings:
C:
stf_launchable_graph_handle (opaque)
stf_stackable_pop_prologue / _epilogue
stf_launchable_graph_launch / _exec / _stream / _graph
stf_launchable_graph_destroy
Python (Cython):
stackable_context.launchable_graph_scope()
-> context manager with .launch(), .exec_graph, .stream, .graph
Error semantics mirror the C++ implementation: misuse of a stale /
post-epilogue handle aborts with a diagnostic on stderr. graph()
returns the underlying cudaGraph_t without triggering
cudaGraphInstantiate, enabling callers to embed the nested graph as a
child node into their own graph via cudaGraphAddChildGraphNode while
still getting the lazy dep-A sync on stream() so it can serve as an
event source for cross-stream ordering.
Also add C (Catch2) and Python (pytest) test suites covering:
* relaunch accumulation (N launches accumulate N times)
* zero-launch reusability (pop_prologue immediately followed by
pop_epilogue still unfreezes cleanly)
* exec / stream accessors are non-null between prologue/epilogue
* embedding graph() into an outer graph with
cudaGraphAddChildGraphNode + manual cudaGraphLaunch on the
support stream
* graph-only lifecycle (never calling launch or exec) leaves data
unchanged
Made-with: Cursor
…exts Allow building a stream_ctx (including via stackable_ctx) on a caller stream that is participating in a CUDA graph capture, and ensure that two back-to-back stream_ctx instances constructed on the same caller stream serialize correctly without any explicit sync. - stream_pool: gate cuStreamGetId / cudaStreamGetDevice on cudaStreamIsCapturing (both queries are capture-unsafe and would otherwise invalidate the in-progress capture); use the stream pointer as a fallback stable ID. Add a destructor that closes the pool-owned streams, leaving externally-supplied decorated_stream untouched. - backend_ctx: initialize the CUDA runtime (cudaFree(0)) exactly once per process via std::call_once instead of on every context construction -- cudaFree(0) is rejected and poisons the capture under ThreadLocal/Global capture modes. Record user_provided_handle before the handle is moved so downstream code can distinguish caller-supplied pools from freshly-minted ones. - acquire_release: always merge the context's start events into tasks that have no "input" dependency (no read/rw/reduce/relaxed dep). Previously only truly empty tasks got them, which meant tasks dispatched to pool streams during a capture never issued the cudaStreamWaitEvent forking the pool stream into the capture, so their work was considered uncaptured. - stream_ctx: reject stream_ctx(user_stream, handle) when user_stream is currently capturing and the caller supplied an async_resources_handle -- that handle's pool may carry uncaptured work that cannot legally join the on-going capture. Contexts built without an explicit handle get a fresh empty pool and remain the supported in-capture configuration. Add tests covering both scenarios: cudax/local_stf/legacy_to_stf_in_capture (stream_ctx inside a cudaStreamBeginCapture region, including the fork/join multi-stream property and the negative handle+capture case); cudax/local_stf/stream_ctx_lifetime_btb (back-to-back stream_ctx on the same caller stream, with and without a shared handle); and c/stf/test_stream_ctx_override (C-API equivalent of the back-to-back scenario via stf_ctx_create_ex + has_stream=1, asserting the chaining contract). Made-with: Cursor
pop_prologue() + launchable_graph_scope already cover the manual and the lexical-RAII flavors of a re-launchable stackable pop, but neither lets a caller build a graph, stash it as a data member / in a std::vector / std::unordered_map, and release it later when the last copy dies. Introduce stackable_ctx::launchable_graph, a copyable / movable handle whose shared internal state holds both a stackable_ctx (to keep the impl alive) and the existing launchable_graph_handle. When the last shared copy is destroyed, the state dtor runs ctx.pop_epilogue() - unless the user already did so manually, in which case handle.valid() is false and the dtor is a no-op (no double-epilogue). Factory: stackable_ctx::pop_prologue_shared() wraps pop_prologue() into a launchable_graph; exec()/stream()/graph() forward to the underlying handle so lazy instantiation / dep-A sync semantics are preserved. Includes unit tests for: basic build+launch+implicit release, shared copies driving the same graph (one reset, the other still launches), storage in std::vector across function boundaries, and tolerating a manual pop_epilogue() while shared copies are still alive. Made-with: Cursor
Mirror the C++ stackable_ctx::launchable_graph shared-ownership flavor
behind an opaque stf_launchable_graph_shared:
- stf_stackable_pop_prologue_shared / _dup / _free (last _free fires
pop_epilogue automatically),
- _launch / _exec / _stream / _graph accessors,
- _valid for NULL-safe introspection.
Internally each opaque wraps one heap-allocated launchable_graph value;
_dup copy-constructs a new one, bumping the shared_ptr inside. _free
delete's the wrapper, dropping one refcount; the epilogue runs on the
last C-side handle.
Also appends a C test exercising dup + alternating launches + staggered
free, plus a smoke test that NULL handles are accepted by _free and
_valid.
Made-with: Cursor
Expose the new C stf_launchable_graph_shared surface as a Cython cdef
class LaunchableGraph that holds exactly one C-level shared reference:
Python refcounting drives release (__dealloc__ calls _free), and the
last live Python copy fires pop_epilogue through the C layer.
Adds:
- LaunchableGraph with launch / reset / __enter__/__exit__ /
exec_graph / stream / graph / valid,
- stackable_context.pop_prologue_shared() factory,
- stackable_context.push() / pop() as the decoupled counterparts of
graph_scope() - required so a LaunchableGraph can outlive the
lexical scope that created it.
Tests cover: basic build+launch+drop, stashing the handle in a list
that outlives its producing function, idempotent reset() plus raising
accessors after reset, and the with-block shorthand.
Made-with: Cursor
wp_stf.task(...) is the Warp counterpart of pytorch_task(...). It caches per-raw-stream wp.Stream handles (re-registering the same cudaStream_t with Warp corrupts its bookkeeping), pushes a wp.ScopedStream so default-stream allocations and launches inside the body land on the task stream, auto-detects whether the task's stream is part of an active CUDA graph capture not already tracked by Warp and opens wp.capture_begin(external=True) / wp.capture_end only when needed, and exposes each non-token dep as a zero-copy wp.array view. test_warp_pytorch_dag.py covers three properties of mixing wp_stf.task and pytorch_task in one cuda.stf DAG: - a sequential round-trip Warp -> PyTorch -> Warp on a shared logical_data, - two siblings on disjoint write sets (one Warp, one PyTorch) joined by a third Warp task, - the mixed pipeline captured into a stackable_context and replayed. Made-with: Cursor
stf_stackable_token() returns a stf_logical_data_handle that the rest of
the C API was treating as stackable_logical_data<slice<char>>* via the
old from_opaque_sld(). Tokens are stackable_logical_data<void_interface>,
so every entry point that called sld->get_ld() / validate_access() /
push() through that bare cast was reading the wrong concrete type. When
a task dep was bound to a token inside a stackable push() / pop_prologue
scope, downstream freeze machinery saw void_interface claims layered on
top of an mdspan<char,...,layout_stride> payload (or vice versa) and
aborted with "Data interface type mismatch."
Replace the bare pointee with a discriminated wrapper:
struct stackable_ld_opaque {
bool is_token;
void* impl; // stackable_ld_t* or stackable_token_t*
};
and route every entry point through visit_sld(handle, generic-lambda),
which dispatches on is_token and forwards the concrete
stackable_logical_data<T>& so get_ld() / validate_access() / push() /
set_symbol() / set_read_only() all see the right T. The five producers
(stf_stackable_logical_data{,_with_place,_empty,_no_export_empty},
stf_stackable_token) allocate the inner stackable_*_t and the wrapper
with the unique_ptr/release pattern so a stf_try_allocate failure on
the wrapper does not leak the inner. _destroy and _token_destroy now
assert the wrapper kind so a mismatched destroy entry point is caught
instead of silently UB-deleting through the wrong type.
Adds two reproducers that previously aborted:
- c/experimental/stf/test/test_stackable_token_push.cu:
push_graph + token-only task (with and without pop_prologue,
shared and non-shared) through the C facade.
- cudax/test/stf/local_stf/stackable_token_push_prologue.cu:
same shapes through the native cudax::stf C++ API to confirm the
bug surface is at the C boundary, not in stackable_ctx itself.
Made-with: Cursor
A task wrapper owns per-instance in-flight state -- the capture stream, the captured frontier / done_nodes vector, the held graph_mutex during stream capture, and bookkeeping like must_destroy_child_graph -- on top of the pimpl `task` base. Copying that wrapper had no meaningful semantics: two copies would race on the same in-flight state and double-decrement task counts on destruction. The two surviving "copyable graph_task<>" / "copyable stream_task" UNITTESTs only ever exercised default construction + immediate copy, never an actual copy of a started task, so no real caller relies on this. Mark both as move-only (`= delete` on the copy ctor / copy-assignment, explicit `= default` on the move ctor / move-assignment) and drop the two now-stale UNITTESTs in graph_ctx.cuh and stream_ctx.cuh. Contexts (`stream_ctx`, `graph_ctx`, `context`) are pimpl handles and remain copyable; only the per-task wrapper is restricted. Made-with: Cursor
Pure output of the project's pre-commit hooks on the working tree:
- clang-format: parameter-comment alignment, long-signature wrapping,
include-order fixup, removal of blank lines after namespace {
in c/experimental/stf/{include,test} and in cudax/{include,test}.
- ruff-format / isort: dict literal multi-line wrapping, long argument
list reflow, import ordering in
python/cuda_cccl_experimental/tests/stf/*.py.
- ruff F401: drop a redundant inline `import ctypes` inside
`cai_to_numpy` in example_cholesky.py / example_potri.py (the
top-level `import ctypes` already in scope is the one actually used
a few lines below for `ctypes.c_byte`).
No behavioral change. Pre-commit is idempotent on the result.
Made-with: Cursor
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change adds structured CUDA placement and geometry-aware allocation APIs, external CUDA-context execution places, a standalone Linux-only ChangesStructured placement and allocation
Python package and CI
Validation and documentation
Possibly related PRs
Suggested reviewers: ✨ Finishing Touches⚔️ Resolve merge conflicts
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
🟠 Major comments (19)
cudax/include/cuda/experimental/__places/localized_array.cuh-95-104 (1)
95-104: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: derive the default granularity from the active allocation context rather than device 0.
localized_array::initusescuCtxGetDevice()at Line 347, so evaluation under another device context can report different blocks and owners than the eventual allocation.Also applies to: 563-566
cudax/include/cuda/experimental/__places/localized_array.cuh-162-164 (1)
162-164: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: sample only the valid elements in the final partial block. Clamping every out-of-range probe to
total_elems - 1gives the last element all unused-tail probability, which can select the wrong block owner in both evaluation and allocation.Also applies to: 381-382, 578-589
cudax/test/places/placement.cu-191-209 (1)
191-209: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: keep these general shaped-allocation tests on repeated device 0. Passing
ndevsdistributes blocks onto device 1, but both tests touch the entire range from device 0 without checking device-1 VMM support or peer access; valid multi-GPU systems can therefore fail before the guarded residency test.Also applies to: 223-243
cudax/include/cuda/experimental/__places/exec/cuda_context.cuh-63-65 (1)
63-65: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winimportant: A nonnegative
devidbypasses validation, allowing a device-1 context to advertise and allocate throughdata_place::device(0). Always derive the actual device, reject explicit mismatches, and restore the pushed context with RAII ifcuda_trythrows.Also applies to: 163-175
python/cuda_stf/tests/stf/examples/burger.py-323-349 (1)
323-349: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winimportant: Both Newton solvers make their continuation decision using a residual computed before updating the solution.
python/cuda_stf/tests/stf/examples/burger.py#L323-L349: recompute the residual norm afterlUis updated before writinglnewton_cond.python/cuda_stf/tests/stf/examples/burger_reference.py#L169-L177: check convergence before the solve or recomputenorm2after updatingtU.python/cuda_stf/tests/stf/examples/cholesky.py-744-789 (1)
744-789: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: Synchronize timing with STF completion. Both events use CuPy’s current stream, while factorization runs on STF task streams, and
stop_event.record()occurs beforectx.finalize(). The reported GFLOPS therefore measures submission/default-stream time rather than factorization.python/cuda_stf/tests/stf/examples/neural_ode_dopri5.py-186-190 (1)
186-190: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: operator precedence bug in step-size growth factor.
**binds tighter than*, so.clamp_max(10.0)applies to1/err_normbefore** 0.2, not after. That caps the achievable factor at0.9 * 10**0.2 ≈ 1.43, never approaching the outer.clamp(0.2, 10.0)upper bound the code clearly intends. The adaptive solver will take far more steps than necessary once error is small.- factor = (safety * (1.0 / err_norm.clamp(min=1e-20)).clamp_max(10.0) ** 0.2).clamp( - 0.2, 10.0 - ) + factor = ( + safety * (1.0 / err_norm.clamp(min=1e-20)) ** 0.2 + ).clamp(0.2, 10.0)python/cuda_stf/tests/stf/examples/stackable_branch_while_warp.py-151-163 (1)
151-163: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant:
cudart.cudaGraphDebugDotPrint(...)returns a 1-tuple(err,), soerr != cudart.cudaError_t.cudaSuccessis always true andint(err)will raiseTypeError. Unpack the status before checking it.python/cuda_stf/tests/stf/interop/test_fdtd.py-97-100 (1)
97-100: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: Use the corresponding spacing for each curl derivative. Each stencil divides two derivatives along different axes by one spacing, so anisotropic grids produce incorrect fields; this currently remains hidden because the defaults set
dx == dy == dz.Also applies to: 112-115, 127-130, 142-145, 157-160, 172-175
python/cuda_stf/tests/stf/interop/test_fdtd.py-376-390 (1)
376-390: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftimportant: Add a numerical reference assertion. A no-op stencil, incorrect curl signs, or incorrect axis spacing can still leave every array finite and pass this test.
python/cuda_stf/tests/stf/interop/test_jacobi_numba.py-117-119 (1)
117-119: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: Both Jacobi tests can pass without verifying that Jacobi iteration performed a correct update.
python/cuda_stf/tests/stf/interop/test_jacobi_numba.py#L117-L119: assert the final residual and representative solution values.python/cuda_stf/tests/stf/interop/test_jacobi_pytorch.py#L62-L65: add equivalent convergence and output assertions.python/cuda_stf/tests/stf/interop/test_jacobi_numba.py-69-76 (1)
69-76: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: add a CUDA 12.4+ runtime skip before the conditional-graph paths. The bindings shim only filters on CUDA major, so 12.0–12.3 environments still import these tests and then hit
while_loop/repeat/stackable_context:
python/cuda_stf/tests/stf/interop/test_jacobi_numba.pypython/cuda_stf/tests/stf/interop/test_jacobi_pytorch.pypython/cuda_stf/tests/stf/interop/test_fdtd.pypython/cuda_stf/tests/stf/interop/test_jacobi_warp.py-196-196 (1)
196-196: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: Gate both while_loop tests on CUDA 12.4+ conditional-graph support.
test_jacobi_stackable_warp()andtest_launchable_graph_k_branches_then_while_warp()callctx.while_loop()unconditionally, so older drivers will fail instead of skip.python/cuda_stf/tests/stf/interop/test_scoped_capture.py-147-161 (1)
147-161: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: Cache
wp.Streamwrappers by(device, raw_ptr)here, as intest_jacobi_warp.py. STF can reuse task stream handles, and creating a new wrapper for eacht.stream_ptr()can double-register the same CUDA stream in Warp.python/cuda_stf/tests/stf/test_from_context.py-46-67 (1)
46-67: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: Destroy
resourcesandplacebefore releasing the retained primary context. Both remain live at Line 67, so their later destructors may clean up streams/place state after the underlying non-owning CUDA context has been released. Clear them infinallybeforecuDevicePrimaryCtxRelease(dev).python/cuda_stf/cuda/stf/_experimental/green_places.py-73-108 (1)
73-108: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant:
dev.set_current()leaves the caller’s current CUDA context ondevice_id. Save the previous current device/context and restore it infinally; otherwise later CUDA work on the same thread can run on the wrong device.python/cuda_stf/merge_cuda_wheels.py-119-131 (1)
119-131: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winimportant: reject wheels missing their declared CUDA payload.
A wheel whose filename declares
cu12orcu13is silently skipped when that directory is absent, allowing an incomplete “multi-CUDA” wheel to be produced. Validate every input, including the base wheel, and raise whencuda/stf/_experimental/cu<version>is missing.python/cuda_stf/merge_cuda_wheels.py-152-159 (1)
152-159: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winimportant: select only the wheel created by this invocation.
output_dir.glob("*.whl")[0]may return a stale pre-existing wheel, causing the script to report or publish the wrong artifact. Snapshot the directory before packing, use a clean temporary destination, or require exactly one newly created wheel.ci/build_cuda_stf_python.sh-59-59 (1)
59-59: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winimportant: Line 59 preserves stale wheel outputs, and Lines 108-109 select an arbitrary matching file. A retry or reused workspace can merge an old CUDA wheel into the newly published artifact. Clear
wheelhouse,wheelhouse_merged, andwheelhouse_finalbefore building, then require exactly one wheel per CUDA major.-mkdir -p wheelhouse +rm -rf wheelhouse wheelhouse_merged wheelhouse_final +mkdir -p wheelhouseAs per path instructions, focus on cache/artifact handling and clear failures.
Also applies to: 108-109
Source: Path instructions
🟡 Minor comments (8)
AGENTS.md-210-210 (1)
210-210: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winimportant: Document the standalone package installation in this section. The module list now says
cuda.stf._experimentalships incuda-stf, but the PyPI command immediately below installs onlycuda-cccl; following this guidance will not install the documented STF module. Add thecuda-stf[cu12]/cuda-stf[cu13]command or explicitly separate the package installation instructions.docs/python/stf.rst-126-129 (1)
126-129: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Make this interop snippet executable as documented. Lines 128-129 reference
cuda.compute,OpKind, andNwithout imports or definitions, so the example cannot be copied and run as written. Add the missing setup or label the block as abbreviated.Source: Path instructions
cudax/examples/stf/partitioned_axpy.cu-32-35 (1)
32-35: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winimportant: include
<algorithm>and<vector>directly;std::minandstd::vectorshould not rely on transitive includes fromcuda/experimental/stf.cuh.Source: Coding guidelines
python/cuda_stf/tests/test_examples.py-71-89 (1)
71-89: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winsuggestion: Both paths convert unsupported execution into a passing test instead of a skipped test.
python/cuda_stf/tests/test_examples.py#L71-L89: callpytest.skip(...)for successful skip exits and missing optional dependencies.python/cuda_stf/tests/stf/examples/burger_reference.py#L314-L318: callpytest.skip(...)when no CUDA device is available.python/cuda_stf/tests/stf/test_stream_utils.py-76-87 (1)
76-87: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winsuggestion: Gate the device creation here with the existing no-usable-device skip pattern;
pytest.importorskip(...)only covers imports, notcore.Device()on GPU-less machines.python/cuda_stf/tests/stf/test_nested_scopes.py-195-196 (1)
195-196: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Tighten these tolerances below one loop increment. The current thresholds accept exactly the regressions these tests should detect:
1.8passes the2.0 ± 0.2check,2.75passes3.0 ± 0.35, and1.5passes2.0 ± 0.6. Use a numerical-noise tolerance such as1e-6, or assert the executed iteration count separately.Also applies to: 391-392, 468-470, 525-525, 574-576
python/cuda_stf/tests/stf/test_place_support.py-184-197 (1)
184-197: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Synchronize
streamdirectly instead of the whole Numba context. Context-wide synchronization also waits for work mistakenly submitted to another stream, so this test cannot detect a brokencuda.computestream handoff. Use thecudaStreamSynchronizepattern already used at Lines 312-316.python/cuda_stf/cuda/stf/_experimental/device_array.py-114-126 (1)
114-126: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winimportant: Empty arrays skip allocate() but still register a finalizer that calls deallocate(0, 0, ...). Allocator implementations that track allocations may reject this deallocation and emit spurious warnings during garbage collection. Avoid registering the finalizer when
_nbytes == 0(e.g., checkif self._nbytes > 0:before theweakref.finalize()call at line 119).
🧹 Nitpick comments (5)
docs/python/stf.rst (1)
3-4: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion: Use one canonical CUDASTF expansion across the documentation. The two changed locations currently disagree.
docs/python/stf.rst#L3-L4: update the title and introduction to the confirmed canonical expansion.AGENTS.md#L210-L210: use the same expansion in the module guidance.python/cuda_stf/tests/test_examples.py (1)
91-99: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion: Call the discovered
__main__callable directly instead of constructing code forexec. The module name is repository-derived rather than request-controlled, but removingexeceliminates unnecessary dynamic evaluation and the static-analysis failure.Source: Linters/SAST tools
python/cuda_stf/tests/stf/interop/test_jacobi_warp.py (1)
268-270: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winsuggestion: Assert that Jacobi actually converged.
This test only finalizes and prints, so incorrect iteration output or a prematurely terminated loop still passes. Validate the exported residual or compare the final grid with a numerical oracle.
python/cuda_stf/tests/stf/interop/test_numba.py (1)
326-330: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winsuggestion: Vectorize the CPU Laplacian reference.
These million Python-level iterations unnecessarily inflate GPU-test runtime. Compute the interior with NumPy slicing while retaining the explicit boundary copies.
python/cuda_stf/tests/stf/interop/test_warp_pytorch_dag.py (1)
145-188: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftsuggestion: The test does not verify concurrent sibling execution. A scheduler that serializes both tasks would produce the same values and pass. Use stream/event instrumentation to demonstrate overlap, or rename the test and docstring to cover dependency ordering only.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f3a9a0ca-7183-478d-9cd7-c3d94e3acbfa
📒 Files selected for processing (102)
.github/CODEOWNERSAGENTS.mdc/experimental/stf/include/cccl/c/experimental/stf/stf.hc/experimental/stf/src/stf.cuc/experimental/stf/test/test_placement.cppc/experimental/stf/test/test_places.cppci/build_cuda_stf_python.shci/build_cuda_stf_wheel.shci/matrix.yamlci/project_files_and_dependencies.yamlci/test/inspect_changes/c2h_dependency.outputci/test_cuda_stf_python.shci/util/workflow/get_wheel_artifact_name.shcudax/examples/places/thrust_device_data_place_allocator.cucudax/examples/stf/CMakeLists.txtcudax/examples/stf/partitioned_axpy.cucudax/include/cuda/experimental/__places/cute_partition.cuhcudax/include/cuda/experimental/__places/data_place_interface.cuhcudax/include/cuda/experimental/__places/exec/cuda_context.cuhcudax/include/cuda/experimental/__places/exec/green_context.cuhcudax/include/cuda/experimental/__places/localized_array.cuhcudax/include/cuda/experimental/__places/partitions/cyclic_shape.cuhcudax/include/cuda/experimental/__places/places.cuhcudax/include/cuda/experimental/__stf/graph/interfaces/slice.cuhcudax/include/cuda/experimental/__stf/internal/stf_places_extended_exports.cuhcudax/include/cuda/experimental/__stf/stream/interfaces/slice.cuhcudax/include/cuda/experimental/__stf/utility/dimensions.cuhcudax/include/cuda/experimental/places.cuhcudax/include/cuda/experimental/stf.cuhcudax/test/places/CMakeLists.txtcudax/test/places/placement.cudocs/conf.pydocs/cudax/places.rstdocs/python/api_reference.rstdocs/python/index.rstdocs/python/setup.rstdocs/python/stf.rstdocs/python/stf_api.rstpython/cuda_stf/.gitignorepython/cuda_stf/CMakeLists.txtpython/cuda_stf/LICENSEpython/cuda_stf/README.mdpython/cuda_stf/cuda/stf/_experimental/__init__.pypython/cuda_stf/cuda/stf/_experimental/_cuda_version_utils.pypython/cuda_stf/cuda/stf/_experimental/_stf_bindings.pypython/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyxpython/cuda_stf/cuda/stf/_experimental/_stream_utils.pypython/cuda_stf/cuda/stf/_experimental/device_array.pypython/cuda_stf/cuda/stf/_experimental/fill_utils.pypython/cuda_stf/cuda/stf/_experimental/green_places.pypython/cuda_stf/cuda/stf/_experimental/interop/__init__.pypython/cuda_stf/cuda/stf/_experimental/interop/numba.pypython/cuda_stf/cuda/stf/_experimental/interop/pytorch.pypython/cuda_stf/cuda/stf/_experimental/paths.pypython/cuda_stf/cuda/stf/_experimental/task_graph.pypython/cuda_stf/merge_cuda_wheels.pypython/cuda_stf/pyproject.tomlpython/cuda_stf/tests/stf/examples/__init__.pypython/cuda_stf/tests/stf/examples/bicgstab.pypython/cuda_stf/tests/stf/examples/burger.pypython/cuda_stf/tests/stf/examples/burger_reference.pypython/cuda_stf/tests/stf/examples/cg.pypython/cuda_stf/tests/stf/examples/cholesky.pypython/cuda_stf/tests/stf/examples/fhe.pypython/cuda_stf/tests/stf/examples/fhe_decorator.pypython/cuda_stf/tests/stf/examples/neural_ode_dopri5.pypython/cuda_stf/tests/stf/examples/neural_ode_rk4.pypython/cuda_stf/tests/stf/examples/potri.pypython/cuda_stf/tests/stf/examples/stackable_branch_while_warp.pypython/cuda_stf/tests/stf/interop/__init__.pypython/cuda_stf/tests/stf/interop/test_cuda_compute.pypython/cuda_stf/tests/stf/interop/test_decorator.pypython/cuda_stf/tests/stf/interop/test_fdtd.pypython/cuda_stf/tests/stf/interop/test_jacobi_numba.pypython/cuda_stf/tests/stf/interop/test_jacobi_pytorch.pypython/cuda_stf/tests/stf/interop/test_jacobi_warp.pypython/cuda_stf/tests/stf/interop/test_legacy_to_stf.pypython/cuda_stf/tests/stf/interop/test_local_stf_capture.pypython/cuda_stf/tests/stf/interop/test_numba.pypython/cuda_stf/tests/stf/interop/test_pytorch.pypython/cuda_stf/tests/stf/interop/test_pytorch_task_context.pypython/cuda_stf/tests/stf/interop/test_scoped_capture.pypython/cuda_stf/tests/stf/interop/test_stencil_decorator.pypython/cuda_stf/tests/stf/interop/test_warp_pytorch_dag.pypython/cuda_stf/tests/stf/test_cai.pypython/cuda_stf/tests/stf/test_composite_places.pypython/cuda_stf/tests/stf/test_context.pypython/cuda_stf/tests/stf/test_cuda_kernel.pypython/cuda_stf/tests/stf/test_fill_utils.pypython/cuda_stf/tests/stf/test_from_context.pypython/cuda_stf/tests/stf/test_graph_scope.pypython/cuda_stf/tests/stf/test_host_launch.pypython/cuda_stf/tests/stf/test_launchable_graph.pypython/cuda_stf/tests/stf/test_lifecycle.pypython/cuda_stf/tests/stf/test_nested_scopes.pypython/cuda_stf/tests/stf/test_packaging.pypython/cuda_stf/tests/stf/test_place_support.pypython/cuda_stf/tests/stf/test_placement.pypython/cuda_stf/tests/stf/test_stream_utils.pypython/cuda_stf/tests/stf/test_task_graph.pypython/cuda_stf/tests/stf/test_token.pypython/cuda_stf/tests/test_examples.py
👮 Files not reviewed due to content moderation or server errors (3)
- python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx
- python/cuda_stf/cuda/stf/_experimental/_stream_utils.py
- ci/test_cuda_stf_python.sh
| tol_sq = tol * tol | ||
|
|
||
| # --- CG while loop (conditional CUDA graph node) ---------------------- | ||
| with ctx.while_loop() as loop: | ||
| lAp = ctx.logical_data_empty((N,), np.float64, name="Ap") | ||
| lpAp = ctx.logical_data_empty((1,), np.float64, name="pAp") | ||
| lrsnew = ctx.logical_data_empty((1,), np.float64, name="rsnew") | ||
|
|
||
| # Ap = A @ P | ||
| stf_matvec(ctx, lA, lP, lAp) | ||
|
|
||
| # pAp = P'Ap | ||
| stf_dot(ctx, lP, lAp, lpAp) | ||
|
|
||
| # X += alpha·P (alpha = rsold / pAp) | ||
| # Alpha is recomputed in the R update task below because each | ||
| # pytorch_task is an independent graph node with its own closure. | ||
| with pytorch_task(ctx, lX.rw(), lrsold.read(), lpAp.read(), lP.read()) as ( | ||
| tX, | ||
| tRsold, | ||
| tPAp, | ||
| tP, | ||
| ): | ||
| alpha = tRsold.squeeze() / tPAp.squeeze() | ||
| tX += alpha * tP | ||
|
|
||
| # R -= alpha·Ap | ||
| with pytorch_task(ctx, lR.rw(), lrsold.read(), lpAp.read(), lAp.read()) as ( | ||
| tR, | ||
| tRsold, | ||
| tPAp, | ||
| tAp, | ||
| ): | ||
| alpha = tRsold.squeeze() / tPAp.squeeze() | ||
| tR -= alpha * tAp | ||
|
|
||
| # rsnew = R'R | ||
| stf_dot(ctx, lR, lR, lrsnew) | ||
|
|
||
| # Condition: continue while residual norm² exceeds tolerance². | ||
| # This sets the predicate for the *next* replay — the P and rsold | ||
| # updates below still execute in the current iteration. | ||
| loop.continue_while(lrsnew, ">", tol_sq) | ||
|
|
||
| # P = R + beta·P (beta = rsnew / rsold) | ||
| with pytorch_task(ctx, lP.rw(), lR.read(), lrsnew.read(), lrsold.read()) as ( | ||
| tP, | ||
| tR, | ||
| tRsnew, | ||
| tRsold, | ||
| ): | ||
| beta = tRsnew.squeeze() / tRsold.squeeze() | ||
| tP[:] = tR + beta * tP | ||
|
|
||
| # rsold = rsnew | ||
| with pytorch_task(ctx, lrsold.write(), lrsnew.read()) as (tRsold, tRsnew): | ||
| tRsold.copy_(tRsnew) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
critical: Add an iteration cap to this conditional loop. Residual-only termination can replay forever if CG stagnates or produces NaN, hanging the example job. Track an iteration scalar and combine it with the residual predicate as the BiCGSTAB examples do.
| # One handle per process is enough: all nvmath submissions are serialized by | ||
| # the Python GIL, and ``set_stream`` is called at the top of each task body so | ||
| # the asynchronous work lands on ``t.stream_ptr()``. | ||
| _cublas_handle = 0 | ||
| _cusolver_handle = 0 | ||
|
|
||
|
|
||
| def _cublas(): | ||
| global _cublas_handle | ||
| if _cublas_handle == 0: | ||
| _cublas_handle = _cb.create() | ||
| return _cublas_handle | ||
|
|
||
|
|
||
| def _cusolver(): | ||
| global _cusolver_handle | ||
| if _cusolver_handle == 0: | ||
| _cusolver_handle = _cdn.create() | ||
| return _cusolver_handle |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the example first, then read the relevant slices.
ast-grep outline python/cuda_stf/tests/stf/examples/cholesky.py --view expanded || true
printf '\n--- lines 1-120 ---\n'
sed -n '1,120p' python/cuda_stf/tests/stf/examples/cholesky.py
printf '\n--- lines 180-280 ---\n'
sed -n '180,280p' python/cuda_stf/tests/stf/examples/cholesky.py
printf '\n--- lines 700-820 ---\n'
sed -n '700,820p' python/cuda_stf/tests/stf/examples/cholesky.pyRepository: NVIDIA/cccl
Length of output: 12846
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search for handle creation/teardown and stream usage in this subtree.
rg -n "_cublas_handle|_cusolver_handle|set_stream|create\(|destroy\(|stream_ptr\(|ndevs|device" python/cuda_stf/tests/stf/examples/cholesky.py python/cuda_stf/tests/stf/examples -SRepository: NVIDIA/cccl
Length of output: 19342
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('python/cuda_stf/tests/stf/examples/cholesky.py')
text = p.read_text()
for needle in ['_cublas_handle', '_cusolver_handle', 'set_stream', 'stream_ptr', 'ndevs', 'destroy', 'create(']:
print(f'## {needle}')
for i, line in enumerate(text.splitlines(), 1):
if needle in line:
print(f'{i}: {line}')
print()
PYRepository: NVIDIA/cccl
Length of output: 1763
🌐 Web query:
cuBLAS handle CUDA context bound multi-GPU stream setStream current device context documentation
💡 Result:
In the cuBLAS library, the handle (cublasHandle_t) acts as the primary opaque object representing the cuBLAS library context [1][2]. This context is tightly coupled to the specific CUDA context that was current at the time the handle was created via cublasCreate [1][3]. Because of this binding, a cuBLAS handle is associated with a single CUDA device [1][2]. For multi-GPU applications, you must create one unique cuBLAS handle for each device you intend to use [1][3][2]. Before creating a handle for a specific device, ensure that the target device is set as the current device in the thread (e.g., by calling cudaSetDevice) [1][3]. Furthermore, the cuBLAS handle must never outlive the CUDA context with which it is associated [3][2]. To manage parallelism or control execution within a specific CUDA device, you can use cublasSetStream [4][5]. This function associates a user-created CUDA stream (cudaStream_t) with the cuBLAS handle [4][5]. All subsequent cuBLAS calls using that handle will then execute within the specified stream [4][5]. If a stream is not explicitly set, cuBLAS defaults to using the legacy NULL stream [4][5]. When using cublasSetStream, note the following: - The function call replaces the current stream associated with the handle [4][5]. - It unconditionally resets the cuBLAS library workspace to the default workspace pool, which may require you to re-call cublasSetWorkspace if you were using a custom workspace [4][5]. - In multi-threaded applications, the recommended practice is to maintain one handle per thread, ensuring that the thread's current device and stream associations are managed consistently throughout the life of the thread [2].
Citations:
- 1: https://docs.nvidia.com/cuda/cublas/index.html
- 2: https://docs.nvidia.com/cuda/archive/12.2.2/cublas/
- 3: https://docs.nvidia.com/cuda/archive/12.6.2/cublas/
- 4: https://docs.nvidia.com/cuda/archive/11.6.1/pdf/CUBLAS_Library.pdf
- 5: https://docs.nvidia.com/cuda/archive/12.1.1/cublas/index.html
critical: python/cuda_stf/tests/stf/examples/cholesky.py:55-73, 213-239 Cache cuBLAS/cuSOLVER handles per device/context, not once per process. Blocks are scheduled across all ndevs, so a handle created on one device can be reused on another and hit the wrong CUDA context.
| @pytest.mark.parametrize("context_type", [stf.context, stf.stackable_context]) | ||
| def test_logical_data_accepts_explicit_c_contiguous_cai_strides(context_type): | ||
| array = np.arange(8, dtype=np.float32).reshape(2, 4) | ||
| assert array.strides is not None | ||
|
|
||
| ctx = context_type() | ||
| ld = ctx.logical_data(_CudaArrayInterfaceWrapper(array)) | ||
| assert ld.shape == array.shape | ||
| ctx.finalize() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
critical: _CudaArrayInterfaceWrapper(array) exposes the host address from array.ctypes.data as a CUDA pointer. Unlike the rejection tests, this path accepts and may use that pointer during finalize(), potentially causing an invalid-address CUDA failure. Build this case from genuinely device-backed storage and add explicit contiguous strides to its CAI metadata.
Without replication, an unbound grid axis pins every element to coordinate 0 on that axis and leaves the other places idle. Fail at construction time so callers must collapse the grid or bind every axis to a tensor dimension.
Keep documentation and implementation comments focused on the package's current behavior and requirements.
Address the CUDASTF Python-bindings review: retain source buffers and reject writes to read-only imports, validate CAI layout, propagate place/data-place ownership so external CUDA resources outlive their wrappers, capture host-callback exceptions and re-raise them from blocking wait()/finalize()/check_errors(), guard stackable finalization against open scopes, drop the CuPy-only fill path in favor of the driver API, and harden the wheel merge. Export task-argument CUDA Array Interface views with stream=None: inside a task the caller launches on the task stream(s), which STF has already ordered behind the data's producers, so the CAI stream handoff is redundant and an integer stream would force an illegal-during-capture host sync in consumers like Numba. This also makes a single view valid for multi-place grids. Documentation and tests updated accordingly.
…ontexts A caller-stream context finalizes asynchronously, so its queued resource release runs later on the caller's stream. Dropping the async_resources keep-alive at finalize() could destroy a temporary handle=async_resources() before that work runs. Retain it on the context object until __dealloc__ for caller-stream contexts (default contexts still release at finalize(), which blocks); the caller is contractually required to synchronize their stream before releasing or reusing resources.
… stream STF does not yet order imported data behind an external producer stream, so a non-None CUDA Array Interface stream field is a potential race: the first task could read the buffer on another stream before the producer's work completes. Rather than silently ignore it, reject such inputs with a clear NotImplementedError in both the regular and stackable logical_data import paths, pointing the caller to synchronize the producer stream before registering. This does not fire for common producers (PyTorch exports CAI v2 without a stream; NumPy/CuPy/Numba default to stream=None).
… 8-byte fill The CI run flagged five stale tests that still asserted pre-change behavior: - test_task_arg_cai_v3 expected the task-argument CAI view to advertise the task stream; it now advertises stream=None (STF orders the task stream behind producers, and an integer stream would make Numba host-synchronize). - test_logical_data_rejects_non_contiguous matched "not contiguous"; the message is now "not C-contiguous". - The fill_utils tests monkeypatched the removed _fill_8byte_cupy helper; nonzero 8-byte fills now go through the CUDA-driver strided memset (_fill_8byte_driver), so the tests assert that path instead.
# Conflicts: # c/experimental/stf/src/stf.cu # c/experimental/stf/test/test_places.cpp # ci/matrix.yaml # cudax/examples/stf/partitioned_axpy.cu # cudax/include/cuda/experimental/__places/cute_partition.cuh # cudax/include/cuda/experimental/__places/exec/cuda_context.cuh # cudax/include/cuda/experimental/__places/localized_array.cuh # cudax/include/cuda/experimental/__stf/graph/interfaces/slice.cuh # cudax/include/cuda/experimental/__stf/internal/stf_places_extended_exports.cuh # cudax/include/cuda/experimental/__stf/stream/interfaces/slice.cuh # cudax/test/places/placement.cu # docs/cudax/places.rst # docs/python/setup.rst
The C++ runtime partition class merged through NVIDIA#9804/NVIDIA#9808 as cute_partition_descriptor (cute_partition now names the typed template), and the runtime factory as make_partition_descriptor. Constructor and factory signatures are unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Python shapes, per-dimension specifications, callback coordinates, and grid axes are now C/row-major, matching NumPy: no caller ever reverses a shape. The C ABI and C++ implementation remain dimension-0-fastest; the conversion is private to the binding and rank-aware (tensor and grid ranks are stored on wrapper objects, never inferred by trimming extent-1 dimensions). - cute_partition.from_spec/from_leaves reverse dims, specs, and leaf order, and remap grid axes; dims accessors return rank-aware C-order tuples; leaves are public last-leaf-fastest. - Add cute_partition.grid_place_offset(i): the offset of the place at a C-order grid-linear index, which differs from place-mode place_offset(i) when tensor dimensions map to grid axes in swapped order. - Python mappers see C-order tuples of explicit rank and return C-order grid coordinates; data_place.composite and exec_place_grid.create require data_rank for Python callables (shape-free callbacks cannot infer rank). partition_fn_blocked takes a public axis. - data_place.allocate extents, placement_evaluate data_dims, exec_place(.grid) dims, and task.get_grid_dims follow the same contract. - Extend the placement tests with non-square shapes, per-axis blocking, swapped tensor/grid axes (pinning grid_place_offset vs place_offset), rank preservation with extent-1 dims, public round trips, uneven padding, and the tensor-of-tiles pattern (spec with trailing whole payload dims), including its shaped allocation. - Document the C-order contract and the tensor-of-tiles construction. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Harden the Cython bindings and interop, correct numerical examples, and apply review-driven test/doc cleanup: - Bound the Dopri5/CG adaptive loops with device-side caps and finite checks; reset persistent RK4 state per invocation and key its warmup on step size. - Make the Jacobi (numba/pytorch/warp) tests assert finite residuals below tolerance, unchanged boundaries, and an evolved interior. - FHE examples support arbitrary lengths, preserve keys in empty_like(), and reject mismatched contexts/lengths/keys. - Fix Cholesky timing to synchronize STF work before measuring, and move the STF-binding skip before optional CuPy/nvmath imports in Cholesky/POTRI. - Vectorize the stencil reference, compile the shared AXPY kernel once, and replace dynamic exec in the example runner with a direct call. - Correct "Stream Task Flow" -> "Sequential Task Flow" references, drop the untested Python 3.14 classifier, and quote STF pip extras in the docs.
# Conflicts: # c/experimental/stf/src/stf.cu # docs/python/setup.rst # python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx
- Export the structured-partition symbols (cute_partition, partition_fn_blocked/cyclic, placement_evaluate, placement_stats) through the explicit _BINDING_EXPORTS allowlist introduced on stf_c_api. - Port the new composite-mapper tests to the C-order contract (data_rank keyword, grid-rank-sized mapper results). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Store the public shape (int or C-order tuple, like a NumPy shape) instead of exposing every allocation as one-dimensional: - Emit the shape with compact C-contiguous strides through __cuda_array_interface__; copy_to_host returns a shaped array and from_host preserves the source shape. - Add contiguous reshape() as a non-owning view (with NumPy-style -1 inference) that keeps the owning allocation alive; slicing stays restricted to 1-D arrays. - Default the allocation geometry to the shape with the dtype's item size, so a shaped allocation on a partitioned composite place needs no extra argument; dims=/elemsize= remain as an explicit override. Byte counts passed to allocate/deallocate are unchanged (plain places forward allocate_nd to the byte path). - Cover the rank-4 tensor-of-tiles allocation consumed through a rank-2 collapse, shaped CAI round trips, and reshape view lifetime. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
green_places() may be the first CUDA call in the process, in which case cuCtxGetCurrent reports CUDA_ERROR_NOT_INITIALIZED. Call the idempotent cuInit first so no explicit setup is required of the caller. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The C++ core's cute_partition owner(pos4) was reachable only from C++. Add stf_cute_partition_owner() (native dimension-0-fastest order) and cute_partition.owner(coords) in Python (C-order coordinates in, C-order grid coordinates out, ranks validated), so adapters can reason about element placement without re-implementing any policy. Tests pin the 1-D blocked mapping, the swapped tensor/grid-axis case, rank validation, and the tensor-of-tiles property that ownership depends only on the tile coordinates; the docs show that property via owner() and note that physical placement is page-granular (placement_evaluate quantifies the difference). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
wp.CaptureMode members are upper-case (RELAXED); the test was written against an older spelling and only runs when warp-lang is installed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
3 participants
Description
Python/C surface for the placement facilities of the STF places layer.
Stacked on #5315 and #9804 — their commits appear in this diff until they land on
main; review the placement-specific commits:[STF] C API for placement evaluation and structured partitions[STF] Python: cute_partition, placement_evaluate, shaped allocation[STF] Follow the allocate_nd rename in the C API and Python bindings[STF] Harden the placement C API and Python bindings (review findings)[STF] Allocate DeviceArray through the extents formThis does not include the
parallel_forintegration from #9808; that remains a separate stacked PR.C API (
c/experimental/stf/include/cccl/c/experimental/stf/stf.h):stf_data_place_allocate_nd()— geometry-aware allocation on any place handle.stf_placement_evaluate()/stf_placement_evaluate_partition()— score a candidate mapping without allocating:stf_placement_statsplus optional per-grid-position byte counts.stf_cute_partition_*— build structured partitions from a JAX-like per-dimension spec or raw leaves, with accessors (dims, leaves, per-place offsets);stf_data_place_composite_cute()wraps one as a composite data place.stf_partition_fn_blocked()/stf_partition_fn_cyclic()— native partition functions usable wherever a mapper is expected, avoiding FFI callback cost for the common policies.Python (
cuda.stf._experimental):cute_partition.from_spec((nx, ny, nz), (None, ("blocked", 0), None), (nplaces,))andfrom_leaves, with accessor properties.placement_evaluate(grid, mapper, dims, elemsize)where the mapper is acute_partition, a native function pointer, or a Python callable (documented as the slow path — each probe crosses the GIL); returnsplacement_statswithaccuracyand per-position bytes.data_place.allocate()accepts extents with keyword-onlyelemsizein addition to a byte count (existing callers unaffected —streamstays the second positional);composite()additionally accepts native function pointers;composite_cute()wraps a structured partition.Hardened by adversarial review with every crash reproduced first: NULL-handle accessors,
bool-as-function-pointer,elemsize=0SIGFPE, an out-of-range mapper corrupting the caller's stats buffer, and raising Python mappers being silently swallowed — all now clean Python exceptions, each pinned by a pytest. The multi-GPU residency test queries the device's real allocation granularity and self-skips below two devices.Checklist