[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 635 commits into
Conversation
The comment referenced a non-existent _bindings.pyi and the wrong package (cuda.cccl.parallel). It was copy-pasted and does not apply to cuda.stf. Made-with: Cursor
The buffer-protocol fallback passed view.buf/view.len to STF assuming contiguous memory. For strided views this registered the wrong byte range. Add PyBUF_ANY_CONTIGUOUS to the request flags so PyObject_GetBuffer fails on non-contiguous inputs, and add a test. Made-with: Cursor
Coerce shape to a tuple of Python ints so numpy scalars and lists are handled consistently with the buffer-protocol path. Reject empty shapes and non-positive dimensions early with clear error messages. Made-with: Cursor
Wrap cupy/cupyx.scipy imports in try/except so users without CuPy get a helpful install message instead of a raw ImportError. Made-with: Cursor
Both FDTD tests now check that all field components are finite after the simulation loop, catching NaN/Inf regressions. Also clean up the simplified variant (remove unused params/imports/type alias, fix return annotation). Add TODO for host_launch in show_slice. Made-with: Cursor
Implement ctx.host_launch() which schedules a Python callable as a host-side task graph node with full dependency tracking. Dependencies are auto-unpacked as numpy arrays and passed to the callback. Key changes in _stf_bindings_impl.pyx: - Declare stf_ctx_finalize as nogil and release the GIL during finalize() to prevent deadlocks with host_launch callbacks - Add cdef extern declarations for all host_launch C API functions - Add _host_launch_trampoline (C callback -> Python) and _python_payload_destructor (prevent leaks on cleanup) - Add context.host_launch() method with try/finally for handle safety Adapt existing tests to use host_launch: - test_fhe.py / test_fhe_decorator.py: replace exec_place.host() + managed + synchronize workaround with host_launch for print_values and final assertion - test_fdtd_pytorch.py / test_fdtd_pytorch_simplified.py: replace GPU-side torch.isfinite task with host_launch using np.isfinite New test file test_host_launch.py with six tests: basic, multiple_deps, write_back, chained, no_deps, and loop_value_capture. Made-with: Cursor
Exposes the C API's stf_fence() as context.fence(), returning a CUDA stream handle that is signaled when all pending tasks complete. Unlike finalize(), the context stays alive so more work can be submitted. Made-with: Cursor
The C declaration was already present but no Python method used it. Made-with: Cursor
…ches Expose stf_cuda_kernel C API in Python, enabling native CUDA graph kernel nodes instead of stream capture. Uses cuda.core.ParamHolder for argument marshalling (matching cuda.core.launch conventions) and handles the CUkernel-to-CUfunction conversion needed by modern cuda.core. Includes tests for AXPY (stream and graph), chained tasks, raw handles, and a loop-accumulate test that stresses per-iteration argument lifetime. Made-with: Cursor
…-launch Previously, each call to launch() replaced self._arg_holder, so only the last ParamHolder survived until end(). When multiple launches were chained inside a single cuda_kernel task, earlier argument buffers could be freed before STF executed the kernels. Switch to a list (_arg_holders) that accumulates every ParamHolder and is cleared in end(). Add test_cuda_kernel_multi_launch to verify two launches with different scalar args in one task produce the correct result. Made-with: Cursor
Add NULL checks after every handle-returning C API call (context, task, cuda_kernel, logical_data, exec_place, data_place, host_launch), raising RuntimeError on failure instead of proceeding with a null handle. Wrap all __dealloc__ methods in try/except so cleanup failures are printed rather than raised, matching the cuda.compute pattern. Add tests for double-finalize safety, borrowed-context rejection, and post-finalize error propagation. Made-with: Cursor
Change partition_fn_t (C++ and C API) from return-by-value to out-pointer convention so the mapper callback is trivially representable in ctypes/cffi/Rust. Update all C++ partitioners (blocked, tiled, cyclic) and call sites. Add Python bindings for exec_place_grid (from_devices, create), data_place.composite with ctypes mapper bridge, exec_place introspection (dims, size, set_affine_data_place), and current_device factories. Made-with: Cursor
…stream_ptrs Expose grid-task query APIs across the full stack for multi-place dispatch: - C++ unified_task: add get_grid_dims(dim4*) and get_stream(size_t) - C API: add stf_task_get_grid_dims() and stf_task_get_custream_at_index() - Python task: add get_grid_dims(), get_stream_at_index(), get_stream_ptrs() - C tests: grid dims/stream query + error case for non-grid exec_place - Python tests: grid task dims, stream queries, scalar fallback, affine dep Made-with: Cursor
Made-with: Cursor
Made-with: Cursor
Port C++ STF examples that use CUB/Thrust to Python using cuda.compute: - Reduce (cf. 08-cub-reduce.cu): cuda.compute.reduce_into inside STF task - Inclusive/exclusive scan (cf. scan.cu): cuda.compute.inclusive_scan - Binary/unary transform (cf. thrust_zip_iterator.cu): cuda.compute transforms - Multi-task pipeline: binary_transform then reduce_into with auto dependency The bridge is a lightweight StfStream wrapper that adapts task.stream_ptr() to the __cuda_stream__ protocol expected by cuda.compute algorithms. Made-with: Cursor
…ified tests - Add stf_ctx_wait() C API and Python ctx.wait() to read logical data contents without finalizing the context (enables convergence checks). - Fix use-after-free in logical_data: store reference to source buffer object (_source_buf) to prevent GC before async H2D copy completes. - Add numba_task context manager (tests/stf/numba_task.py) that wraps ctx.task() and yields (numba_arrays, stf_stream) for cleaner tests. - Add simplified tests using numba_task + ctx.wait() + logical_data_empty. - Replace numba.cuda.as_cuda_array with from_cuda_array_interface(sync=False) to avoid spurious stream synchronization (perf bug in stream mode, correctness bug in graph capture mode). - Fix wrong expected value in test_stf_binary_transform (double-counted h_b). - Skip graph test: cuda.compute uses cudaMallocAsync internally, which creates mem-alloc graph nodes incompatible with cuGraphAddChildGraphNode. - Add C test for stf_ctx_wait. Made-with: Cursor
…, machine_init Extend the STF C API and Python bindings with place primitives needed for task-free usage (e.g. Python torq). This enables activating a place's CUDA context, obtaining streams from the place's pool, querying sub-places and affine data places, and initializing the machine singleton -- all without requiring an STF task graph. C API: stf_exec_place_scope_enter/exit, stf_exec_place_get_affine_data_place, stf_exec_place_pick_stream, stf_exec_place_get_place, stf_machine_init. Python: exec_place context manager (with place:), pick_stream(), get_place(), affine_data_place property, __getitem__, and stf.machine_init(). Made-with: Cursor
…, and DeviceArray Add stf_data_place_allocate(), stf_data_place_deallocate(), and stf_data_place_allocation_is_stream_ordered() to the C API, with corresponding Python bindings on the data_place class. Introduce DeviceArray, a lightweight 1D device array backed by data_place.allocate() that implements __cuda_array_interface__ for cuda.compute compatibility. This replaces direct numba.cuda usage for device memory management, ensuring allocations flow through the places abstraction (critical for green contexts, composite places, and other future place types). Made-with: Cursor
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
make_partition and its surroundings introduce several non-obvious concepts that only lived in code comments: the reference-shape binding (and how it differs from the scale-free classic policies), padding to divisibility and the predication idiom, evaluation before allocation, the per-tensor nature of partition-backed composite places, and the linearization convention. Give them a section in the places guide, anchored against the existing partitioning-policies documentation. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
allocate_nd(dim4(nbytes), 1) only works for composites built from scale-free partitioners; a structured-partition composite validates the extents against the partition's tensor and rejects a flat byte count. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Cedric AUGONNET <caugonnet@nvidia.com>
sin/cos/fabs were relying on a transitive include. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Cedric AUGONNET <caugonnet@nvidia.com>
…tion Reversing only the extents would re-target each dim_spec at the wrong axis; the spec list and owner() coordinates reverse together. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Cedric AUGONNET <caugonnet@nvidia.com>
The byte-count allocate() throws for composite places since the allocate_nd change. Default to the explicit flat byte geometry ((nbytes,), elemsize=1) - equivalent for regular places, and the byte-domain contract callback composites expect - and add dims=/ elemsize= keywords so structured-partition composites (which require their tensor's exact extents) can allocate through DeviceArray too. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Cedric AUGONNET <caugonnet@nvidia.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.
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