[STF] Add C and Python execution-grid reshape bindings#9899
[STF] Add C and Python execution-grid reshape bindings#9899caugonnet wants to merge 660 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>
…nCaptureToGraph On CTK 12.4+, conditional graph nodes are not allowed inside a child graph that is later embedded into another graph. The legacy capture path used by graph_task captured into a fresh per-task graph and embedded it as a child node of ctx_graph, which made wp.capture_while inside an STF task abort with cudaErrorNotSupported in end_uncleared(). Use cudaStreamBeginCaptureToGraph (available since CTK 12.3) to capture directly into ctx_graph so conditional nodes land at the top level. The shared ctx_graph is mutated for the full capture interval, so graph_mutex must be held throughout; ownership is handed from start() into the task-owned capture_lock_ member and moved back into a scope-bound lock in end_uncleared() so RAII releases it even if user code throws between the two calls. End-of-capture snapshots the frontier and emits a single empty done node, keeping the done_nodes invariant intact and avoiding fragile empty/single-tail detection. The legacy CTK 12.0-12.2 path is preserved unchanged; conditional graph nodes are unreachable on those toolkits anyway (Warp stubs them out below CTK 12.4), so only ordinary STF capture has to keep working there. Co-authored-by: Cursor <cursoragent@cursor.com>
`task::acquire` only merged the context's start events into the task's prereqs when no dependency had a read access mode (`!has_input_dep`). That gate is too narrow: a task whose only dep is a `.read()` (or `.rw()`) on a logical data with no previous writer in this context -- e.g. an in-capture `token().read()` of a token nobody has written -- sets `has_input_dep == true`, but `enforce_stf_deps_before` returns an empty event list and `fetch_data` adds no MSI/allocation prereqs for void-interface (token) data, so the resulting prereq list is still empty. When the context is bound to a user stream that is participating in a CUDA graph capture, the pool stream picked for that task therefore issues no `cudaStreamWaitEvent` and never gets forked into the on-going capture. Later, `ctx.finalize() -> fence() -> join_with_stream` issues `cudaStreamWaitEvent(user_stream /*captured*/, event /*recorded on uncaptured pool stream*/)` and CUDA aborts with `cudaErrorStreamCaptureIsolation`. Replace the `!has_input_dep` gate with the actual invariant the existing comment was trying to express: merge the start events when `result` is still empty after the dependency loop. This covers both the original write-only case and the read-with-no-prior-writer case in the same condition. Because `start_event` was recorded on the user stream (in capture), the implicit `cudaStreamWaitEvent(pool_stream, start_event)` issued at task setup forks the pool stream into the same capture, so the finalize-time wait on its events is a normal intra-capture sync. Drop the now-dead `has_input_dep` tracking. Co-authored-by: Cursor <cursoragent@cursor.com>
When a stackable graph_scope is popped, graph_ctx_node::finalize_after_launch records an event behind support_stream after the body cudaGraphLaunch and returns it as finalize_prereqs. _pop_epilogue then propagates that event exclusively through pushed_data via pop_after_finalize. For a graph_scope used purely for token ordering (no logical_data pushed into it), pushed_data is empty and finalize_prereqs is dropped on the floor. The parent stream_ctx's submitted_stream therefore has no order relationship with support_stream, and the eventual cudaStreamSynchronize in stream_ctx::finalize returns while the body graph is still running. This manifests as host reads of wp.array values written from inside the body returning stale data after the with-context exits. Forward finalize_prereqs into the parent context's dangling_events so that the next insert_fence (called by fence() / submit() and therefore finalize()) picks them up. The pushed_data path is unchanged. Co-authored-by: Cursor <cursoragent@cursor.com>
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesCUDASTF package and grid APIs
Suggested reviewers: 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: 2
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
🟠 Major comments (27)
cudax/include/cuda/experimental/__places/places.cuh-1415-1421 (1)
1415-1421: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: reject dimension-product overflow before comparing against the place count.
Unsigned multiplication in
dims.size()can wrap; for example,{UINT64_MAX, UINT64_MAX, 2, 1}produces2modulo 2⁶⁴ and can incorrectly validate a two-place grid. Use checked multiplication and add C/C++ regression coverage for overflowed shapes.Also applies to: 1520-1526
python/cuda_stf/tests/test_examples.py-143-166 (1)
143-166: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winimportant: Move this STF test runner under
python/cuda_stf/tests/stf.Relocate it to
python/cuda_stf/tests/stf/test_examples.pyand adjusttests_dirso module discovery still producesstf.examples.*. As per path instructions, “Keep Stream Task Flow tests underpython/cuda_stf/tests/stf.”Source: Path instructions
python/cuda_stf/tests/stf/examples/fhe.py-91-112 (1)
91-112: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: Define and enforce ciphertext size in both FHE implementations.
Both variants permit out-of-bounds reads for unequal operands and process at most 512 elements.
python/cuda_stf/tests/stf/examples/fhe.py#L91-L112: preserve element count throughempty_like(), validate equal lengths, and derive the Numba launch grid from the count.python/cuda_stf/tests/stf/examples/fhe_decorator.py#L72-L83: apply the same size contract and dynamic launch calculation to the decorator path.python/cuda_stf/merge_cuda_wheels.py-119-131 (1)
119-131: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winimportant: Fail when any declared CUDA payload is missing.
Validate every wheel—including the base—contains
cuda/stf/_experimental/cu<version>. Silently continuing produces a wheel that omits an advertised CUDA backend while reporting success.python/cuda_stf/tests/stf/examples/neural_ode_rk4.py-305-338 (1)
305-338: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: Reset
l_ybefore every persistent forward.Warmup and timed calls currently continue from the preceding trajectory, whereas the eager path clones
y0each time. Stage the initial state before enteringgraph_scope()so every sample measures the same integration.Also applies to: 380-398
python/cuda_stf/merge_cuda_wheels.py-136-159 (1)
136-159: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: Return the wheel created by
wheel pack, notoutput_dir.glob("*.whl")[0].output_diris reused, so the glob can pick a stale wheel; use the expectedbase_wheel_namepath or diff the directory contents and require exactly one new artifact.ci/build_cuda_stf_python.sh-164-169 (1)
164-169: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: Inconsistent path resolution for the artifact-name/upload calls.
Every other cross-script call in this file uses
$ci_dir(lines 9, 20, 101), and the siblingci/test_cuda_stf_python.shcorrectly does the same for the identical calls ("$ci_dir/util/workflow/get_wheel_artifact_name.sh","$ci_dir/util/artifacts/download.sh"). Here,ci/util/workflow/get_wheel_artifact_name.shandci/util/artifacts/upload.share invoked with plain relative paths, which only work if the script happens to be run from the repo root.Fix: use $ci_dir consistently
- wheel_artifact_name="$(CCCL_WHEEL_KIND=stf ci/util/workflow/get_wheel_artifact_name.sh)" - ci/util/artifacts/upload.sh "$wheel_artifact_name" 'wheelhouse/.*' + wheel_artifact_name="$(CCCL_WHEEL_KIND=stf "$ci_dir/util/workflow/get_wheel_artifact_name.sh")" + "$ci_dir/util/artifacts/upload.sh" "$wheel_artifact_name" 'wheelhouse/.*'ci/build_cuda_stf_python.sh-39-53 (1)
39-53: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant:
cuda13_version=13.1.1is out of sync withci/matrix.yaml’s13.XCTK alias, which maps to13.3. Bump the CUDA 13 image tag here so the STF build uses the same toolkit line as thetest_py_stflanes.python/cuda_stf/tests/stf/test_place_support.py-171-174 (1)
171-174: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: Guard the optional Numba dependency.
A valid environment with CUDA Compute but without Numba fails this test at import time. Use
pytest.importorskip("numba.cuda"), as the other Numba-dependent tests do.python/cuda_stf/tests/stf/test_place_support.py-13-19 (1)
13-19: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: catch only the documented unsupported-platform condition in
_require_green_context_helper; swallowing every exception turns binding bugs and bad results into skips. Also guard thenumba.cudaimport intest_scope_with_cuda_computewithpytest.importorskip("numba")/pytest.importorskip("numba.cuda")so missing Numba skips cleanly instead of failing the test.python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx-476-490 (1)
476-490: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winimportant: validate
shapebefore returning for implicit strides. A CAI producer can supply negative extents withstrides=None; both context variants then compute_lenfrom unchecked dimensions instead of producing the intendedValueError. Move shape normalization and extent validation above the early return and add regression coverage.python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx-1904-1924 (1)
1904-1924: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftimportant: preserve Python callback failures across the
noexceptFFI boundary. An exception fromfn(...)becomes unraisable, allowing synchronization or finalization to report success after host work failed. Capture the exception in context-owned state and re-raise it at the next synchronization point.python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx-3302-3317 (1)
3302-3317: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: apply the same shape validation as
logical_data.init_by_shape. This path treats an empty shape as one element and relies on C conversions for non-positive or non-integral extents, producing inconsistent sizes and exceptions between context types.python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx-710-725 (1)
710-725: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winimportant: both buffer paths request
PyBUF_ANY_CONTIGUOUSbut discard strides and expose the bytes with C-order shape metadata. Fortran-contiguous arrays are therefore accepted and interpreted in the wrong order.
python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx#L710-L725: requirePyBUF_C_CONTIGUOUS.python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx#L3268-L3279: apply the same restriction to stackable data.python/cuda_stf/tests/stf/test_context.py#L95-L105: add a Fortran-contiguous rejection test for both context types.python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx-1373-1380 (1)
1373-1380: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: validate
grid_dimsbefore indexing or conversion. Empty dimensions currently raiseIndexError, extents beyond four are silently ignored, and fractional values are truncated byint().
python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx#L1373-L1380: require one to four positive integral extents.python/cuda_stf/tests/stf/test_composite_places.py#L89-L103: cover empty, oversized, non-integral, zero, and negative shapes.python/cuda_stf/tests/stf/test_nested_scopes.py-195-196 (1)
195-196: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: these tolerances equal or exceed one full iteration's state change, so off-by-one repeat/while execution can still pass. Use tolerances below half an iteration increment and assert residual or target state where available.
python/cuda_stf/tests/stf/test_nested_scopes.py#L195-L196: tighten below the scaled0.2step.python/cuda_stf/tests/stf/test_nested_scopes.py#L391-L392: tighten below the0.25step.python/cuda_stf/tests/stf/test_nested_scopes.py#L468-L470: tighten below the five-step0.5batch.python/cuda_stf/tests/stf/test_nested_scopes.py#L525-L525: tighten below the0.25step.python/cuda_stf/tests/stf/test_nested_scopes.py#L574-L576: tighten below the five-step0.5batch.python/cuda_stf/cuda/stf/_experimental/device_array.py-62-97 (1)
62-97: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: Retain the stream object until deallocation completes.
Only
_stream_intsurvives construction. Ifstreamowns its CUDA handle and is collected first, the finalizer passes a stale handle todplace.deallocate. Store the stream object and retain it through theweakref.finalizecallback arguments.python/cuda_stf/cuda/stf/_experimental/interop/numba.py-133-170 (1)
133-170: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftimportant: Return an immutable configured launcher instead of mutating the kernel.
kernel[grid, block]stores_launch_cfgon the shared decorator and returnsself. A second configuration overwrites previously retained launchers, and concurrent calls can launch with the wrong grid, context, or execution place. Return a separate launcher object containingcfg.python/cuda_stf/cuda/stf/_experimental/green_places.py-71-75 (1)
71-75: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: Restore the caller’s CUDA context before returning.
green_places()switches the calling thread todevice_idand never switches it back, so later CUDA work in the same thread can run on the wrong context. Save the previous context and restore it in afinallyblock.python/cuda_stf/cuda/stf/_experimental/interop/pytorch.py-83-89 (1)
83-89: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: Bind the external stream to the STF task device. PyTorch accepts an optional
devicehere; omitting it leaves the stream on the current CUDA device, which can diverge from STF’sexec_placeon multi-GPU tasks. Resolve the task’s device and pass it totorch.cuda.ExternalStream(...)before enteringtc.stream(...).python/cuda_stf/tests/stf/examples/bicgstab.py-115-137 (1)
115-137: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftimportant: All three BiCGSTAB implementations omit the standard convergence exit after calculating
s, allowing exact convergence to proceed into an omega0/0.
python/cuda_stf/tests/stf/examples/bicgstab.py#L115-L137: testdot(s, s)before computingtandomega.python/cuda_stf/tests/stf/examples/burger.py#L233-L255: stop the captured iteration after applyingalpha * pwhenshas converged.python/cuda_stf/tests/stf/examples/burger_reference.py#L132-L149: preserve batched checks while handling exactsconvergence immediately.python/cuda_stf/tests/stf/examples/burger.py-325-351 (1)
325-351: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftimportant: Both Newton implementations decide continuation from a residual measured before the update they just applied.
python/cuda_stf/tests/stf/examples/burger.py#L325-L351: check convergence before solving, or recompute the residual after updatinglU.python/cuda_stf/tests/stf/examples/burger_reference.py#L169-L178: move the check beforebicgstab_solve, or calculatenorm2from the updatedtU.python/cuda_stf/tests/stf/examples/cg.py-90-145 (1)
90-145: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winimportant: Add an iteration limit to the CG loop.
The only exit condition is the absolute residual tolerance. If finite-precision CG stagnates above that threshold, the conditional graph replays indefinitely. Track an iteration scalar and combine
residual > tol²withiteration < max_iter.python/cuda_stf/tests/stf/examples/cholesky.py-744-759 (1)
744-759: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winimportant: Synchronize benchmark timing with the STF task streams.
Both events are recorded on CuPy’s current stream, while the factorization runs on independent STF streams; the stop event can complete before the factorization. Use an STF fence or device-synchronized wall-clock timing before reporting GFLOPS.
Also applies to: 777-789
python/cuda_stf/tests/stf/examples/potri.py-1029-1092 (1)
1029-1092: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winimportant: Stop POTRI timing before scheduling result verification.
With
--check, the measured interval also includes constructing test matrices, norm tasks,PDSYMM, andPDGEMM, yet Lines 1105-1106 report it as POTRI performance. Synchronize and capture the POTRI duration before Line 1047.python/cuda_stf/tests/stf/interop/test_fdtd.py-88-175 (1)
88-175: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: Use the spacing for each derivative axis.
Each curl currently applies one spacing to both terms. For example,
_update_exmust divide thehzy-difference bydyand thehyz-difference bydz; the other five components need the analogous axis-specific divisors. The current implementation is only correct whendx == dy == dz. Pass both required spacings to every kernel, or explicitly reject nonuniform grids.Also applies to: 280-329
python/cuda_stf/tests/stf/interop/test_jacobi_warp.py-196-270 (1)
196-270: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winimportant: Assert the Jacobi result or terminal residual.
This test never observes
A_host,Anew_host, orlresidual, so a no-op loop or broken Warp launch still passes. Back the residual with a host array and assert it is finite and<= tolafterfinalize(); also consider checking that the interior grid changed.
🟡 Minor comments (13)
AGENTS.md-274-275 (1)
274-275: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winsuggestion: Distinguish the STF example runner from the example sources. Line 275 says STF examples live in
python/cuda_stf/tests/test_examples.py, but the documented source tree ispython/cuda_stf/tests/stf/examples; describetest_examples.pyas the runner and keep the source/test locations underpython/cuda_stf/tests/stf. As per path instructions, “Keep Stream Task Flow tests underpython/cuda_stf/tests/stf.”Source: Path instructions
docs/conf.py-243-247 (1)
243-247: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winimportant: Mock
cuda.stf._experimental.cu12._stf_bindings_impl/cuda.stf._experimental.cu13._stf_bindings_implhere instead of the unversioned path, or drop the misleading entry and rely on the shim mock.AGENTS.md-210-210 (1)
210-210: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winimportant: Use one STF expansion. This entry says “Stream Task Flow,” but the STF docs and package metadata use “Sequential Task Flow.” Align the terminology so the guidance stays consistent.
docs/cudax/places.rst-437-460 (1)
437-460: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Document the rejection conditions for both transformations.
reshape()also requires positive extents, whilecollapse_axes()requires0 <= first <= last < 4; callers should also be told that invalid inputs are rejected. This keeps the C++ documentation aligned with the public Python contract and rejection tests.Source: Path instructions
python/cuda_stf/tests/stf/test_graph_scope.py-181-211 (1)
181-211: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: This test never waits on
fence_stream, so it does not verify that the stream completes after the first scope. The second scope’s dependency ordering andfinalize()can produce8.0even iffence()is broken. Synchronize the returned stream and validate the first scope’s observable result before submitting the second scope.python/cuda_stf/cuda/stf/_experimental/_stf_bindings.py-52-54 (1)
52-54: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Do not recommend an extra that the package does not provide.
For CUDA 14, this reports
pip install cuda-stf[cu14], but packaging only declarescu12andcu13. Report the unsupported CUDA major and list the actually supported extras instead.Also applies to: 72-74
python/cuda_stf/tests/stf/interop/test_jacobi_numba.py-117-119 (1)
117-119: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Assert the Jacobi result instead of treating completion as success.
Both tests pass after finalization without checking convergence or output correctness.
python/cuda_stf/tests/stf/interop/test_jacobi_numba.py#L117-L119: validate the resulting host arrays against a reference or convergence invariant.python/cuda_stf/tests/stf/interop/test_jacobi_pytorch.py#L62-L65: add the equivalent numerical assertion so stream and tensor-conversion regressions fail the test.python/cuda_stf/cuda/stf/_experimental/interop/numba.py-55-58 (1)
55-58: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Avoid mutating and returning the same launcher object.
__getitem__stores the launch config onself, so a savedkernel[grid, block]handle is overwritten by the next indexing call. Return a fresh wrapper per config.python/cuda_stf/tests/stf/interop/test_pytorch_task_context.py-11-16 (1)
11-16: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Assert that
FakeStreamContext.__exit__()was invoked. The test still passes ifpytorch_taskomits stream cleanup entirely, so it does not prove that the secondary cleanup failure was encountered and suppressed. Record anexitedflag and assert it alongsidefake_task.ended.Also applies to: 54-61
python/cuda_stf/tests/stf/interop/test_scoped_capture.py-93-112 (1)
93-112: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winimportant: Destroy
s_rawinfinallyblocks in both helpers. Any exception aftercudaStreamCreate()currently leaks the CUDA stream and can cascade into later tests.Also applies to: 130-170
python/cuda_stf/tests/stf/examples/stackable_branch_while_warp.py-174-179 (1)
174-179: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Assert the branch result instead of only printing it.
The example computes
expectedbut never validatesout_host; graph ordering, loop-count, or host-copy regressions therefore complete successfully. Add annp.allclose(out_host, expected)assertion after finalization.python/cuda_stf/tests/stf/interop/test_legacy_to_stf.py-311-395 (1)
311-395: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Invoke the advertised benchmark entry point.
python test_legacy_to_stf.pycurrently defines_benchmark()and exits without running it.print("\ncorrectness: OK for all sizes") + + +if __name__ == "__main__": + _benchmark()python/cuda_stf/tests/stf/interop/test_jacobi_warp.py-22-34 (1)
22-34: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winimportant: Add a CUDA 12.4+ skip before exercising
ctx.while_loop()andctx.repeat().pytest.importorskip()only covers optional deps, so these tests can still hit unsupported runtimes inpython/cuda_stf/tests/stf/interop/test_jacobi_warp.pyandpython/cuda_stf/tests/stf/interop/test_fdtd.py.
🧹 Nitpick comments (6)
cudax/include/cuda/experimental/__places/places.cuh (1)
654-667: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion: apply the required host API annotation to both new methods and their definitions.
These exception- and vector-based APIs should follow the repository’s host-function annotation convention. As per coding guidelines, “Functions must be marked with
_CCCL_HOST_API,_CCCL_DEVICE_API,_CCCL_HOST_DEVICE_API,_CCCL_TILE_API, or_CCCL_API.”Source: Coding guidelines
ci/build_cuda_stf_python.sh (1)
59-59: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winsuggestion:
wheelhouse/isn't cleared before the two docker builds populate it, so a leftover wheel from a previous local invocation could be picked up byfind ... | head -1instead of the freshly-built one.Fix: clean before building
mkdir -p wheelhouse +rm -f wheelhouse/cuda_stf-*.whlAlso applies to: 108-121
ci/build_cuda_stf_wheel.sh (2)
56-57: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion: CUDA major-version parsing from
nvcc --versionis duplicated with a different implementation inci/test_cuda_stf_python.sh(line 10). Worth extracting to a shared helper to avoid the two diverging ifnvcc --versionoutput format changes.
1-1: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion: Two independent parsers extract the same CUDA major version from
nvcc --versionoutput. One root cause: no shared helper for this. Ifnvcc --version's format ever changes, only one of these may get fixed, silently breaking the other.
ci/build_cuda_stf_wheel.sh#L56-57: replace with a call to a shared helper (e.g. added toci/util/) instead of the localgrep -oP/cutpipeline.ci/test_cuda_stf_python.sh#L10-10: replace theawk/tr/cutpipeline with the same shared helper.python/cuda_stf/tests/stf/interop/test_stencil_decorator.py (1)
74-86: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winsuggestion: Vectorize the interior reference calculation with NumPy slicing. The nested loops execute about one million Python iterations per test and add avoidable runtime across the CI matrix.
python/cuda_stf/tests/stf/interop/test_warp_pytorch_dag.py (1)
171-188: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftsuggestion: Validate concurrency explicitly or narrow the test’s stated contract.
The output assertions verify dependency ordering and the join, but a scheduler that serializes both sibling tasks produces identical results. Inspect captured graph topology/streams or add timing instrumentation if concurrent sibling execution is intended coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2a86a29c-0bcf-4b30-b1fa-1b6ea8b94982
📒 Files selected for processing (83)
.github/CODEOWNERSAGENTS.mdc/experimental/stf/include/cccl/c/experimental/stf/stf.hc/experimental/stf/src/stf.cuc/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/include/cuda/experimental/__places/places.cuhdocs/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_stream_utils.pypython/cuda_stf/tests/stf/test_task_graph.pypython/cuda_stf/tests/stf/test_token.pypython/cuda_stf/tests/test_examples.py
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>
A grid_dims product that does not match the number of places now raises a ValueError in the Python binding (stf_c_api review change) instead of a RuntimeError from the C layer. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
python/cuda_stf/README.md (2)
29-36: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion: The minimal variants section only shows
cu13andsysctk13examples. Consider addingminimal-cu12andminimal-sysctk12for completeness, since the non-minimal section lists all four variants.
55-56: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion: The phrase "cccl.c.experimental.stf / cudax" may confuse readers — the package namespace is
cuda.stf, notcudax. Ifcudaxrefers to the internal C++ extension name, clarify that distinction.python/cuda_stf/cuda/stf/_experimental/device_array.py (1)
51-60: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion:
_finalizerusesprint()for deallocation warnings. Consider usingwarnings.warn()or theloggingmodule so warnings are capturable and filterable rather than going to stdout.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1534bfa9-4a96-4601-bb9a-76caa962f093
📒 Files selected for processing (49)
AGENTS.mdci/build_cuda_stf_combined_python.shci/build_cuda_stf_python.shci/matrix.yamlci/project_files_and_dependencies.yamlci/test_cuda_stf_python.shdocs/python/setup.rstdocs/python/stf.rstdocs/python/stf_api.rstpython/cuda_stf/CMakeLists.txtpython/cuda_stf/README.mdpython/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/numba.pypython/cuda_stf/cuda/stf/_experimental/interop/pytorch.pypython/cuda_stf/cuda/stf/_experimental/paths.pypython/cuda_stf/merge_cuda_wheels.pypython/cuda_stf/pyproject.tomlpython/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/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_pytorch_task_context.pypython/cuda_stf/tests/stf/interop/test_stencil_decorator.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_lifecycle.pypython/cuda_stf/tests/stf/test_packaging.pypython/cuda_stf/tests/stf/test_place_support.pypython/cuda_stf/tests/test_examples.py
💤 Files with no reviewable changes (2)
- ci/project_files_and_dependencies.yaml
- python/cuda_stf/tests/stf/test_packaging.py
🚧 Files skipped from review as they are similar to previous changes (34)
- python/cuda_stf/cuda/stf/_experimental/_cuda_version_utils.py
- docs/python/setup.rst
- AGENTS.md
- python/cuda_stf/pyproject.toml
- python/cuda_stf/tests/stf/interop/test_stencil_decorator.py
- ci/test_cuda_stf_python.sh
- python/cuda_stf/merge_cuda_wheels.py
- python/cuda_stf/tests/stf/interop/test_jacobi_pytorch.py
- python/cuda_stf/tests/stf/test_cuda_kernel.py
- python/cuda_stf/tests/stf/interop/test_fdtd.py
- python/cuda_stf/tests/stf/interop/test_jacobi_numba.py
- python/cuda_stf/tests/stf/examples/cg.py
- python/cuda_stf/cuda/stf/_experimental/interop/pytorch.py
- python/cuda_stf/tests/stf/test_from_context.py
- docs/python/stf.rst
- python/cuda_stf/tests/stf/examples/stackable_branch_while_warp.py
- python/cuda_stf/cuda/stf/_experimental/green_places.py
- python/cuda_stf/tests/stf/examples/neural_ode_rk4.py
- ci/build_cuda_stf_python.sh
- python/cuda_stf/cuda/stf/_experimental/paths.py
- python/cuda_stf/CMakeLists.txt
- python/cuda_stf/tests/stf/examples/burger.py
- ci/matrix.yaml
- python/cuda_stf/tests/stf/test_place_support.py
- python/cuda_stf/tests/stf/test_lifecycle.py
- python/cuda_stf/tests/stf/examples/burger_reference.py
- python/cuda_stf/tests/stf/interop/test_jacobi_warp.py
- python/cuda_stf/cuda/stf/_experimental/interop/numba.py
- python/cuda_stf/tests/stf/examples/cholesky.py
- python/cuda_stf/tests/stf/test_context.py
- python/cuda_stf/tests/stf/examples/neural_ode_dopri5.py
- python/cuda_stf/tests/test_examples.py
- python/cuda_stf/tests/stf/examples/potri.py
- python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx
| .. py:method:: set_affine_data_place(dplace) | ||
| .. py:attribute:: affine_data_place | ||
| .. py:method:: pick_stream(resources, for_computation=True) | ||
| .. py:method:: get_place(idx) | ||
|
|
||
| .. py:class:: exec_place_grid | ||
|
|
||
| A grid of execution places (subclass of :py:class:`exec_place`). | ||
|
|
||
| .. py:staticmethod:: from_devices(device_ids) | ||
| .. py:staticmethod:: create(places, grid_dims=None, mapper=None) | ||
|
|
||
| Build a grid from device ordinals, or from explicit places with an optional | ||
| ``grid_dims`` shape (validated for rank, positivity, and product) and a ``mapper`` | ||
| partition function. The grid retains its sub-places. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
important: Document exec_place.reshape(dims) and collapse_axes(first_axis, last_axis). These are primary APIs introduced by this PR, but neither appears in the explicit compiled-extension reference.
As per path instructions, documentation review must verify that public API changes have matching documentation updates.
Source: Path instructions
| def _launch_grid(n: int): | ||
| return (n + _THREADS_PER_BLOCK - 1) // _THREADS_PER_BLOCK, _THREADS_PER_BLOCK |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== repo files ==\n'
git ls-files | rg '(^|/)(CONTRIBUTING\.md|CI|ci|README|python/cuda_stf/tests/stf/examples/fhe(_decorator)?\.py|python/cuda_stf/tests/stf/examples/)' || true
printf '\n== file sizes ==\n'
wc -l python/cuda_stf/tests/stf/examples/fhe.py python/cuda_stf/tests/stf/examples/fhe_decorator.py 2>/dev/null || true
printf '\n== outline fhe.py ==\n'
ast-grep outline python/cuda_stf/tests/stf/examples/fhe.py --view expanded || true
printf '\n== outline fhe_decorator.py ==\n'
ast-grep outline python/cuda_stf/tests/stf/examples/fhe_decorator.py --view expanded || true
printf '\n== relevant excerpts ==\n'
sed -n '1,220p' python/cuda_stf/tests/stf/examples/fhe.py
printf '\n---\n'
sed -n '1,220p' python/cuda_stf/tests/stf/examples/fhe_decorator.py
printf '\n--- tests/search ---\n'
rg -n "empty|zero-length|len\\(|_launch_grid|ciphertext|block" python/cuda_stf/tests/stf/examples python/cuda_stf/tests -g '!**/__pycache__/**' || trueRepository: NVIDIA/cccl
Length of output: 50367
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== docs ==\n'
for f in CONTRIBUTING.md .github/workflows/* docs/*; do
[ -e "$f" ] || continue
done
git ls-files | rg '(^|/)(CONTRIBUTING\.md|.*CI.*\.md|.*ci.*\.md|docs/.+\.md|\.github/workflows/.+\.ya?ml)$' || trueRepository: NVIDIA/cccl
Length of output: 1629
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== CONTRIBUTING excerpt ==\n'
sed -n '1,220p' CONTRIBUTING.md
printf '\n== ci overview excerpt ==\n'
sed -n '1,220p' docs/infrastructure/ci/references/ci_overview.rst
printf '\n== fhe.py relevant lines ==\n'
sed -n '60,250p' python/cuda_stf/tests/stf/examples/fhe.py
printf '\n== fhe_decorator.py relevant lines ==\n'
sed -n '45,210p' python/cuda_stf/tests/stf/examples/fhe_decorator.py
printf '\n== zero-length searches ==\n'
rg -n "length.*0|len\\(.*\\)== 0|\\[\\]|empty input|empty ciphertext|zero[- ]length|parametrize\\(\"length\", \\[0|0\\]" python/cuda_stf/tests/stf/examples python/cuda_stf/tests -g '!**/__pycache__/**' || true
printf '\n== launch grid uses ==\n'
rg -n "_launch_grid\\(" python/cuda_stf/tests/stf/examples/fhe.py python/cuda_stf/tests/stf/examples/fhe_decorator.pyRepository: NVIDIA/cccl
Length of output: 50369
🌐 Web query:
Numba CUDA kernel launch zero blocks invalid configuration grid size 0 documentation
💡 Result:
In CUDA, attempting to launch a kernel with zero blocks (a grid size of 0 in any dimension) is an invalid configuration that causes the underlying CUDA driver to return an error, typically manifesting as a CUDA_ERROR_INVALID_VALUE or cudaErrorInvalidConfiguration [1][2][3]. The CUDA architecture requires each dimension of the grid to be at least 1 [3]. If your Numba code encounters this error, it is generally because the grid dimension calculation results in zero when the input data size is empty or zero [1][2]. To resolve this: 1. Add explicit guards in your host code to check for zero-sized inputs before calculating grid dimensions or launching the kernel [4][2]. 2. If the logic is intended to handle zero-element work, implement an early return or skip the kernel launch entirely when the workload is zero [1][5][4]. For example, when calculating the number of blocks per grid for a 1D kernel: blockspergrid = (data_size + threadsperblock - 1) // threadsperblock If data_size is 0, this might result in 0 blocks [6]. Always ensure that the grid dimensions passed to the kernel launch are at least 1 [4]. Numba has previously addressed internal issues where certain operations (such as slicing or utility functions like forall) could inadvertently attempt to launch kernels with zero blocks [1][5]. If you are using standard Numba kernel launches, you are responsible for ensuring the configuration conforms to CUDA requirements [7][6]. Additionally, note that cudaErrorInvalidConfiguration is a sticky error; if a launch fails, the CUDA runtime state may be corrupted, and subsequent valid CUDA API calls may also fail until the error is cleared (e.g., via synchronization or driver reset) [2].
Citations:
- 1:
cuda.to_device()cannot copy zero-element data numba/numba#5017 - 2: [Bug]
silu_and_muland other activation kernels crash withcudaErrorInvalidConfigurationwhennum_tokens=0sgl-project/sglang#23609 - 3: [BUG] Compute Sanitizer reports cudaErrorInvalidConfiguration error when writing ORC rapidsai/cudf#13887
- 4: https://stackoverflow.com/questions/19022725/running-zero-blocks-in-cuda
- 5: Don't launch ForAll kernels with 0 elements (Fixes #5017) numba/numba#5023
- 6: https://numba.readthedocs.io/en/stable/cuda/kernels.html
- 7: https://numba.readthedocs.io/en/latest/cuda-reference/kernel.html
important: n == 0 still yields a 0-block launch here, and CUDA rejects that configuration. Guard the empty case in _launch_grid or short-circuit the arithmetic paths, and add length-0 test coverage in both python/cuda_stf/tests/stf/examples/fhe.py and python/cuda_stf/tests/stf/examples/fhe_decorator.py.
📍 Affects 2 files
python/cuda_stf/tests/stf/examples/fhe.py#L68-L69(this comment)python/cuda_stf/tests/stf/examples/fhe_decorator.py#L50-L51
Source: Coding guidelines
| def test_pytorch_task_exit_body_error_wins_over_cleanup(monkeypatch): | ||
| """A body failure is preserved even when both cleanups also fail.""" | ||
| FakeTask, FakeContext = _make_pytorch_env( | ||
| monkeypatch, | ||
| stream_exit_error=RuntimeError("stream cleanup failed"), | ||
| end_error=RuntimeError("task end failed"), | ||
| ) | ||
| task = FakeTask() | ||
|
|
||
| with pytest.raises(ValueError, match="body boom"): | ||
| with pytorch_interop.pytorch_task(FakeContext(task)): | ||
| raise ValueError("body boom") | ||
|
|
||
| # Both cleanups still ran even though the body raised. | ||
| assert task.ended |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
important: Track and assert that FakeStreamContext.__exit__() ran. This test currently proves only task.end() executed, so it would pass if body exceptions caused stream cleanup to be skipped.
Source: Coding guidelines
|
|
||
| fill_utils.init_logical_data(_FakeScalarContext(np.float32), _FakeLogicalData(), 7) | ||
|
|
||
| assert len(fill_calls) == 1 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
important: Assert the scalar fill payload and stream, not only the call count. The current test still passes if init_logical_data fills the scalar with the wrong byte pattern.
- assert len(fill_calls) == 1
+ expected = np.array([7], dtype=np.float32).tobytes()
+ assert fill_calls == [(expected, "stream")]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| assert len(fill_calls) == 1 | |
| expected = np.array([7], dtype=np.float32).tobytes() | |
| assert fill_calls == [(expected, "stream")] |
Source: Coding guidelines
Follow the C-order Python contract from the structured-partition work: reshape() takes a public C-order shape (linear place order is preserved because the public C-order enumeration and the native dimension-0-fastest enumeration coincide after reversal), collapse_axes() takes public axes mapped onto the native inclusive range [rank-1-last, rank-1-first] and validated against the grid's rank, and both stamp the result's grid rank so dims stays rank-aware. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
4 participants
Summary
exec_place.reshape()andcollapse_axes()Python methodsextern \"C\"Stacked on #5315 and #9898. Their commits appear in this diff until they land on
main; the bindings-specific change is[STF] Bind execution grid reshape operations.Test plan
pre-commit run --files c/experimental/stf/include/cccl/c/experimental/stf/stf.h c/experimental/stf/src/stf.cu c/experimental/stf/test/test_places.cpp python/cuda_stf/cuda/stf/_experimental/_stf_bindings_impl.pyx python/cuda_stf/tests/stf/test_composite_places.py docs/python/stf.rstcmake --build build/cccl-c-stf --target cccl.c.experimental.stf.test.places.cpp -j 4ctest --test-dir build/cccl-c-stf --output-on-failure -R '^cccl\.c\.experimental\.stf\.test\.places\.cpp$'python -m pip install -e python/cuda_stf --no-depspython -m pytest -v python/cuda_stf/tests/stf/test_composite_places.py(24 passed)