Skip to content

codegen: support grid-constant kernel parameters - #1225

Open
lucifer1004 wants to merge 4 commits into
NVlabs:mainfrom
lucifer1004:feat/grid-constant-tma-1223
Open

codegen: support grid-constant kernel parameters#1225
lucifer1004 wants to merge 4 commits into
NVlabs:mainfrom
lucifer1004:feat/grid-constant-tma-1223

Conversation

@lucifer1004

Copy link
Copy Markdown
Contributor

Summary

Adds parameter-local #[grid_constant] support so one Rust kernel declaration controls both the device grid-constant ABI and generated host by-value marshalling.

#[kernel]
fn copy(#[grid_constant] descriptor: &TmaDescriptor, output: *mut u32) {
    // descriptor is a read-only parameter-space reference
}

Closes #1223.

Changes

  • Validate #[grid_constant] on immutable, sized reference parameters.
  • Generate a by-value host argument from the same kernel declaration.
  • Carry source parameter index, pointee type, and ABI alignment through MIR import and lowering.
  • Map source parameters through the existing scalarized kernel ABI before LLVM export.
  • Emit LLVM byval(T) align N and the one-based NVVM grid_constant parameter list.
  • Extend the existing cuda_module_contract example with a 128-byte, 64-byte-aligned SM120 execution check.

Testing

  • just check passes (all stages through CUDA-linked tests and the pre-deny guards passed; the local run stopped because cargo-deny is not installed)
  • cargo oxide run cuda_module_contract --arch sm_120 passes
  • Existing example extended; no new example directory added.
  • Targeted cuda-macros, llvm-export, mir-importer, mir-lower, and reserved-oxide-symbols tests pass.
  • Targeted clippy with -D warnings passes.
  • just check-intrinsics passes for all 1,025 routes and 2,374 target configurations.
  • just doc-check passes.
  • A downstream CUDA 13.3 release CUBIN build and GPU integration suite pass.

Checklist

  • All commits signed off
  • SPDX headers on new source files (no new source files)

Let #[kernel] declarations mark immutable reference parameters as grid constants, derive the generated host launch ABI from the same declaration, and carry pointee layout through MIR and LLVM lowering.

Emit LLVM byval alignment and NVVM grid_constant metadata, with compiler and SM120 contract coverage.

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
@nihalpasham

Copy link
Copy Markdown
Collaborator

The generic-kernel path needs to be part of the design before this lands:

  • For a generic #[kernel], the macro splits your fn into a generated #[inline(always)] helper that keeps the body and a concrete entry wrapper per instantiation that calls it. The marker lands in both. When rustc's MIR inliner folds the generated helper into the entry, detect_grid_constant_params hard-errors on the duplicate.
  • Repro:
#[kernel] 
pub fn k<T: Copy>(#[grid_constant] d: &Desc, out: *mut u32, _tag: T)  { 
   unsafe { *out = d.values[3] } 
 } 

fails with kernel contains duplicate grid-constant marker for source parameter 0. Bigger bodies only compile because they exceed the inliner's cost threshold.

  • More importanly, the helper's marker leaks: a plain kernel
other(d: &Desc, out: *mut u32) { 
  k::<u8>(d, out, 0) 
} 

with no attribute comes out with .param .align 64 .b8 other_param_0[128] and a grid_constant annotation. Host passes an 8-byte pointer, device expects 128 bytes by value.

  • unchecked_indexing already handles this by stripping its marker from the re-emitted helper, see strip_unchecked_indexing_config_marker. grid_constant needs the same in both generic paths. A duplicate-tolerant detector alone fixes the error but not the leak.

And we need tests in the PR exercise a generic kernel with #[grid_constant].

To land this: the strip in both generic paths, a generic #[kernel] + #[grid_constant] test with a body small enough to be inlined, and a negative test that a non-opted kernel calling the helper still gets a plain pointer param.

@nihalpasham nihalpasham added enhancement New feature or request codegen Device code-generation pipeline (Rust MIR to IR to PTX) cuda-feature A CUDA hardware/toolkit feature (TMA, tcgen05, WGMMA) labels Sep 6, 2026

@nihalpasham nihalpasham left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changes requested, details in #1225 (comment)

@lucifer1004

Copy link
Copy Markdown
Contributor Author

Addressed the generic-helper ABI concern in f2b6b15.

The macro now captures the grid-constant marker for the generated entry and strips it from the re-emitted callable helper in both generic expansion paths. The contract example now covers both sides:

  • an instantiated generic grid-constant entry retains the 128-byte by-value PTX parameter and launches successfully on sm_120;
  • an ordinary kernel calling the same generic helper retains a .param .u64 pointer ABI and does not acquire the 128-byte parameter.

Measured validation:

  • pixi run cargo test -p cuda-macros
  • strict clippy for cuda-macros and cuda_module_contract
  • pixi run cargo oxide run cuda_module_contract --arch sm_120 -- --verify-ptx
  • pixi run cargo oxide run cuda_module_contract --arch sm_120
  • the full just check code/test portion passed; its cargo-deny guard initially hit the local NFS advisory-lock limitation, then all deny/license guards passed with the advisory DB placed on local /tmp storage.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

codegen Device code-generation pipeline (Rust MIR to IR to PTX) cuda-feature A CUDA hardware/toolkit feature (TMA, tcgen05, WGMMA) enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support grid-constant by-value kernel parameters for TMA descriptors

2 participants