tilelang.cuda.language.atomic¶

CUDA dialect of the atomic operators: the common ops plus CUDA knobs.

Functions¶

atomic_add(dst, value[, memory_order, return_prev, ...])

Atomically add value into dst, with CUDA lowering knobs.

Module Contents¶

tilelang.cuda.language.atomic.atomic_add(dst, value, memory_order=None, return_prev=False, use_tma=False, annotations=None)¶

Atomically add value into dst, with CUDA lowering knobs.

Same semantics as the common tilelang.language.atomic.atomic_add(). use_tma selects the sm90+ TMA cp.reduce lowering for the tile-region path; targets without TMA reject it at compile time.

Parameters:

  • dst (Buffer) – Destination buffer/address to apply the atomic add.

  • value (PrimExpr) – Value to add atomically.

  • memory_order (Optional[str]) – Memory-order name controlling the atomic operation’s ordering (“relaxed”, “consume”, “acquire”, “release”, “acq_rel”, “seq_cst”).

  • return_prev (bool) – Return the previous value (scalar path only).

  • use_tma (bool) – If True, lower the tile-region atomic add through TMA cp.reduce. Available on sm90+ only (default False).

  • annotations (Optional[dict]) – Extra annotations for the tile-region path; values in it take precedence over the individual keywords.

Returns:

A handle to the atomic operation.

Return type:

PrimExpr