feat(cuda): cuda_device_count, and two device probes that still aborted - #297
Merged
Conversation
Verified on 2x RTX 3090 (driver 570.207, CUDA 12.8, LLVM 15.0.7 + NVPTX, rustc 1.97.1). Everything below was checked by running it, which is the point: none of the 3.8.0 CUDA changes had executed on hardware, because no CI job builds the extension with the feature and the dev box has neither GPU nor LLVM. The 3.8.0 work holds up. The `assert_ran_on_gpu` hard gate passes — no prime silently falls back — and the `GpuBackendReport` counters do count: a genuine run reports 16 reductions on GPU and 0 on CPU on either card, `device_id: None` reports 0/16, and an absent ordinal reports 0/16 with `first_gpu_error` latched and the basis still correct. `ran_on_gpu()` is therefore observably true and observably false, not permanently one or the other. `cuda_device_count()` is added, which 3.8.0 deliberately deferred for want of a machine that could compile and run it. `catch_unwind` is load-bearing: `cudarc` panics rather than returning `Err` when `libcuda.so` cannot be dlopen'd, so a naive binding aborts the process on exactly the driverless machines a capability probe exists to report on. It reports 0 for every "no GPU here" shape and never raises. Both directions are pinned by tests that run: every ordinal below the count launches, the ordinal at the count is refused. That same panic still aborted two probes. The 3.8.0 fix landed on `groebner_cuda.rs::gpu_available` only; `nvptx_gpu.rs::device_available` kept the bare `CudaContext::new(0)`, and `nvptx_multi_device_both_3090s` called `CudaContext::device_count()` direct. Both aborted where they meant to skip, on the machines their own doc comments promise to support. Both now go through the guarded path. racecheck is promoted to `continue-on-error: false`, but the honest reason to read it is on the step, not in this message: it detects `__shared__` hazards and no kernel here uses shared memory — `ELIMINATE_ROW_PTX` declares none, `compile_cuda`'s output declares none, launches pass `shared_mem_bytes: 0`. It cannot fail today. It is a gate for the first kernel that introduces `__shared__`, which is exactly when nobody would remember to add one. A green tick from it means "no shared memory", not "no races", and the step now says so. memcheck, by contrast, is doing real work, and this was checked rather than assumed: removing the bounds guard from `eliminate_row_kernel` turned `ERROR SUMMARY: 0 errors` into 272 errors reported as `Invalid __global__ read of size 8 bytes at eliminate_row_kernel+0xf0`. It is analysing kernel instructions by name and offset, not the build tool. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
📝 WalkthroughWalkthroughThe change adds safe CUDA device discovery across Rust and Python APIs. It updates exports and GPU documentation, adds fallback and ordinal tests, and makes the CUDA nightly racecheck a required workflow gate. ChangesCUDA device discovery
Estimated code review effort: 3 (Moderate) | ~25 minutes Mergeability Score: 🟡 Moderate · up to The PR adds CUDA device validation, but an invalid device ordinal can still be accepted for empty batches, masking configuration errors and violating the documented behavior. Merge should wait for validation to occur before the empty-batch return or for the contract to be explicitly narrowed. Sequence Diagram(s)sequenceDiagram
participant PythonCaller
participant PythonExtension
participant RustAPI
participant CUDADriver
PythonCaller->>PythonExtension: call cuda_device_count()
PythonExtension->>RustAPI: call cuda_device_count()
RustAPI->>CUDADriver: initialize CUDA and query devices
CUDADriver-->>RustAPI: count or initialization failure
RustAPI-->>PythonExtension: count or 0
PythonExtension-->>PythonCaller: return integer count
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@alkahest-core/src/jit/nvptx.rs`:
- Around line 90-115: Update CudaCompiledFn::call_batch_on to validate
device_ordinal against cuda_device_count() before returning early for an empty
output, so out-of-range ordinals never succeed. Preserve the existing
empty-batch success behavior for valid ordinals and keep cuda_device_count’s
documented boundary consistent.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 5a4d1401-7ade-423f-a2ba-41b55a88d467
📒 Files selected for processing (11)
.github/workflows/cuda_nightly.ymlalkahest-core/src/jit/mod.rsalkahest-core/src/jit/nvptx.rsalkahest-core/src/lib.rsalkahest-core/tests/groebner_cuda.rsalkahest-core/tests/nvptx_gpu.rsalkahest-py/src/lib.rsdocs/mdbook/src/gpu.mdpython/alkahest/__init__.pytests/test_agent_contract.pytests/test_cuda.py
Comment on lines
+90
to
+115
| /// Number of CUDA devices visible to this process, or `0` when none are. | ||
| /// | ||
| /// Answers "which ordinals may I pass to [`CudaCompiledFn::call_batch_on`]?" | ||
| /// without the caller having to probe by launching and catching | ||
| /// `E-CUDA-003`, which is what `docs/mdbook/src/gpu.md` had to recommend | ||
| /// while no verified implementation existed. | ||
| /// | ||
| /// Returns `0` rather than an error for every "no GPU here" shape — no | ||
| /// driver, no device, driver too old. The distinction a caller acts on is | ||
| /// "can I use a GPU", and each of those answers it identically; an ordinal | ||
| /// is valid iff it is `< cuda_device_count()`. | ||
| /// | ||
| /// `catch_unwind` is load-bearing for the same reason it is in | ||
| /// `groebner_cuda.rs::gpu_available`: `cudarc` *panics* rather than | ||
| /// returning `Err` when `libcuda.so` cannot be dlopen'd at all, which is the | ||
| /// state of any machine with no driver installed. A capability probe that | ||
| /// aborts the process on the exact configuration it exists to report on | ||
| /// would be worse than useless. | ||
| #[cfg(feature = "cuda")] | ||
| pub fn cuda_device_count() -> usize { | ||
| std::panic::catch_unwind(|| { | ||
| cudarc::driver::CudaContext::device_count() | ||
| .map(|n| n.max(0) as usize) | ||
| .unwrap_or(0) | ||
| }) | ||
| .unwrap_or(0) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Make the ordinal boundary true for empty batches.
Line 100 states that an ordinal is valid only when it is below cuda_device_count(). However, CudaCompiledFn::call_batch_on returns Ok(()) for an empty output before it loads or validates device_ordinal. An out-of-range ordinal can therefore succeed for an empty batch.
Validate the ordinal before the empty-batch return, or restrict the documented contract to nonempty launches.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@alkahest-core/src/jit/nvptx.rs` around lines 90 - 115, Update
CudaCompiledFn::call_batch_on to validate device_ordinal against
cuda_device_count() before returning early for an empty output, so out-of-range
ordinals never succeed. Preserve the existing empty-batch success behavior for
valid ordinals and keep cuda_device_count’s documented boundary consistent.
Merging this PR will not alter performance
Comparing 
Footnotes
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Verified on 2× RTX 3090 (driver 570.207, CUDA 12.8, LLVM 15.0.7 + NVPTX, rustc 1.97.1 to match CI). Every claim below was checked by running it — none of the 3.8.0 CUDA changes had executed on hardware, because no CI job builds the extension with the feature.
Verification of what 3.8.0 shipped
assert_ran_on_gpuhard gateGpuBackendReportcounters countcapabilities()contract v3groebner_cuda/numpygonegpu_available()catch_unwindThe counters are observably both true and false, so
ran_on_gpu()is not permanently one or the other:device_idran_on_gpu()fell_back_to_cpu()first_gpu_errorNoneSome(0)Some(1)Some(4096)memcheck is doing real work (item 4), checked rather than assumed. Removing the bounds guard from
eliminate_row_kernelturnedERROR SUMMARY: 0 errorsinto 272 errors, reported asInvalid __global__ read of size 8 bytes at eliminate_row_kernel+0xf0. It analyses kernel instructions by name and offset, not the build tool. Guard restored.What this changes
cuda_device_count()(item 6), which 3.8.0 deferred for want of a machine that could compile and run it.catch_unwindis load-bearing:cudarcpanics rather than returningErrwhenlibcuda.socannot bedlopen'd, so a naive binding aborts the process on exactly the driverless machines a capability probe exists to report on. Reports0for every "no GPU here" shape, never raises. Both directions pinned by tests that run on hardware: every ordinal below the count launches, the ordinal at the count is refused. Added to_FEATURE_EXCLUSIVE_NAMES["cuda"]so the v3 contract test covers it.Two probes still aborted on driverless machines. The 3.8.0
catch_unwindfix landed ongroebner_cuda.rs::gpu_availableonly.nvptx_gpu.rs::device_availablekept the bareCudaContext::new(0), andnvptx_multi_device_both_3090scalledCudaContext::device_count()direct — both abort where they mean to skip, on the machines their own doc comments promise to support. Both now go through the guarded path.A test for the error-fallback branch.
first_gpu_errorand theErr(e)arm of thereduce_gpudispatch had never executed anywhere: the existing coverage is theNonepath, which never touches the driver.racecheck (item 3) — promoted, with a caveat on the step
Flipped to
continue-on-error: false. It is clean on real hardware (0 hazards, twice) and was never the false-positive risk the old comment assumed.But read the coverage note before trusting a green tick. racecheck detects
__shared__memory hazards, and no kernel in this workspace uses any:ELIMINATE_ROW_PTXdeclares no.shared,compile_cuda's generated PTX declares none, launches passshared_mem_bytes: 0. So it cannot fail today, and its green tick is a statement about the absence of shared memory, not the absence of races. It is a gate for the first kernel that introduces__shared__— exactly when nobody would think to add one. The step is renamed to say so rather than left to be misread, since a ninth gate that passes while inspecting nothing is the thing this release exists to stop.Gates
cargo fmt --all -- --check— cleanruff format python/ tests/— 167 files unchanged;ruff checkall passedcargo clippy --all-targets --features cuda,groebner-cuda -- -D warnings— exit 0cargo test --features cuda,groebner-cuda— 2083 passed, 0 failedcargo check --workspace(no CUDA) — passes, confirming thecfg-gating does not break default buildspytest -q— 2709 passed, 57 skipped, 0 failedtests/test_cuda.py+test_agent_contract.pyunderALKAHEST_GPU_TESTS=1— 32 passed🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
cuda_device_count()in the Python and core APIs.Bug Fixes
Documentation
Tests