Skip to content

[FEA] cuda.core: support stream batch memory operations (incl. CUDA 13.1 atomic reductions) #2364

Description

@leofang

Summary

CUDA 13.1 extends the stream batch memory operations with atomic reductions: a new op type
CU_STREAM_MEM_OP_ATOMIC_REDUCTION with params struct CUstreamMemOpAtomicReductionParams
(operation, flags, reductionOp, dataType, address, value, alias), plus enums
CUstreamAtomicReductionOpType (ADD/AND/OR) and CUstreamAtomicReductionDataType
(UNSIGNED_32/UNSIGNED_64), joined into CUstreamBatchMemOpParams_union.atomicReduction.
Device support is advertised by CU_DEVICE_ATTRIBUTE_ATOMIC_REDUCTION_SUPPORTED (13.1, see
also #2361).

cuda.core has no batch mem-op surface at all today — the only mentions are a
graph debug-dot print flag
(

if self.batch_mem_op_node_params:
)
and a DeviceProperties support query
(
"""bool: 64-bit operations are supported in cuStreamBatchMemOp and related MemOp APIs."""
).
So this issue covers designing the base surface (the pre-12.8 wait/write-value ops) and
including the 13.1 atomic-reduction op in it — adding the new op alone is not possible.

Underlying C APIs to cover

Symbol Purpose
cuStreamBatchMemOp(CUstream, unsigned int count, CUstreamBatchMemOpParams* paramArray, unsigned int flags) execute a batch of memory operations on a stream (pre-12.8, unexposed)
CU_STREAM_MEM_OP_ATOMIC_REDUCTION + CUstreamMemOpAtomicReductionParams (13.1) atomic-reduction op entry
CUstreamAtomicReductionOpType, CUstreamAtomicReductionDataType (13.1) reduction op / operand type
CU_DEVICE_ATTRIBUTE_ATOMIC_REDUCTION_SUPPORTED (13.1) support gating

Design sketch (draft — needs design-meeting review)

Important

Starting point only, not a settled design — review in the cuda.core design meeting.

stream.batch_mem_op([                          # names TBD
    WaitValue(addr, value, condition=">="),
    WriteValue(addr, value),
    AtomicReduction(address=addr, value=1, op="add", dtype="u64"),  # 13.1+
])

Open questions for the meeting:

  1. Scope: stream-level only, or also the graph batch-mem-op node type?
  2. Typed per-op dataclasses (above) vs. lower-level pass-through of binding structs.
  3. Semantics/purpose of the alias field (clarify against driver docs during design).
  4. Version/device gating (op availability differs across 12.x/13.x and devices).

References

-- Leo's bot

Metadata

Metadata

Assignees

No one assigned

    Labels

    cuda.coreEverything related to the cuda.core modulefeatureNew feature or requesttriageNeeds the team's attention

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions