Skip to content

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

Description

@lucifer1004

What problem does this solve?

CUDA recommends passing a tensor map to a kernel as a read-only, by-value
__grid_constant__ parameter. cuda-oxide currently exposes TMA operations over
*const TmaDescriptor, so callers that need a runtime descriptor must allocate
device memory for its 128 bytes, copy the encoded descriptor to that allocation,
and pass an ordinary device pointer. This adds storage and transfer work and
does not express the intended immutable kernel-parameter ABI.

Proposed API / behavior

Provide one Rust kernel-parameter declaration whose device ABI and generated
host launcher both derive from it. The exact spelling is open to design; for
example:

#[kernel]
pub unsafe fn copy(
    input_tma: GridConstant<TmaDescriptor>,
    // ...
) {
    cp_async_bulk_tensor_2d_g2s(
        dst,
        input_tma.as_ptr(),
        coord0,
        coord1,
        barrier,
    );
}

The implementation would:

  • lower the aggregate kernel parameter to a pointer-typed LLVM parameter with
    byval(TmaDescriptor) and the required alignment;
  • add that parameter index to the kernel's grid_constant entry in
    !nvvm.annotations;
  • keep the device view read-only and addressable so it can be passed directly
    to TMA intrinsics;
  • make #[cuda_module] marshal the descriptor bytes as one by-value kernel
    argument, rather than requiring a device allocation;
  • derive both sides from the one parameter declaration rather than adding an
    independent launcher option.

Ordinary pointers to tensor maps in global memory should remain available for
the genuinely mutable/dynamically selected descriptor case.

Alternatives considered

  • Allocate a 128-byte device slot, copy the descriptor H2D, and pass its
    pointer. This works today but adds storage/transfer work and a separate
    lifetime to retain.
  • A constant-memory global is not per invocation and conflicts with concurrent
    kernels carrying different descriptors.
  • Inline PTX cannot change an already-established kernel parameter ABI.

Additional context

cuda-oxide already has aggregate-by-value launch machinery and emits other
kernel properties through nvvm.annotations; this request is for the missing
end-to-end parameter contract, including host marshalling.

Activity

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions