Primitives#
NVVM wrapper namespace — hand-maintained wrappers over raw MLIR NVVM dialect ops.
Purpose#
This module gives users a single nvvm.* namespace for all
NVVM-level GPU operations (barriers, TMA copies, tcgen05 tensor-core
ops, PTX special-register reads, etc.) without requiring them to deal
with raw MLIR ceremony.
How ops are exposed#
There are three categories of entries in this module:
Wrapped ops (
@dsl_user_opfunctions): A wrapper exists when it adds genuine value over the raw MLIR op:Return-type hiding — the raw MLIR op requires the caller to pass the result type as the first positional argument (e.g.
T.i32()). The wrapper inserts it automatically and wraps the return value in a typed wrapper (Int32,Boolean, …). Examples:shfl_sync.Python-to-DSL type coercion — the raw MLIR op expects every operand to be an
ir.Value. The wrapper accepts plain Pythonint/booland converts them to the correct DSL type (Int32,Int64,Boolean, …) so the proxy can turn them intoir.Value. The proxy alone cannot do this because it does not know which MLIR type a Python literal should become. Because MLIR integers are signless, coerced parameters also accept the unsigned counterpart (e.g.count: int | Int32 | Uint32). The coercion always uses the signed type internally —Int32andUint32both producei32. Examples:mbarrier_init,tcgen05_alloc,tcgen05_mma.
Every wrapper maps 1:1 to a single MLIR NVVM dialect op with the same argument order. Higher-level convenience ops (e.g. computing derived parameters or specialising generic ops) belong in the higher-level namespace, not here.
Direct aliases (plain attribute assignments): When a raw NVVM op needs no coercion, no return-type hiding, and no default parameters, it is re-exported as-is so that it still appears in the
nvvm.*namespace.fence_mbarrier_inithas a thin wrapper below for documentation; other fence aliases are bare.Auto-converting proxy (
nvvm.dialect): For any NVVM op that is not listed in this module at all, the proxy can be used directly. It auto-converts any argument that has an.ir_value()method before forwarding to the raw dialect:nvvm.dialect.some_unlisted_nvvm_op(T.i32(), my_int32_value, ...)The proxy does not handle
int -> Int32coercion (it cannot guess the target MLIR type), so callers must wrap Python literals themselves when using it.
- cutlass.experimental.primitives.nvvm_wrapper.add_packed_f32x2(
- src_a: tuple | Vector,
- src_b: tuple | Vector,
- *,
- rnd: FPRoundingMode | None = None,
- ftz: bool | None = None,
Wrapper over
nvvm.add_packed_f32x2.Accepts a 2-tuple of f32 scalars or a
Vectorfor each operand and returns a tuple when called with tuples, else aVector
- cutlass.experimental.primitives.nvvm_wrapper.atomicrmw(
- op: ~cutlass.experimental.primitives.nvvm_wrapper.AtomicOp,
- ptr: ~cutlass.Array | ~cutlass.Pointer,
- a: int | float | ~cutlass.Int32 | ~cutlass.Uint32 | ~cutlass.Int64 | ~cutlass.Uint64 | ~cutlass.Float32 | ~cutlass.Float64,
- *,
- b: int | float | ~cutlass.Int32 | ~cutlass.Uint32 | ~cutlass.Int64 | ~cutlass.Uint64 | ~cutlass.Float32 | ~cutlass.Float64 | None = None,
- mem_order: ~cutlass.experimental.primitives.nvvm_wrapper.MemOrder | None = None,
- syncscope: ~cutlass.base_dsl.array.MemScope | None = None,
- space: _install_cutlass_mlir_autodoc_stub.<locals>._DocDialectObject | None = None,
- results: list | None = None,
Atomic read-modify-write on a memory location.
Emits
atom.{op}.{mem_order}.{scope}. The op is performed atomically on*ptr: the prior value is returned, and the new value (a function of the old value,a, and optionallyb) is written back. No data race is observable by other threads insyncscope.- Parameters:
op – One of
AtomicOp. For"add"/"min"/"max"the wrapper picks the dialect-level FADD / UMIN / UMAX variant from the operand dtype: float operands pick FADD; unsigned-integer operands pick UMIN / UMAX; signed-integer operands pick the signed variant.ptr – Pointer/Array to the target memory cell.
a – First operand — for
"cas"this is the expected old value; for everything else the value combined with*ptr.b – Second operand — only used for
"cas"(the new value).mem_order – Memory ordering —
"relaxed"/"acquire"/"release"/"acq_rel".syncscope – Scope across which
mem_orderis enforced —"cta"/"cluster"/"gpu"/"sys".results – Optional preallocated result list (advanced).
- Returns:
The old value at
*ptrbefore the op was applied.- Raises:
ValueError – if op is not a valid
AtomicOp, or ifop="cas"is used without the second operand b (the new value).
- cutlass.experimental.primitives.nvvm_wrapper.bar_warp_sync(mask: int | Int32 | Uint32) None#
Rendezvous all lanes named in mask at a warp-level barrier.
Maps to PTX
bar.warp.sync membermask. Every lane whose bit is set in mask must execute this call before any of them may proceed. The barrier also acts as an acquire-release memory fence: stores to shared or global memory issued by any lane in mask before the call are visible to all other lanes in mask after it, and loads issued after the call observe those stores. Usecute.arch.FULL_MASK(0xFFFFFFFF) to rendezvous all 32 lanes; pass a narrower mask to synchronize only a known active subset.This is the warp-level equivalent of
__syncwarp(mask)in CUDA C++.Constraints:
Every lane named in mask must reach
bar_warp_syncwith the same mask value. If any named lane diverges (e.g. is inside a branch only some lanes take), the remaining lanes stall indefinitely.Do not call inside a branch unless all lanes in mask are guaranteed to enter that branch.
# All 32 lanes rendezvous in a uniform region (SMEM read-after-write) nvvm.bar_warp_sync(cute.arch.FULL_MASK) # Warp-specialization: all lanes converge before diverging by role. # bar_warp_sync guarantees every lane sees any prior register/SMEM # writes (e.g. from nvvm.setmaxregister) before the if-branch. warp = cute.arch.warp_idx() is_tma_warp = warp == cutlass.Int32(0) nvvm.setmaxregister(40, "decrease") nvvm.bar_warp_sync(cute.arch.FULL_MASK) # all lanes rendezvous here if is_tma_warp: ... # TMA producer path else: ... # compute consumer path # Partial mask — only lanes 0–15 synchronize (must all be active) nvvm.bar_warp_sync(0x0000FFFF)
- Parameters:
mask (int or cutlass.Int32 or cutlass.Uint32) – 32-bit member mask; bit i set means lane i participates. All participating lanes must execute this call with the same mask value or the remaining lanes stall indefinitely. Pass
cute.arch.FULL_MASK(0xFFFFFFFF) for all 32 lanes.- Raises:
ValueError – if a static
intmask does not fit in 32 bits (outside[0, 0xFFFFFFFF]). RuntimeInt32/Uint32values pass through unchecked.
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cluster_arrive() None#
Register an acquire-release arrival on the cluster-wide barrier.
Emits
barrier.cluster.arrive(non-aligned, the ordered.releaseform): each CTA registers its arrival, and the arrival also orders the issuing CTA’s prior memory writes so they become visible to every other CTA once it returns from the paired wait. This is the ordered counterpart of the relaxedbarrier_cluster_arrive_relaxed()(which registers arrival without any memory ordering); prefer the relaxed form plus an explicit fence when you do not need the built-in acquire-release ordering. Pair every arrive withbarrier_cluster_wait()from every CTA.The aligned variant
barrier_cluster_arrive_aligned()additionally asserts that every thread in the issuing warp executes the instruction convergently.# Cluster rendezvous with built-in acquire-release ordering. if nvvm.elect_sync(): nvvm.barrier_cluster_arrive() nvvm.barrier_cluster_wait()
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cluster_arrive_aligned() None#
Aligned variant of
barrier_cluster_arrive().Emits
barrier.cluster.arrive.aligned. Same acquire-release cluster arrival, but the.alignedqualifier additionally asserts that every thread in the issuing warp executes this instruction convergently (behaviour is undefined if any lane in the warp does not reach it). Do not combine with single-thread election (e.g. insideelect_sync): the other lanes would never reach the instruction.
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cluster_arrive_relaxed() None#
Register a relaxed arrival on the cluster-wide barrier.
Emits
barrier.cluster.arrive.relaxed(non-aligned): each CTA in the cluster registers its arrival, but no memory ordering is implied. Pair every arrive with abarrier_cluster_wait()from every CTA to block until all have arrived. The relaxed variant is cheaper than the acquire-releasebarrier_cluster_arrive(); prefer relaxed + an explicit fence where ordering is genuinely needed.The aligned variant
barrier_cluster_arrive_relaxed_aligned()additionally asserts that every thread in the issuing warp executes the instruction convergently.# CTA_2 GEMM cleanup — both CTAs rendezvous before dealloc. if nvvm.elect_sync(): nvvm.barrier_cluster_arrive_relaxed() nvvm.barrier_cluster_wait()
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cluster_arrive_relaxed_aligned() None#
Aligned variant of
barrier_cluster_arrive_relaxed().Emits
barrier.cluster.arrive.relaxed.aligned. Same relaxed cluster arrival (no memory ordering), but the.alignedqualifier additionally asserts that every thread in the issuing warp executes this instruction convergently (behaviour is undefined if any lane in the warp does not reach it). Do not combine with single-thread election (e.g. insideelect_sync): the other lanes would never reach the instruction.
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cluster_wait() None#
Block until every CTA in the cluster has called
barrier.cluster.arrive*.Emits
barrier.cluster.wait(non-aligned): the issuing thread stalls until the cluster-wide arrival counter reaches the cluster’s CTA count. The counter is set by the paired arrive call from every CTA (seebarrier_cluster_arrive_relaxed()for the relaxed form, orbarrier_cluster_arrive()when acquire-release ordering is needed).The aligned variant
barrier_cluster_wait_aligned()additionally asserts that every thread in the issuing warp executes the instruction convergently.# See barrier_cluster_arrive_relaxed for the paired usage. if nvvm.elect_sync(): nvvm.barrier_cluster_arrive_relaxed() nvvm.barrier_cluster_wait()
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cluster_wait_aligned() None#
Aligned variant of
barrier_cluster_wait().Emits
barrier.cluster.wait.aligned. Same cluster-scope wait, but the.alignedqualifier additionally asserts that every thread in the issuing warp executes this instruction convergently (behaviour is undefined if any lane in the warp does not reach it). Do not combine with single-thread election (e.g. insideelect_sync): the other lanes would never reach the instruction.
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cta_arrive( ) None#
Signal arrival at a named CTA barrier without waiting.
Emits
barrier.cta.arrive a, b;(non-aligned). Works in any control flow, including divergent on sm_70+. Producer/consumer pairs use the arrive/sync split: the producer arrives and runs ahead while the consumerbarrier_cta_sync()blocks until the count is reached. Each CTA has 16 named barrier slots (barrier_idin 0..15).For the aligned variant (
barrier.cta.arrive.aligned, equivalent to the legacybar.cta.arrive) usebarrier_cta_arrive_aligned()— that form promises that every CTA thread executes the barrier and is undefined behavior under divergent control flow.- Parameters:
barrier_id (int or Int32 or Uint32) – Barrier slot ID in 0..15. Must match the consumer’s id.
thread_count (int or Int32 or Uint32) – Number of participating threads. Required by PTX for
barrier.cta.arrive; must be a non-zero multiple of the warp size (32) and consistent across all arrive/sync calls on this slot.
- Raises:
ValueError – if a static
barrier_idis outside[0, 15]or a staticthread_countis not a positive multiple of 32. RuntimeInt32/Uint32values pass through unchecked.
# Producer warps signal arrival, then continue work if warp < N_PRODUCERS: nvvm.barrier_cta_arrive(0, (N_PRODUCERS + N_CONSUMERS) * 32) # ... continue producing ... else: # Consumer warps wait at the same id nvvm.barrier_cta_sync(0, thread_count=(N_PRODUCERS + N_CONSUMERS) * 32) # ... safe to read producer outputs ...
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cta_arrive_aligned( ) None#
Aligned variant of
barrier_cta_arrive().Emits
barrier.cta.arrive.aligned a, b;— equivalent to the legacybar.cta.arrive. Promises that every CTA thread executes this barrier in convergence; undefined behavior on sm_70+ when a strict subset of CTA threads reaches the instruction. Use this only when the call site is provably all-CTA-converged; otherwise fall back to the non-alignedbarrier_cta_arrive().- Parameters:
- Raises:
ValueError – if a static
barrier_idis outside[0, 15]or a staticthread_countis not a positive multiple of 32. RuntimeInt32/Uint32values pass through unchecked.
# Every CTA thread reaches this barrier, no divergent guards above nvvm.barrier_cta_arrive_aligned(0, threads_per_cta)
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cta_red(
- pred: int | Boolean,
- barrier_id: int | Int32 | Uint32,
- kind: BarrierRedux,
- *,
- thread_count: int | Int32 | Uint32 | None = None,
Synchronize a CTA barrier and reduce a predicate.
Emits the
barrier.cta.red.{popc,and,or}family (non-aligned). The call marks the issuing thread’s arrival at a named CTA barrier, waits until the barrier’s participant count is reached, then broadcasts the predicate reduction result to every waiting thread.For the aligned variant (
.alignedmodifier, equivalent to the legacybar.cta.red) usebarrier_cta_red_aligned()— that form promises all CTA threads execute the barrier and is undefined behavior under divergent control flow.Reduction kinds (selected by
kind):"and"→.and.pred— returnsBoolean,Trueiff every participant contributedpred=True."or"→.or.pred— returnsBoolean,Trueiff any participant contributedpred=True."popc"→.popc.u32— returnsInt32, the count of participants whosepredwasTrue.
Do not mix
barrier_cta_redwith the non-reducing variants (barrier_cta_sync(),barrier_cta_arrive()) on the same active barrier generation. PTX marks that use as unpredictable; use a differentbarrier_idor wait for the barrier to complete and reinitialize before reusing it.- Parameters:
pred (int or Boolean) – Per-thread predicate contributed to the reduction.
barrier_id (int or Int32 or Uint32) – CTA barrier slot ID in 0..15.
kind (BarrierRedux) – Reduction kind —
"and"/"or"yieldBoolean,"popc"yieldsInt32.thread_count (int or Int32 or Uint32, optional) – Number of participating threads. Omit for all CTA threads; otherwise pass a non-zero multiple of the warp size and keep it consistent across participants on this slot.
- Returns:
For AND/OR: the reduced
Booleanbroadcast to all participants. For POPC: the countInt32.- Raises:
ValueError – if
kindis not one of"and"/"or"/"popc", if a staticbarrier_idis outside[0, 15], or a staticthread_countis not a positive multiple of 32. RuntimeInt32/Uint32values pass through unchecked.
tx, _, _ = cute.arch.thread_idx() any_lane_zero = nvvm.barrier_cta_red( tx == 0, barrier_id=0, kind="or", thread_count=64, ) all_in_range = nvvm.barrier_cta_red( tx < 64, barrier_id=1, kind="and", thread_count=64, ) n_true = nvvm.barrier_cta_red( # returns Int32 tx % 2 == 0, barrier_id=2, kind="popc", thread_count=64, )
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cta_red_aligned(
- pred: int | Boolean,
- barrier_id: int | Int32 | Uint32,
- kind: BarrierRedux,
- *,
- thread_count: int | Int32 | Uint32 | None = None,
Aligned variant of
barrier_cta_red().Emits
barrier.cta.red.{popc,and,or}.aligned— equivalent to the legacybar.cta.red. Promises that every CTA thread executes this barrier in convergence; undefined behavior on sm_70+ when a strict subset of CTA threads reaches the instruction. Use this only when the call site is provably all-CTA-converged; otherwise fall back to the non-alignedbarrier_cta_red().Reduction kinds and return-type rules match
barrier_cta_red().- Parameters:
pred (int or Boolean) – Per-thread predicate contributed to the reduction.
barrier_id (int or Int32 or Uint32) – CTA barrier slot ID in 0..15.
kind (BarrierRedux) – Reduction kind —
"and"/"or"yieldBoolean,"popc"yieldsInt32.thread_count (int or Int32 or Uint32, optional) – Number of participating threads. Omit for all CTA threads; otherwise a non-zero multiple of the warp size and consistent across the barrier slot.
- Returns:
For AND/OR: the reduced
Boolean. For POPC: the countInt32.- Raises:
ValueError – if
kindis not one of"and"/"or"/"popc", if a staticbarrier_idis outside[0, 15], or a staticthread_countis not a positive multiple of 32. RuntimeInt32/Uint32values pass through unchecked.
# Every CTA thread reaches this reduction barrier all_true = nvvm.barrier_cta_red_aligned( tx < threads_per_cta, barrier_id=0, kind="and", )
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cta_sync( ) None#
Synchronize threads at a named CTA barrier.
Emits
barrier.cta.sync a{, b};(non-aligned). Works in any control flow, including divergent on sm_70+. All participants at slotbarrier_idwait untilthread_countof them have arrived, then proceed together. Omitthread_countfor the “all CTA threads” rendezvous.nvvm.barrier_cta_sync()with no arguments is the__syncthreads()equivalent (slot 0, all CTA threads).For the aligned variant (
barrier.cta.sync.aligned, equivalent to the legacybar.cta.sync) usebarrier_cta_sync_aligned()— that form promises every CTA thread executes the barrier and is undefined behavior under divergent control flow.- Parameters:
- Raises:
ValueError – if a static
barrier_idis outside[0, 15]or a staticthread_countis not a positive multiple of 32. RuntimeInt32/Uint32values pass through unchecked.
# All CTA threads sync (equivalent to __syncthreads at slot 0) nvvm.barrier_cta_sync(0) # Only warps 2–3 (64 threads) sync at slot 1 nvvm.barrier_cta_sync(1, thread_count=64)
Note
In warp-specialized kernels, a slot-0 all-CTA sync stalls every warp, including idle producer warps. Use a named barrier scoped to the relevant warps, or use per-consumer
mbarriersignals instead.
- cutlass.experimental.primitives.nvvm_wrapper.barrier_cta_sync_aligned( ) None#
Aligned variant of
barrier_cta_sync().Emits
barrier.cta.sync.aligned a{, b};— equivalent to the legacybar.cta.sync. Promises that every CTA thread executes this barrier in convergence; undefined behavior on sm_70+ when a strict subset of CTA threads reaches the instruction. Use this only when the call site is provably all-CTA-converged; otherwise fall back to the non-alignedbarrier_cta_sync().- Parameters:
- Raises:
ValueError – if a static
barrier_idis outside[0, 15]or a staticthread_countis not a positive multiple of 32. RuntimeInt32/Uint32values pass through unchecked.
# Every CTA thread reaches this barrier (no divergent guards above) nvvm.barrier_cta_sync_aligned(0)
- cutlass.experimental.primitives.nvvm_wrapper.breakpoint() None#
Suspend the executing thread for an attached debugger.
Emits
brkpt. Suspends the issuing thread so a debugger can inspect state; it is effectively a no-op when no debugger is attached.nvvm.breakpoint()
- cutlass.experimental.primitives.nvvm_wrapper.cluster_ctarank() Int32#
Read
%cluster_ctarank— this CTA’s linear rank within its cluster.rank = prims.cluster_ctarank()
- cutlass.experimental.primitives.nvvm_wrapper.clusterlaunchcontrol_query_cancel(
- query_type: ClusterLaunchControlQueryType,
- try_cancel_response: int | Int128 | Uint128,
1:1 wrapper over
nvvm.clusterlaunchcontrol_query_cancel.Returns
BooleanforIS_CANCELED,Int32forGET_FIRST_CTA_ID_{X,Y,Z}.
- cutlass.experimental.primitives.nvvm_wrapper.convert(
- src: Vector,
- dst_dtype: object,
- *,
- rnd: FPRoundingMode | None = None,
- sat: SaturationMode | None = None,
- relu: bool | None = None,
- scale_factor: int | Int16 | Uint16 | None = None,
- scale_factor_kind: ConvertScale | None = None,
- random_bits: int | Int32 | Uint32 | None = None,
- result_type: object | None = None,
Convert a packed float vector to
dst_dtype, dispatching on types.A single entry point over the typed
convert_*ops: it reads the element type and lane count ofsrcplus the requesteddst_dtypeand routes to the matching NVVM convert, verifying both the combination and the arguments. Covers the float<->float packed conversions:f32x2->f16/bf16/f8/f6/f4f32x4->f8/f6/f4(requiresrandom_bits)f16x2/bf16x2->f8/f6/f4f8x2/f6x2/f4x2->f16
Exotic / non-float-float conversions (
s2f6, scaledf8 -> bf16, float<->integer,tf32) are not routed here; call the explicitnvvm.convert_*wrapper for those.- Parameters:
src (Vector) – Packed source vector (e.g.
Vector[Float32, 2],Vector[Float8E4M3FN, 2]); its element type and lane count drive dispatch.dst_dtype – Target element type (e.g.
Float16,Float8E4M3FN).rnd – Rounding mode, where the matched convert accepts one.
sat – Saturation mode (f8 / f16 / bf16 narrowing only).
relu – Clamp negatives to zero, where supported.
scale_factor – Block scale, for the scaled narrowing converts.
scale_factor_kind – Scale-factor kind paired with
scale_factor.random_bits – Stochastic-rounding bits; required for
f32x4narrowing and optional forf32x2 -> f16 / bf16.result_type – Packed return-shape override (
Int16/Vector[Int8, 2]) for the narrowing converts that support it.
- Raises:
ValueError – the (source dtype, lane count, destination dtype) triple is not a supported float<->float conversion, an argument is not accepted by the matched convert, or
random_bitsis missing for anf32x4narrowing.
Note
Narrowing returns a packed integer carrier (
Int16/Vector[Int8, N]), whereas widening expects a typed narrow vector (e.g.Vector[Float8E4M3FN, 2], such as one loaded from an FP8 tensor) and bitcasts it to the byte carrier internally. A packed narrowing result is therefore not directly re-widenable: reinterpret it as the typed narrow vector first (carrier.bitcast(<narrow dtype>)).# Narrowing: f32 pair -> packed FP8x2 (an Int16 carrier). packed = nvvm.convert(f32x2, cutlass.Float8E4M3FN, rnd="rn", sat="satfinite") # Widening: a *typed* FP8x2 vector (e.g. loaded from memory) -> f16x2. f16x2 = nvvm.convert(fp8_vec, cutlass.Float16)
- cutlass.experimental.primitives.nvvm_wrapper.convert_and_pack_integer(
- src_a: int | ~cutlass.Int32 | ~cutlass.Uint32,
- src_b: int | ~cutlass.Int32 | ~cutlass.Uint32,
- convert_type: _install_cutlass_mlir_autodoc_stub.<locals>._DocMlirType,
- *,
- src_c: int | ~cutlass.Int32 | ~cutlass.Uint32 | None = None,
- is_signed: bool | None = None,
Wrapper over
nvvm.convert_and_pack_integer.
- cutlass.experimental.primitives.nvvm_wrapper.convert_bf16x2_to_s2f6x2(
- *args: Any,
- **kwargs: Any,
Gated 1:1 wrapper over
nvvm.convert_bf16x2_to_s2f6x2(PTX ISA 9.1).The
.s2f6x2cvt instruction type was introduced in PTX ISA 9.1 and is unavailable on CTK 12.9 (PTX ISA 8.8).
- cutlass.experimental.primitives.nvvm_wrapper.convert_f32x2_to_s2f6x2(
- a: float | ~cutlass.Float32,
- b: float | ~cutlass.Float32,
- *,
- result_type: type[~cutlass.Int16] | type[~cutlass.Vector] = <class 'cutlass.Int16'>,
- scale_factor: int | ~cutlass.Int16 | ~cutlass.Uint16 | None = None,
- relu: bool | None = None,
1:1 wrapper over
nvvm.convert_f32x2_to_s2f6x2.result_type selects the packed return shape:
Int16(default, one i16 with the two converted values packed into the high and low bytes) orVector[Int8, 2](each lane holds one converted value). Both shapes hold equivalent bits.
- cutlass.experimental.primitives.nvvm_wrapper.convert_f8x2_to_bf16x2(
- *args: Any,
- **kwargs: Any,
Gated 1:1 wrapper over
nvvm.convert_f8x2_to_bf16x2(PTX ISA 9.2).The
.bf16x2destination from an.e4m3x2/.e5m2x2source was introduced in PTX ISA 9.2 and is unavailable on CTK 12.9 (PTX ISA 8.8).
- cutlass.experimental.primitives.nvvm_wrapper.convert_float_to_integer(
- src: float | ~cutlass.Float32,
- *,
- result_type: type[~cutlass.Int8] | type[~cutlass.Int32] = <class 'cutlass.Int32'>,
- rnd: ~cutlass.experimental.primitives.nvvm_wrapper.IntRoundingMode | None = None,
- sat: bool | None = None,
- ftz: bool | None = None,
- is_signed: bool | None = None,
1:1 wrapper over
nvvm.convert_float_to_integer.
- cutlass.experimental.primitives.nvvm_wrapper.convert_float_to_tf32(
- src: float | Float32,
- *,
- rnd: FPRoundingMode | None = None,
- sat: SaturationMode | None = None,
- relu: bool | None = None,
Wrapper over
nvvm.convert_float_to_tf32.
- cutlass.experimental.primitives.nvvm_wrapper.convert_s2f6x2_to_bf16x2(
- src: Vector,
- *,
- scale_factor: int | Int16 | Uint16 | None = None,
- sat: SaturationMode | None = None,
- relu: bool | None = None,
1:1 wrapper over
nvvm.convert_s2f6x2_to_bf16x2.
- cutlass.experimental.primitives.nvvm_wrapper.cp_async_bulk_commit_group() None#
Commits all prior initiated but uncommitted cp.async.bulk instructions.
See the PTX documentation.
- dst_mem: Array | Pointer,
- src_mem: Array | Pointer,
- size: int | Int32 | Uint32,
- *,
- l2_cache_hint: int | Int64 | Uint64 | None = None,
- byte_mask: int | Int16 | Uint16 | None = None,
Async bulk-copy a byte range from CTA shared memory to global memory.
Emits the
.shared::cta -> .globalform ofcp.async.bulkwith.bulk_groupcompletion:cp.async.bulk.global.shared::cta.bulk_group [dst], [src], size;. Copies a byte range from SMEM to GMEM without a tensor-map descriptor; useful for 1-D / flat buffers where TMA setup is not justified.sizemust be a positive multiple of 16 and bothdst_memandsrc_memmust be 16-byte aligned (the PTX ISA leaves non-conforming values undefined). The wrapper exposes the full option set this direction has in PTX (unchanged across ISA 8.8 and 9.3): the.L2::cache_hintcache policy (l2_cache_hint) and the.cp_maskbyte mask (byte_mask).- Parameters:
dst_mem (cutlass.Array or cutlass.Pointer) – GMEM destination pointer/array (the
[dst]operand); must be 16-byte aligned.src_mem (cutlass.Array or cutlass.Pointer) – CTA-scope SMEM source pointer/array (the
[src]operand); must be 16-byte aligned.size (int or cutlass.Int32 or cutlass.Uint32) – Number of bytes to copy; must be a positive multiple of 16.
l2_cache_hint (int or cutlass.Int64 or cutlass.Uint64, optional) – Optional 64-bit L2 cache-eviction policy descriptor (emits
.L2::cache_hintwith the policy operand). Defaults to None.byte_mask (int or cutlass.Int16 or cutlass.Uint16, optional) – Optional 16-bit
.cp_maskselecting which bytes of each 16-byte source chunk are written: bit i set copies byte i of every 16-byte chunk, bit i clear skips it. Defaults to None (all bytes copied).
- Raises:
TypeError – if
src_memexposes an address space that is not shared memory.ValueError – if a statically known
sizeis not a positive multiple of 16.
# Drain-on-completion SMEM -> GMEM bulk store of `nbytes` bytes. if nvvm.elect_sync(): nvvm.cp_async_bulk_global_shared_cta(gmem_dst, smem_src, nbytes) nvvm.cp_async_bulk_commit_group() nvvm.cp_async_bulk_wait_group(0)
- cutlass.experimental.primitives.nvvm_wrapper.cp_async_bulk_prefetch(
- src_mem: Array | Pointer,
- size: int | Int32 | Uint32,
- *,
- l2_cache_hint: int | Int64 | Uint64 | None = None,
Prefetch a byte range from global memory into L2.
Emits
cp.async.bulk.prefetch.L2.global [src], size;. Pullssizebytes starting atsrc_meminto L2 without writing any destination; the GMEM access is asynchronous and best-effort.- Parameters:
src_mem (cutlass.Array or cutlass.Pointer) – Global-memory source pointer/array.
size (int or cutlass.Int32 or cutlass.Uint32) – Number of bytes to prefetch; must be a positive multiple of 16.
l2_cache_hint (int or cutlass.Int64 or cutlass.Uint64, optional) – Optional 64-bit L2 cache-eviction policy descriptor.
- Raises:
ValueError – if a statically known
sizeis not a positive multiple of 16.
- dst_mem: Array | Pointer,
- src_mem: Array | Pointer,
- mbar: Array | Pointer,
- size: int | Int32 | Uint32,
- *,
- multicast_mask: Int16 | Int32 | None = None,
- l2_cache_hint: int | Int64 | Uint64 | None = None,
Async bulk-copy a byte range from global memory into cluster shared memory.
Emits
cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [dst], [src], size, [mbar];. Copies a byte range (sizea positive multiple of 16) from GMEM to SMEM without a tensor map; hardware fires the mbarrier’scomplete_txautomatically when the transfer finishes. For per-CTA self-delivery, omitmulticast_mask(or passNone); passing1 << cta_rankstill emits the.multicast::clusterPTX modifier and pays the multicast-routing overhead even though every byte only lands in the issuing CTA. The tensor-descriptor variantcp_async_bulk_tensor_shared_cluster_global()follows the same rule.- Parameters:
dst_mem (cutlass.Array or cutlass.Pointer) – Cluster-scope SMEM destination pointer/array; must be 16-byte aligned. A
shared::ctapointer is auto-cast toshared::cluster(the cluster bulk-copy intrinsic requires it).src_mem (cutlass.Array or cutlass.Pointer) – GMEM source pointer/array; must be 16-byte aligned.
mbar (cutlass.Array or cutlass.Pointer) – Pointer/Array to the SMEM mbarrier signalled on completion.
size (int or cutlass.Int32 or cutlass.Uint32) – Number of bytes to copy (the same value the consumer arms via
arrive_expect_tx); must be a positive multiple of 16.multicast_mask (cutlass.Int16 or cutlass.Int32, optional) – Optional per-bit mask over CTA ranks. Defaults to None: omit it for per-CTA self-delivery (each issuing CTA delivers to itself only). A non-None value emits the
.multicast::clustermodifier, gated on mask presence not value, so1 << cta_rankis a footgun (same delivery as omitted, but pays the multicast-routing overhead). Set it only for genuine cluster broadcast: e.g.3(0b11) on a 2-CTA cluster delivers identical bytes to both CTAs from one issuer.l2_cache_hint (int or cutlass.Int64 or cutlass.Uint64, optional) – Optional 64-bit L2 cache-eviction policy descriptor.
- Raises:
TypeError – if
dst_memis not shared or cluster-shared memory, ormbaris not shared memory.ValueError – if a statically known
sizeis not a positive multiple of 16, or a statically knownmulticast_maskdoes not fit in 32 bits.
# GMEM -> cluster-SMEM bulk load of `nbytes`, signalled on `mbar`. if nvvm.elect_sync(): nvvm.cp_async_bulk_shared_cluster_global(smem_dst, gmem_src, mbar, nbytes)
- dst_mem: Array | Pointer,
- src_mem: Array | Pointer,
- mbar: Array | Pointer,
- size: int | Int32 | Uint32,
Async bulk-copy a byte range between two CTAs’ shared memory.
Emits
cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [dst], [src], size, [mbar];. Copiessizebytes from the issuing CTA’s SMEM (src_mem) to a peer CTA’s SMEM within the same cluster (dst_mem, addressed viashared::cluster); hardware fires the destinationmbarrier’scomplete_txwhen the transfer finishes.- Parameters:
dst_mem (cutlass.Array or cutlass.Pointer) – Cluster-scope SMEM destination pointer/array on the peer CTA; must be 16-byte aligned.
src_mem (cutlass.Array or cutlass.Pointer) – CTA-scope SMEM source pointer/array on the issuing CTA; must be 16-byte aligned.
mbar (cutlass.Array or cutlass.Pointer) – Pointer to the destination CTA’s SMEM mbarrier; signalled on completion.
size (int or cutlass.Int32 or cutlass.Uint32) – Number of bytes to copy (must match the consumer’s
arrive_expect_txcount); positive multiple of 16.
- Raises:
ValueError – if a statically known
sizeis not a positive multiple of 16.
- tma_descriptor: Array | Pointer,
- src_mem: Array | Pointer,
- coordinates: list[int | Int32 | Uint32],
- *,
- l2_cache_hint: int | Int64 | Uint64 | None = None,
- mode: TMAStoreMode | None = None,
Async TMA-store a tile from CTA shared memory to global memory.
Emits
cp.async.bulk.tensor.<N>d.global.shared::cta.bulk_group [tensor_map, {coords}], [src];where<N>is the number ofcoordinates. Stores a CTA-shared tile to the global tensor selected by the TMA descriptor atcoordinates. Completion uses the bulk-group mechanism (no mbarrier signal), so the store must be drained withcp_async_bulk_commit_group()andcp_async_bulk_wait_group().- Parameters:
tma_descriptor (cutlass.Array or cutlass.Pointer) – TMA tensor-map descriptor for the destination global tensor.
src_mem (cutlass.Array or cutlass.Pointer) – CTA-scope SMEM source pointer/array.
coordinates (list of (int or cutlass.Int32 or cutlass.Uint32)) – 1-5D tile coordinate into the descriptor’s tensor.
l2_cache_hint (int or cutlass.Int64 or cutlass.Uint64, optional) – Optional 64-bit L2 cache-eviction policy descriptor.
mode (TMAStoreMode, optional) – Optional TMA store mode (
tiledefault,im2col, ortile_scatter4).
- Raises:
TypeError – if
src_memdoes not reside in shared memory.ValueError – if the
coordinatescount is invalid formode(1-5 for tile, 3-5 for im2col, exactly 2 for scatter4).
Reverse of
cp_async_bulk_tensor_shared_cta_global()(load). TMA hardware reads SMEM and writes the tile to the global tensor described bytma_descriptorat the givencoordinates. Argument order is(desc, smem, coords)— no mbarrier; TMA stores use commit/wait groups instead of mbarriers.Proxy fence required before issue: thread SMEM writes go through the “generic” proxy; the TMA engine reads through the “async” proxy. Issue
nvvm.fence_proxy("async_shared", space=SharedSpace.shared_cta)before this call so the TMA engine sees the latest SMEM data. Without it, TMA may read stale SMEM.# TMA store of an SMEM tile to GMEM at (x, y), then drain. if nvvm.elect_sync(): nvvm.cp_async_bulk_tensor_global_shared_cta(tma_desc, smem_src, [x, y]) nvvm.cp_async_bulk_commit_group() nvvm.cp_async_bulk_wait_group(0)
- cutlass.experimental.primitives.nvvm_wrapper.cp_async_bulk_tensor_prefetch(
- tma_descriptor: Array | Pointer,
- coordinates: list[int | Int32 | Uint32],
- im2col_offsets: list[int | Int16 | Uint16],
- *,
- mode: TMALoadMode | None = None,
- l2_cache_hint: int | Int64 | Uint64 | None = None,
Prefetch a TMA-tensor tile from global memory into L2.
Emits
cp.async.bulk.prefetch.tensor.<N>d.L2.global [tma_desc, {coords}];. Pulls the tile selected bycoordinates(andim2col_offsetsfor im2col mode) into L2 without any destination; the access is asynchronous and best-effort.- Parameters:
tma_descriptor (cutlass.Array or cutlass.Pointer) – TMA tensor-map descriptor for the global tensor.
coordinates (list of (int or cutlass.Int32 or cutlass.Uint32)) – 1-5D tile coordinate into the descriptor’s tensor.
im2col_offsets (list of (int or cutlass.Int16 or cutlass.Uint16)) – Im2col offsets; empty list for tile mode.
mode (TMALoadMode, optional) – Optional TMA load mode (
tiledefault,im2col).l2_cache_hint (int or cutlass.Int64 or cutlass.Uint64, optional) – Optional 64-bit L2 cache-eviction policy descriptor.
- Raises:
ValueError – if the
coordinatescount is invalid formode.
The descriptor-override path is exposed separately as
cp_async_bulk_tensor_prefetch_override().
- cutlass.experimental.primitives.nvvm_wrapper.cp_async_bulk_tensor_reduce(
- tma_descriptor: Array | Pointer,
- src_mem: Array | Pointer,
- red_kind: TMARedux,
- coordinates: list[int | Int32 | Uint32],
- *,
- mode: TMAStoreMode | None = None,
- l2_cache_hint: int | Int64 | Uint64 | None = None,
Issue a TMA async tensor reduction from shared memory to global memory.
Lowers to PTX
cp.reduce.async.bulk.tensor. The source tile in shared::cta memory is reduced into the destination tensor described bytma_descriptoratcoordinatesusingred_kind. The tensor-map element type determines which reduction kinds are valid; PTX supportsADDfor integer/fp elements,MIN/MAXfor integer and fp16/bf16 element types,INC/DECforu32, and bitwise reductions forb32/b64.The operation is non-blocking and uses bulk async-group completion. It does not take an mbarrier operand; issue the operation, commit the bulk group, then wait on the group before consuming completion.
- Parameters:
tma_descriptor (Array or Pointer) – Tensor-map descriptor for the global destination.
src_mem (Array or Pointer) – Shared-memory source tile.
red_kind (TMARedux) – TMA reduction operation, such as
ADD,MIN,MAX,INC,DEC,AND,OR, orXOR.coordinates (list[int or Int32 or Uint32]) – One to five
s32tensor coordinates matching the descriptor rank.mode (TMAStoreMode, optional) – Optional TMA store mode. Omit for tile mode; im2col modes require descriptor-compatible ranks and coordinates.
l2_cache_hint (int or Int64 or Uint64, optional) – Optional 64-bit L2 cache policy.
- Raises:
TypeError – if
src_memdoes not reside in shared memory.ValueError – if
red_kindis not a validTMARedux, ifmodeistile_scatter4(unsupported by reduce), or if thecoordinatescount is invalid formode(1-5 for tile, 3-5 for im2col).
if nvvm.elect_sync(): nvvm.cp_async_bulk_tensor_reduce( desc, smem_tile, "add", [row, col], ) nvvm.cp_async_bulk_commit_group() nvvm.cp_async_bulk_wait_group(0)
- dst_mem: Array | Pointer,
- tma_descriptor: Array | Pointer,
- coordinates: list[int | Int32 | Uint32],
- mbar: Array | Pointer,
- im2col_offsets: list[int | Int16 | Uint16],
- *,
- multicast_mask: Int16 | Int32 | None = None,
- l2_cache_hint: int | Int64 | Uint64 | None = None,
- mode: TMALoadMode | None = None,
- group: CTAGroup | None = None,
Issue a TMA async load into shared memory with optional cluster multicast.
Like
cp_async_bulk_tensor_shared_cta_globalbut supports multicast across all CTAs in a cluster (group="cta_2").Arg order:
(dst_smem, tma_desc, coords, mbar, im2col_offsets, ...)— coords come before mbar. This is the opposite ofcp_async_bulk_tensor_shared_cta_global(mbar before coords) and is the #1 footgun when porting CTA_1 code to CTA_2.CTA_2 mbar routing — bit 24 maps to 2-SM-group leader, NOT cluster leader. When
group="cta_2", this wrapper masks the mbar pointer with& 0xFEFFFFFF(clears bit 24 only). Bit 24 holds the LSB of%cluster_ctarank— clearing it routes the address from a peer CTA back to its 2-SM-group leader (the even-rank CTA in the same pair).In a 2-CTA cluster (
cluster_shape=(2,1,1)), the only group leader is cluster rank 0, so this is equivalent to “route to cluster leader” and is what2cta_mma_basic.pyrelies on.In multi-group clusters (
cluster_shape=(2, n, 1)with n > 1), complete_tx still routes only to the issuer’s group leader (the even-rank CTA in the issuer’s pair). It does NOT route across groups: a TMA issued in group 0 cannot delivercomplete_txto group 1’s mbar through this wrapper. Cross- group multicast TMA (e.g. broadcast B to all groups) requires manually constructing a cluster-shared mbar pointer viamapa.shared::clusterso all issuers target the same cluster-leader mbar; this wrapper does not expose that path.
Shared-tile multicast vs per-CTA unicast — two distinct topologies, commonly confused. The right choice comes from the pseudo-code:
Shared-tile multicast (
multicast @ctain the pseudo-code on a tile declared at the cluster scale, e.g.tile B_smem: f16:K×N @ SMEMwithN = N_TILEcovering the full cluster width). Each CTA gets an identical copy of one full tile. One descriptor withbox_dimsspanning the whole tile;multicast_maskincludes every CTA in the cluster (0b11for 2-CTA,0xfffor 8-CTA). Mbar protocol depends on cluster_size:cluster_size == 2 (with downstream MMA reuse): optionally pass
group="cta_2"to enable bit-24 mbar collapse — only the leader inits + doesarrive_expect_txwith the full-tile byte count, and only the leader waits. All CTAs’complete_txis redirected to the leader’s mbar. This is the right shape when the same mbar is then reused as the MMA’s input-ready bar.cluster_size ≥ 2 (general, including > 2): omit
group=and use per-CTA local mbars — every CTA inits its own mbar (at the same SMEM offset), every CTA doesarrive_expect_txwith the per-CTA tile byte count, only the leader issues the TMA, every CTA waits on its own mbar. Per the PTX ISA, the hardware multicastscomplete_txto that same SMEM offset in every destination CTA’s local SMEM. Works for cluster_size ∈ {2, 4, 8, 16}.
Trying
group="cta_2"with cluster_size > 2 hangs: the bit-24 mask redirects to the issuer’s 2-SM-group leader only (cluster ranks 0+1), so cross-group receivers’ mbars never arrive.Per-CTA unicast (used when each CTA consumes a different slice of the tensor, e.g. each CTA gets a different half of the M-dimension). Descriptor
box_dimscovers only the per-CTA slice, omit ``multicast_mask`` entirely (or passNone), coordinates shift bycta_rank * sliceper CTA, and each CTA callsmbarrier_arrive_expect_txindependently. This is functionally the CTA_1 pattern wrapped ingroup=CTA_2for the routing bookkeeping; it is not a multicast.
Multicast modifier rule — the
multicast::clusterPTX modifier is gated by whether ``multicast_mask`` is present, not by the value the mask carries. A “selfcast” mask such as1 << cta_ranktherefore still emits the modifier and pays the multicast-routing overhead even though the bytes only land in the issuing CTA. For per-CTA unicast topology, omitmulticast_mask. Setmulticast_maskonly when the topology is genuinely shared-tile multicast (one issuer, mask covers every receiving CTA). Applies to bothgroup=CTA_1andgroup=CTA_2.Do not split a shared-tile multicast into per-CTA unicasts to “save bytes” — the multicast is a single HBM read fanned out over the cluster interconnect, so per-CTA unicast increases HBM pressure instead of reducing it, in addition to producing a different SMEM layout than the spec.
# Shared-tile multicast — both CTAs see the same 128×128 tile. # arrive_expect_tx is called by the leader only, with the FULL tile's # byte count (not per-CTA). if is_leader: nvvm.mbarrier_arrive_expect_tx(mbar + s, A_full_bytes + B_full_bytes) if nvvm.elect_sync(): nvvm.cp_async_bulk_tensor_shared_cluster_global( smem_B + s * tile_b, tma_b_desc, (k, n), # full tile origin mbar + s, [], # ← coords BEFORE mbar multicast_mask=Int16(0b11), # every CTA in the 2-CTA cluster group="cta_2", )
- Parameters:
dst_mem (cutlass.Array or cutlass.Pointer) – Cluster-scope SMEM destination tile (shared or cluster-shared); must be 16-byte aligned.
tma_descriptor (cutlass.Array or cutlass.Pointer) – TMA tensor-map descriptor for the source tensor.
coordinates (list of (int or cutlass.Int32 or cutlass.Uint32)) – 1-5D tile coordinate into the descriptor’s tensor (note: coords come BEFORE
mbarin the arg order).mbar (cutlass.Array or cutlass.Pointer) – Shared-memory mbarrier signalled on completion.
im2col_offsets (list of (int or cutlass.Int16 or cutlass.Uint16)) – im2col offsets (empty list for tile mode).
multicast_mask (cutlass.Int16 or cutlass.Int32, optional) – Optional per-bit CTA-rank mask for cluster multicast (omit for per-CTA unicast).
l2_cache_hint (int or cutlass.Int64 or cutlass.Uint64, optional) – Optional 64-bit L2 cache-eviction policy descriptor.
mode (TMALoadMode, optional) – Optional TMA load mode (tile default, or im2col).
group (CTAGroup, optional) – CTA group selector (
cta_1default, orcta_2).
- Raises:
TypeError – if
dst_memis not shared/cluster-shared memory ormbaris not shared memory.ValueError – if the
coordinatescount is invalid formode, or a statically knownmulticast_maskdoes not fit in 32 bits.
- dst_mem: Array | Pointer,
- tma_descriptor: Array | Pointer,
- coordinates: list[int | Int32 | Uint32],
- mbar: Array | Pointer,
- im2col_offsets: list[int | Int16 | Uint16] | None = None,
- *,
- l2_cache_hint: int | Int64 | Uint64 | None = None,
- mode: TMALoadMode | None = None,
Issue a TMA async load from global memory into this CTA’s shared memory.
TMA hardware performs the DMA without stalling the issuing warp. Completion is signaled by an mbarrier
complete_txdecrement, which fires the barrier once all bytes have arrived.Arg order:
(dst_smem, tma_desc, coords, mbar)— coords before mbar, matchingcp_async_bulk_tensor_shared_cluster_global()so porting CTA_1 code to CTA_2 only changes the function name, not the argument order.Calling convention:
Call from exactly one thread (e.g.
if nvvm.elect_sync()).Call
arrive_expect_tx(mbar, nbytes)before this function so the transaction counter is set before TMA can decrement it.nbytes= rows × cols × sizeof(dtype) for the tile being loaded.Coordinates are tensor-space element indices, not byte offsets. For a 2-D descriptor created via
create_tensor_map_tiled_from_viewon a row-major (M, K) tensor, the TMA coord order is column-major:(k_offset, m_offset)— K (innermost) first.
# Separate elect_sync for arrive vs each TMA load (performance) if nvvm.elect_sync(): nvvm.mbarrier_arrive_expect_tx(mbar + s, A_bytes + B_bytes) if nvvm.elect_sync(): nvvm.cp_async_bulk_tensor_shared_cta_global( smem_A + s * tile_a, tma_a_desc, (k, m), mbar + s) if nvvm.elect_sync(): nvvm.cp_async_bulk_tensor_shared_cta_global( smem_B + s * tile_b, tma_b_desc, (k, n), mbar + s)
For multicast / cluster TMA (CTA_2) use
cp_async_bulk_tensor_shared_cluster_globalinstead.modeselects the TMA access pattern (TILEdefault; also theIM2COLfamily andTILE_GATHER4). For im2col modes pass the per-dimim2col_offsets; tile and gather4 modes leave it empty.- Parameters:
dst_mem (cutlass.Array or cutlass.Pointer) – This CTA’s SMEM destination tile; must be 16-byte aligned.
tma_descriptor (cutlass.Array or cutlass.Pointer) – TMA tensor-map descriptor for the source tensor.
coordinates (list of (int or cutlass.Int32 or cutlass.Uint32)) – 1-5D tile coordinate into the descriptor’s tensor (coords come before
mbarin the arg order).mbar (cutlass.Array or cutlass.Pointer) – Shared-memory mbarrier signalled on completion.
im2col_offsets (list of (int or cutlass.Int16 or cutlass.Uint16), optional) – im2col offsets (empty / omitted for tile mode).
l2_cache_hint (int or cutlass.Int64 or cutlass.Uint64, optional) – Optional 64-bit L2 cache-eviction policy descriptor.
mode (TMALoadMode, optional) – Optional TMA load mode (
TILEdefault, im2col, or gather4).
- Raises:
TypeError – if
dst_memormbaris not shared memory.ValueError – if the
coordinatescount is invalid formode.
- cutlass.experimental.primitives.nvvm_wrapper.cp_async_bulk_wait_group(
- group: cutlass.cute.typing.Int,
- *,
- read: bool | None = None,
Waits till only a specified numbers of cp.async.bulk groups are pending.
See the PTX documentation.
- cutlass.experimental.primitives.nvvm_wrapper.cp_async_commit_group() None#
Commits all prior initiated but uncommitted cp.async instructions.
See the PTX documentation.
- cutlass.experimental.primitives.nvvm_wrapper.cp_async_mbarrier_arrive(
- addr: Array | Pointer,
- *,
- noinc: bool | None = None,
Tie outstanding
cp.asyncoperations to an mbarrier.Emits
cp.async.mbarrier.arrive[.noinc][.shared{::cta}]. Causes an asynchronous arrive-on operation to fire on the mbarrier ataddronce all priorcp.asyncops issued by the executing thread have completed. This lets a consumer block on the mbarrier and only wake when the cp.async data is ready, without polling.Two semantic flavours via
noinc:noinc=False(default) — pending count is incremented by 1 before the asynchronous arrive, giving a net-zero effect on the pending count for the current phase.mbarrier.initonly needs to account formbarrier.arrivearrivals.noinc=True— no pre-increment. The asynchronous arrive decrement must be pre-accounted for inmbarrier.init’s thread-count. Use when issuing manycp.asyncoperations and aggregating them into a single arrival.
In a
cp_async_shared_globalpipeline, this is the producer’s completion signal: consumers wait the mbarrier, notcp_async_wait_group. Count the lane-level async-arrives that will be delivered to the mbarrier. For example, a hybrid TMA + per-thread cp.async producer commonly initializes the full barrier with1 + 32: one elected-thread TMA arrival plus onecp_async_mbarrier_arrive(noinc=True)from each lane in a cp.async warp. The TMA transaction byte count should not include the cp.async bytes because these copies complete through the async arrive path.cp_async_mbarrier_arriveis CTA-local. In CTA_2 kernels where a leader CTA’s barrier gates collective MMA, have each CTA’s cp.async lanes arrive on a local mbarrier, then use a completion-forwarder warp to wait the local mbarrier and cross-CTA arrive on the leader barrier with theshared::clusterpointer returned bymapa.# One TMA elected-thread arrival plus 32 cp.async lane arrivals. if nvvm.elect_sync(): nvvm.mbarrier_init(full_mbar + stage, 1 + 32) # Each cp.async lane issues its copies, then contributes one # asynchronous arrive after its prior cp.async operations retire. nvvm.cp_async_shared_global( sfa_smem, sfa_gmem, 8, "ca", ) nvvm.cp_async_mbarrier_arrive(full_mbar + stage, noinc=True)
- dst: Array | Pointer,
- src: Array | Pointer,
- size: int,
- modifier: LoadCacheModifier,
- *,
- cp_size: int | Int32 | Uint32 | None = None,
Issue a per-thread async copy from global to shared memory (SM80+).
Each thread independently copies
sizebytes fromsrc(GMEM) todst(SMEM) without stalling. Unlike TMA, no descriptor is required and every participating thread issues its own copy.- Parameters:
dst (cutlass.Array or cutlass.Pointer) – Destination in shared memory (addr-space 3).
src (cutlass.Array or cutlass.Pointer) – Source pointer in global memory.
size (int) – Bytes per thread. Must be
4,8, or16. Use16(128-bit) for maximum throughput (onefloat4).modifier (LoadCacheModifier) – Cache policy —
"ca"(cache L1+L2, 4/8/16 B) or"cg"(bypass L1, L2 only, 16 B only). Prefer"cg"for streaming loads that won’t be reused.cp_size (int or cutlass.Int32 or cutlass.Uint32, optional) – Source byte count for a masked (zero-fill) copy. When
cp_size < size, the bytes[cp_size, size)indstare zeroed rather than left undefined. Passcp_size=0to write all zeros (useful for out-of-bounds boundary tiles). LeaveNonefor full-size copies.
- Raises:
TypeError – if
dstis not shared memory.ValueError – if
sizeis not 4/8/16, ifmodifieriscgwithsize!= 16, or ifcp_sizefalls outside[0, size].
Synchronization: copies are asynchronous. Use
nvvm.cp_async_commit_group()to mark a batch andnvvm.cp_async_wait_group(n)to drain until ≤ngroups remain.Swizzle requirement for tcgen05.mma: when the SMEM tile will be read by
tcgen05.mma, use the SMEM layout expected by the corresponding matrix descriptor. For the common 128B XOR layout, produce the same layout thatPointer.store_swizzledor a 128B-swizzled tensor map would create.# 16-byte streaming load per thread, bypass L1 nvvm.cp_async_shared_global(smem_dst, gmem_src, 16, "cg") nvvm.cp_async_commit_group() # ... later ... nvvm.cp_async_wait_group(0) # wait for all groups
- cutlass.experimental.primitives.nvvm_wrapper.cp_async_wait_group(n: cutlass.cute.typing.Int) None#
Waits till only a specified numbers of cp.async groups are pending.
See the PTX documentation.
- dst: Array | Pointer,
- src: Array | Pointer,
- size: int | Int32 | Uint32,
- *,
- op: CpReduceOp = CpReduceOp.ADD,
- type: CpReduceType = CpReduceType.BF16,
- noftz: bool = False,
- l2_cache_hint: int | Int64 | Uint64 | None = None,
cp.reduce.async.bulk.global.shared::cta— non-TMA bulk reduction.Asynchronously reduces size bytes from src (shared::cta) into dst (global) using op / type. Unlike
cp_async_bulk_tensor_reduce(), this operates on raw pointers and a byte count (irregular / scatter access, e.g. MoE finalize scatter-reduce). Usesbulk_groupcompletion — bracket withcp_async_bulk_commit_group()/cp_async_bulk_wait_group().There is no NVVM dialect op for this instruction, so it emits inline PTX.
- Parameters:
size – byte count, must be a multiple of 16.
noftz – disable flush-to-zero; only valid with
op=ADDandtypein{F16, BF16}.l2_cache_hint – optional 64-bit L2 eviction policy.
- cutlass.experimental.primitives.nvvm_wrapper.cvta_to(
- addr: Array | Pointer | Int32 | Int64 | Uint32 | Uint64,
- space: CvtaSpace,
- *,
- size: CvtaSize = CvtaSize.U64,
cvta.to.{space}— convert a generic address to a space-specific one.For
Array/Pointerinputs this is anllvm.addrspacecastto the target space (same wrapper type returned); integer inputs round-trip throughinttoptr/addrspacecast/ptrtoint..paramspaces have no addrspacecast and use inline PTX.
- cutlass.experimental.primitives.nvvm_wrapper.cvt_f32x2_to_f4x2(
- a: float | Float32,
- b: float | Float32,
- dst_type: object,
- *,
- is_pzo: bool | None = None,
- scale_factor: int | Int16 | Uint16 | None = None,
- scale_factor_kind: ConvertScale | None = None,
- rnd: FPRoundingMode | None = None,
- relu: bool | None = None,
cvt.{rnd}.{f4x2}.f32— convert anf32pair to packedf4x2.Returns the packed byte in the low 8 bits of an
Int32(callers typically& 0xFFand shift it into a 32-bit word).
- cutlass.experimental.primitives.nvvm_wrapper.cvt_f32x2_to_f8x2(
- a: float | Float32,
- b: float | Float32,
- dst_ty: object,
- *,
- is_pzo: bool | None = None,
- scale_factor: int | Int16 | Uint16 | None = None,
- scale_factor_kind: ConvertScale | None = None,
- rnd: FPRoundingMode | None = None,
- sat: SaturationMode | None = None,
- relu: bool | None = None,
cvt.{rnd}.{f8x2}.f32— convert anf32pair to packedf8x2.Returns the two f8 lanes packed into an
Int16.
- cutlass.experimental.primitives.nvvm_wrapper.cvt_packfloat(
- src_a: int | Int32 | Uint32,
- src_c: int | Int32 | Uint32,
- from_: CVTPackFloat,
- to: CVTPackFloat,
- *,
- rnd: FPRoundingMode | None = None,
- sat: SaturationModeKind | None = None,
- relu: bool | None = None,
- extract_hi: bool | None = None,
Wrapper over
nvvm.cvt_packfloat.
- cutlass.experimental.primitives.nvvm_wrapper.cvt_packfloat_f32(
- src_a: float | Float32,
- src_b: float | Float32,
- src_c: int | Int32 | Uint32,
- to: CVTPackFloat,
- *,
- rnd: FPRoundingMode | None = None,
- sat: SaturationModeKind | None = None,
- relu: bool | None = None,
- extract_hi: bool | None = None,
Wrapper over
nvvm.cvt_packfloat_f32.
- cutlass.experimental.primitives.nvvm_wrapper.dot_accumulate_2way(
- a: Vector,
- a_type: DotAccumulateType,
- b: Vector,
- b_type: DotAccumulateType,
- c: int | Int32 | Uint32,
- b_hi: bool,
Wrapper over
nvvm.dot_accumulate_2way.
- cutlass.experimental.primitives.nvvm_wrapper.dot_accumulate_4way(
- a: Vector,
- a_type: DotAccumulateType,
- b: Vector,
- b_type: DotAccumulateType,
- c: int | Int32 | Uint32,
Wrapper over
nvvm.dot_accumulate_4way.
- cutlass.experimental.primitives.nvvm_wrapper.elect_sync( ) Boolean#
Elect one lane from a warp-convergent group.
Exactly one predicated active lane in
membermaskis elected. That lane receivesTrue; every other participating lane receivesFalse. PTX guarantees deterministic election for a fixed member mask, but does not specify which lane ID wins. Use this as the compiler-visible gate for non-idempotent single-issuer operations such as one TMA issuer per warp, onembarrier_arrive_expect_txcall, or onetcgen05_commitcall.Convergence requirement: every executing lane must be named in
membermask, and all lanes named bymembermaskmust actively execute the instruction. Lanes outside the mask should branch around the call.- Parameters:
membermask (int or Int32 or Uint32) – 32-bit warp participation mask; every set bit identifies a lane that must be executing this instruction. Defaults to
FULL_MASK(0xFFFFFFFF, all 32 lanes).- Returns:
Truein the elected lane;Falsein every other participating lane.- Return type:
- Raises:
ValueError – if
membermaskis a Pythonintoutside[0, 0xFFFFFFFF].
# Elect one lane per warp to perform a non-idempotent op. if nvvm.elect_sync(): # uses FULL_MASK by default nvvm.mbarrier_arrive_expect_tx(full_bar + s, tile_bytes)
- cutlass.experimental.primitives.nvvm_wrapper.exit() None#
Terminate the issuing thread’s execution of the kernel.
Emits PTX
exit: the issuing thread stops running immediately. Other threads in the warp / CTA continue. Rarely needed: a natural return from@cute.kernelis almost always preferable, becauseexitdoes not cooperate with mbarrier-based control-flow staging. CTA barriers (bar.sync/barrier.cta) exclusively waiting on arrivals from exited threads are released automatically by hardware (a PTX guarantee), butmbarrierarrivals are explicit and are NOT auto-completed: a thread that exits before itsmbarrier.arriveleaves any consumer of that arrival hung indefinitely.# Rarely needed — prefer an ``if`` guard around the work. if tid >= num_valid: nvvm.exit() # ... remaining threads continue ...
- cutlass.experimental.primitives.nvvm_wrapper.fence_acq_rel(scope: MemScope) None#
fence.acq_rel.{scope}— acquire-release memory fence (PTX §9.7.13.4).There is no NVVM dialect op for the generic acquire-release fence, so this emits inline PTX.
- Parameters:
scope – memory scope —
MemScope.{CTA,CLUSTER,GPU,SYS}.
prims.fence_acq_rel(prims.MemScope.CTA)
- cutlass.experimental.primitives.nvvm_wrapper.fence_mbarrier_init() None#
Make mbarrier init writes visible to all threads before first use.
PTX address-space semantics require an explicit fence between writing an mbarrier object (via
mbarrier_init()) and the first arrive or wait on that barrier. Without this fence the init write may not be visible to threads in a different warp or address space. The fence itself is a per-thread no-op; visibility comes from the CTA-wide sync that follows.# One elected thread initialises all stages. if warp_idx == 0: if nvvm.elect_sync(): for i in cutlass.range_constexpr(NUM_STAGES): nvvm.mbarrier_init(full_bar + i, 1) nvvm.mbarrier_init(empty_bar + i, 1) # Make init visible before any warp calls arrive/wait. nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync()
- cutlass.experimental.primitives.nvvm_wrapper.fence_proxy(
- kind: ~cutlass.experimental.primitives.nvvm_wrapper.Proxy,
- *,
- space: _install_cutlass_mlir_autodoc_stub.<locals>._DocDialectObject | None = None,
Order writes across memory proxy domains.
Emits
fence.proxy.{kind}. The most common use case isfence_proxy("async_shared", space=SharedSpace.shared_cta)before a TMA store (cp_async_bulk_tensor_global_shared_cta), which makes thread SMEM writes visible to the TMA async-copy engine.Without this fence, the TMA store reads stale SMEM data because regular thread stores go through a different memory proxy domain than TMA async copies.
- Parameters:
kind (Proxy) – Proxy kind. Use
"async_shared"for TMA store fencing (SMEM to global via TMA).space (SharedSpace | None) – Shared-memory space qualifier. Only valid with the
"async"/"async_shared"proxy kinds. UseSharedSpace.shared_ctafor CTA-local SMEM.
- Raises:
ValueError –
spaceis supplied together with a proxy kind that does not accept a space qualifier (anything other than"async"/"async_shared").TypeError –
kindis a raw NVVMProxyKinddialect enum; pass aProxymember or its string alias instead.
# Thread writes to SMEM, then TMA stores SMEM -> global: nvvm.barrier_cta_sync() # all threads done writing SMEM nvvm.fence_proxy( "async_shared", space=nvvm.SharedSpace.shared_cta, ) if nvvm.elect_sync(): nvvm.cp_async_bulk_tensor_global_shared_cta( tma_descriptor=tma_desc, src_mem=smem, coordinates=dst_coords, ) nvvm.cp_async_bulk_commit_group() nvvm.cp_async_bulk_wait_group(0)
- cutlass.experimental.primitives.nvvm_wrapper.fence_proxy_acquire(
- scope: MemScope,
- addr: Array | Pointer,
- size: int | Int32 | Uint32,
- *,
- from_proxy: Proxy | None = None,
- to_proxy: Proxy | None = None,
Acquire a memory region modified in another proxy.
Emits the uni-directional acquire proxy fence
fence.proxy.{to_proxy}::{from_proxy}.acquire.{scope} [addr], size. The canonical use is acquiring a tensormap (to_proxy="tensormap",from_proxy="generic") that was edited in the generic proxy (e.g. viatensormap.replace) before reading it through the tensormap proxy with a TMA copy. Unlike most proxies, the tensormap proxy is not acquired from the generic proxy at kernel start, so this explicit fence is required whenever a tensormap is modified at runtime.- Parameters:
scope (MemScope) – Scope at which the prior writes become visible (
"cta"/"cluster"/"gpu"/"sys").addr (Array | Pointer) – Base address of the region being acquired (e.g. the tensormap object). Generic-addressed; the runtime address must fall within the
.globalwindow (not enforced at trace time).size (int | Int32 | Uint32) – Size of that region in bytes. The only value the instruction supports is
128(the tensormap size), and it must be an immediate.from_proxy (Proxy | None) – Source proxy the writes were performed in. Only
"generic"is valid (the default); leave asNoneto use it.to_proxy (Proxy | None) – Target proxy the subsequent reads use. Only
"tensormap"is valid (the default); leave asNoneto use it.
- Raises:
ValueError – a static
intsizeother than128; an explicitfrom_proxyother than"generic"; or an explicitto_proxyother than"tensormap".
# Acquire a runtime-edited tensormap before using it in a TMA copy: nvvm.fence_proxy_acquire( "gpu", tma_desc, 128, from_proxy="generic", to_proxy="tensormap", ) if nvvm.elect_sync(): nvvm.cp_async_bulk_tensor_shared_cta_global( smem, tma_desc, src_coords, mbar )
- cutlass.experimental.primitives.nvvm_wrapper.fence_proxy_async_acquire_sync_restrict() None#
fence.proxy.async::generic.acquire.sync_restrict::shared::cluster.cluster.Lowers to the
nvvm.fence.proxy.sync_restrictop withacquireorder; per its definition,acquirerestricts the ordering toshared::clusterbetween the generic and async proxies.
- cutlass.experimental.primitives.nvvm_wrapper.fence_proxy_async_release_sync_restrict() None#
fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster.Lowers to the
nvvm.fence.proxy.sync_restrictop withreleaseorder; per its definition,releaserestricts the ordering toshared::ctawith cluster scope between the generic and async proxies.
- cutlass.experimental.primitives.nvvm_wrapper.fence_proxy_release( ) None#
Release a memory region to another proxy.
Emits the uni-directional release proxy fence
fence.proxy.{to_proxy}::{from_proxy}.release.{scope}. It is the release counterpart offence_proxy_acquire(): afence.proxy.releaseforms a release sequence that synchronises with an acquire sequence containing a matchingfence.proxy.acquire. The canonical use is publishing a tensormap (to_proxy="tensormap",from_proxy="generic") edited in the generic proxy before another agent acquires and reads it through the tensormap proxy. Unlikefence_proxy_acquire(), the release form takes no address window.- Parameters:
scope (MemScope) – Scope at which the prior writes are released (
"cta"/"cluster"/"gpu"/"sys").from_proxy (Proxy | None) – Source proxy the writes were performed in. Only
"generic"is valid (the default); leave asNoneto use it.to_proxy (Proxy | None) – Target proxy the subsequent reads use. Only
"tensormap"is valid (the default); leave asNoneto use it.
- Raises:
ValueError – an explicit
from_proxyother than"generic"or an explicitto_proxyother than"tensormap".
# Publish a runtime-edited tensormap to the tensormap proxy: nvvm.fence_proxy_release( "gpu", from_proxy="generic", to_proxy="tensormap", )
- cutlass.experimental.primitives.nvvm_wrapper.fence_proxy_sync_restrict( ) None#
Order memory between the async and generic proxies, sync-restricted.
Emits
fence.proxy.async::generic.{order}.sync_restrict.... Thesync_restrictqualifier narrows the ordering so thatacquireapplies toshared::clusterandreleasetoshared::cta, both at cluster scope. Ordering is supported only between the async and generic proxies.- Parameters:
- Raises:
ValueError –
orderis not"acquire"/"release"; an explicitfrom_proxyother than"generic"; or an explicitto_proxyother than"async".
nvvm.fence_proxy_sync_restrict("acquire")
- cutlass.experimental.primitives.nvvm_wrapper.fence_sc_cluster() None#
Sequentially-consistent memory fence at cluster scope.
Emits
fence.sc.cluster. Contributes this thread’s prior memory accesses to a single total order over the sequentially-consistent operations observed by all threads in the cluster, ordering them before the thread’s subsequent accesses with respect to the whole cluster.nvvm.fence_sc_cluster()
- cutlass.experimental.primitives.nvvm_wrapper.fence_sync_restrict(
- order: MemOrder,
Thread fence restricted to a memory class at cluster scope.
Emits
fence.{order}.sync_restrict::shared::{cluster|cta}.cluster. Thesync_restrictqualifier restrictsacquireordering toshared::clusterandreleasetoshared::cta, both at cluster scope.- Parameters:
order (MemOrder) –
"acquire"or"release".- Raises:
ValueError –
orderis not"acquire"/"release".
nvvm.fence_sync_restrict("release")
- cutlass.experimental.primitives.nvvm_wrapper.fma_packed_f32x2(
- src_a: tuple | Vector,
- src_b: tuple | Vector,
- src_c: tuple | Vector,
- *,
- rnd: FPRoundingMode | None = None,
- ftz: bool | None = None,
Wrapper over
nvvm.fma_packed_f32x2.Accepts a 2-tuple of f32 scalars or a
Vectorfor each operand and returns a tuple when called with tuples, else aVector.
- cutlass.experimental.primitives.nvvm_wrapper.griddepcontrol(
- kind: GridDepAction,
Coordinate execution between consecutive dependent grids.
Emits
griddepcontrol.{launch_dependents|wait}. Used between back-to-back kernels that the runtime has wired with a producer/consumer dependency: the producer kernel issueslaunch_dependentsto let the dependent kernel start as soon as scheduling permits, and the dependent kernel issueswaitto ensure all prerequisite-grid memory operations have drained before reading.- Parameters:
kind (GridDepAction) –
"launch_dependents"(producer side, hint that dependents may start) or"wait"(consumer side, block until prerequisites done).- Raises:
ValueError – if
kindis not aGridDepActionmember or its string value (raw NVVM dialect enums are rejected).
# Producer kernel — at the end, allow the consumer kernel to start if nvvm.elect_sync(): nvvm.griddepcontrol("launch_dependents") # Consumer kernel — at the start, wait for producer's writes if nvvm.elect_sync(): nvvm.griddepcontrol("wait")
- cutlass.experimental.primitives.nvvm_wrapper.inline_ptx_hl(
- ptx_code: str,
- *,
- write_only_types: list | None = None,
- read_only_args: list | None = None,
- read_write_args: list | None = None,
- pred: Boolean | None = None,
Public high-level inline-PTX builder (
{$r0}/{$w0}named refs, DSLwrite_only_types, built-in@ppredication viapred=). Distinct frominline_ptx, which is the rawnvvm.inline_ptxop (write_only_args/ positionalptx_code).
- cutlass.experimental.primitives.nvvm_wrapper.ldmatrix(
- ptr: Array | Pointer,
- num: int,
- layout: MMALayout,
- *,
- shape: LoadShape | None = None,
- src_format: LoadSrcFormat | None = None,
Warp-cooperative load of one to four 8x8 matrix tiles from shared memory.
Emits
ldmatrix.sync.aligned.{shape}.{num}{.trans}{.ss}.{type} d, [a]. All 32 lanes of the issuing warp collectively loadnum8x8 tiles whose row starts each lane holds inptr; the result is the per-thread fragment carried in 32-bit register words as required by subsequentmma.sync/stmatrixinstructions.Lane addressing convention (per
num):num
lanes 0..7
lanes 8..15
lanes 16..31
1 2 4
row starts of tile 0 rows of tile 0 rows of tile 0
rows of tile 1 rows of tile 1
rows of tiles 2..3
- Parameters:
ptr (cutlass.Array or cutlass.Pointer) – Pointer/Array into shared memory; per the PTX ISA the address space must be
.shared{::cta}.num (int) – Number of 8x8 tiles per warp. Must be one of
1,2,4(the PTX.x1/.x2/.x4qualifiers).layout (MMALayout) –
MMALayout.ROWfor the default load.MMALayout.COLselects.trans, which transposes the loaded tile inside the lane registers without reading from a transposed memory layout.shape (LoadShape or None) – Tile shape selector. Defaults to
m8n8(the historical SM75 form).m8n16andm16n16unpack the narrow-float.dst_fmt.src_fmtforms and requiresrc_formatto be set.m16n16additionally requireslayout=MMALayout.COL(.trans).src_format (LoadSrcFormat or None) – Source packing for
m8n16/m16n16.b6x16_p32/b4x16_p64unpacke3m2/e2m1matrices into the.b8x16destination format;b8is the byte form. Must be paired with one ofm8n16/m16n16.
- Returns:
Int32whennum=1;Vector[num x Int32]whennum=2ornum=4. Each element is one 32-bit register word.- Raises:
ValueError –
numis not one of1/2/4.ValueError –
src_formatis given without a matchingshape in {m8n16, m16n16}, orshape in {m8n16, m16n16}is given without a matchingsrc_format.ValueError –
shape=m16n16withoutlayout=MMALayout.COL(the PTX ISA requires.transfor them16n16form).
# 4 x 8x8 b16 tiles, non-transposed, into a vector<4xi32> fragment. smem = cutlass.Array(cutlass.Int16, 4 * 8 * 8, space=cutlass.AddressSpace.smem) regs = nvvm.ldmatrix(smem, num=4, layout=nvvm.MMALayout.ROW) # regs is Vector[4 x Int32], ready as a multiplicand for mma_sync.
- cutlass.experimental.primitives.nvvm_wrapper.load_ext(
- addr: ~cutlass.Array | ~cutlass.Pointer,
- *,
- dtype: type | None = None,
- count: int | None = None,
- l2_cache_hint: int | ~cutlass.Int64 | ~cutlass.Uint64 | None = None,
- order: ~cutlass.experimental.primitives.nvvm_wrapper.MemOrder | None = None,
- scope: ~cutlass.base_dsl.array.MemScope | None = None,
- prefetch: ~cutlass.base_dsl.array.L2PrefetchSize | None = None,
- evict: ~cutlass.base_dsl.array.L1EvictKind | None = None,
- cache_modifier: ~cutlass.base_dsl.array.LoadCacheModifier | None = None,
- shared_space: _install_cutlass_mlir_autodoc_stub.<locals>._DocDialectObject | None = None,
- unified: bool | None = None,
Load a scalar (or a
count-element vector) from generic, global, or shared memory with explicit cache, eviction, and memory-ordering qualifiers (nvvm.load.ext/ PTXld/ld.<vec>).The element type is inferred from the pointer’s
dtype; pass an explicitdtypeto override (e.g. for untyped raw pointers). Passcountto load a vector (PTX.v2/.v4/.v8) and get aVectorback. The underlying op supports onlyb8/b16/b32/b64/b128integer widths andf32/f64floats: load a 16-bit float asInt16and bitcast it yourself.- Parameters:
addr (Array | Pointer) – Address to load from (generic, global, or shared pointer).
dtype (type | None) – DSL type of the loaded value; inferred from
addr.dtypewhen omitted. Must be an 8/16/32/64/128-bit integer (signed or unsigned) orFloat32/Float64; a 16-bit float must be loaded asInt16and bitcast.count (int | None) – If set, load
countelements as aVector(PTX.v2/.v4/.v8); if omitted, load a scalar.l2_cache_hint (int | Int64 | Uint64 | None) – 64-bit L2 cache-eviction policy handle (generic / global space only).
order (MemOrder | None) – Memory ordering (
weakdefault,relaxed,acquire,volatile,mmio).relaxed/acquirerequirescope.scope (MemScope | None) – Memory scope (
cta,cluster,gpu,sys) for an ordered load.prefetch (L2PrefetchSize | None) – L2 prefetch size hint (generic / global space only).
evict (L1EvictKind | None) – L1 eviction-priority hint; mutually exclusive with
cache_modifier.cache_modifier (LoadCacheModifier | None) – Cache operator (
ca/cg/cs/lu/cv); only valid on the defaultweakordering.shared_space (SharedSpace | None) – Shared sub-space (
ctadefault,clusterfor distributed shared memory); for shared-space pointers only.unified (bool | None) – Set the
.unifiedqualifier (generic / global space only).
- Returns:
The loaded scalar as the requested DSL type, or a
Vectorofcountelements whencountis set.- Return type:
Int8 | Uint8 | Int16 | Uint16 | Int32 | Uint32 | Int64 | Uint64 | Int128 | Uint128 | Float32 | Float64 | Vector
- Raises:
TypeError – if
dtypeis omitted andaddrcarries nodtypeto infer from.ValueError – if the qualifier combination is illegal, e.g.
cache_modifierwithevictor with non-weakordering;relaxed/acquirewithoutscope;volatilewith a cache op/hint orunified(prefetchis allowed);mmiowithoutscope=sys; orshared_spacecombined withl2_cache_hint/prefetch/unified/mmio.
ptr = arr.data_ptr() + tx # Stream a global value through L2 only (bypass L1). v = nvvm.load_ext(ptr, dtype=cutlass.Int32, cache_modifier=LoadCacheModifier.CG) # Vectorized load: 4 x f32 in one ld.global.v4.b32, with cache control. v4 = nvvm.load_ext(ptr, dtype=cutlass.Float32, count=4, cache_modifier=LoadCacheModifier.CG)
- cutlass.experimental.primitives.nvvm_wrapper.mapa( ) Array | Pointer#
Translate a local SMEM address to a peer CTA’s distributed-SMEM address.
Emits
mapa.shared::cluster— given a pointer to local shared memory and a peer CTA’s cluster rank, returns a pointer to the same SMEM offset in that peer’s local SMEM, valid in the cluster (shared::cluster) state space. The translation is a single-cycle hardware operation that exposes another CTA’s SMEM through the cluster interconnect. Passingaddrspace=0selects the generic-addressing form (PTXmapawithout.space), where both the source and result are generic addresses pointing to shared memory.Used to construct cluster-shared mbarrier pointers (so all participating CTAs can signal/wait on the same physical mbar) and for direct peer-CTA SMEM reads/writes that bypass GMEM.
- Parameters:
addr – Local-SMEM
Array/Pointer. The same offset will be translated in the peer CTA’s SMEM.cta_rank – Cluster rank of the target peer CTA. Must be a valid rank within the launched cluster (
< cluster_size); out-of-range values produce undefined results.addrspace – Address space of the returned pointer. Default
7is the distributed-shared (shared::cluster) space; pass0for the generic-addressing form. No other address space is representable. (A generic-addressed pointer to shared memory may be used with either form; the NVVM verifier enforces the exact rule.)
- Returns:
Pointer/Array (matching input type) addressing the peer CTA’s SMEM at the same offset, in the requested address space.
- Raises:
ValueError – if
addrspaceis neither7(shared::cluster) nor0(generic), or if a statically knowncta_rankis negative.
- cutlass.experimental.primitives.nvvm_wrapper.match_sync( ) Int32 | tuple[Int32, Boolean]#
1:1 wrapper over
nvvm.match_sync.Returns
Int32forany,(Int32, Boolean)forall.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_arrive(
- addr: Array | Pointer,
- *,
- count: int | Int32 | Uint32 | None = None,
- scope: MemScope | None = None,
- relaxed: bool | None = None,
Decrement an mbarrier’s pending arrival count by
count(default 1).When the total arrivals satisfy the barrier’s
countthreshold, the barrier fires: its internal parity flips and all waiters unblock.- Parameters:
addr – Pointer/Array to the 64-bit SMEM mbarrier.
count –
How many arrival credits to consume in one call. Defaults to 1 (the usual case). Pass
count=Nto let a single thread batch-arrive on behalf of N threads — useful for pre-signaling slots:# Prologue: pre-signal all empty_bar slots so the producer's # first wait(parity=0) passes immediately. if nvvm.elect_sync(): for i in cutlass.range_constexpr(NUM_STAGES): nvvm.mbarrier_arrive(empty_bar + i)
scope – Memory-ordering scope for the arrive (default CTA).
relaxed – When
Trueuse.relaxedordering instead of the default.release.
- Raises:
TypeError –
addris not in shared, shared::cluster, or generic memory.ValueError –
countis a Pythonintoutside[1, 2**20 - 1].
`mbarrier_arrive` vs `mbarrier_arrive_expect_tx`: both count as a software arrive. Use
arrive_expect_txinstead when a TMA load is involved — it additionally registers the byte count that TMA hardware must deliver viacomplete_txbefore the barrier fires. For pure software producer-consumer pipelines (non-TMA), usembarrier_arrive.- Returns:
Opaque 64-bit state token for shared/generic pointers;
Nonefor cluster-space pointers.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_arrive_drop(
- addr: Array | Pointer,
- *,
- count: int | Int32 | Uint32 | None = None,
- scope: MemScope | None = None,
- relaxed: bool | None = None,
1:1 wrapper over
nvvm.mbarrier_arrive_drop.Returns an opaque 64-bit state token for shared/generic pointers. Returns None for shared::cluster pointers.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_arrive_drop_expect_tx(
- addr: Array | Pointer,
- txcount: int | Int32 | Uint32,
- *,
- scope: MemScope | None = None,
- relaxed: bool | None = None,
1:1 wrapper over
nvvm.mbarrier_arrive_drop_expect_tx.Returns an opaque 64-bit state token for shared/generic pointers. Returns None for shared::cluster pointers.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_arrive_drop_nocomplete( ) Int64#
Wrapper over
nvvm.mbarrier_arrive_drop_nocomplete.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_arrive_expect_tx(
- addr: Array | Pointer,
- txcount: int | Int32 | Uint32,
- *,
- scope: MemScope | None = None,
- relaxed: bool | None = None,
Signal an mbarrier’s TMA transaction count and count as the software arrive.
Used in TMA pipelines where
mbarrier_init(bar, count=1)was used: this call serves as the single software arrive and registerstxcountbytes that TMA hardware must deliver viacomplete_txbefore the barrier fires. No separatembarrier_arrivecall is needed.- Parameters:
addr – Pointer/Array to the 64-bit SMEM mbarrier.
txcount –
Bytes that TMA
complete_txsignals must deliver before the barrier fires. The interpretation depends on how many producers share this barrier:Single producer (
mbarrier_init(bar, count=1)): one call, pass the total bytes for all TMA loads sharing this barrier:txcount = A_bytes + B_bytes. The barrier fires once TMA delivers all bytes.Two producers (
mbarrier_init(bar, count=2)): each producer callsarrive_expect_txindependently with its own share only (producer-A passestx_A, producer-B passestx_B). The barrier fires after both arrives andtx_A + tx_Bbytes have been delivered.
Formula per load:
num_rows * num_cols * sizeof(dtype).
- Raises:
ValueError –
txcountis a negative Pythonint.TypeError –
addris not in shared, shared::cluster, or generic memory.
Must be called before the corresponding
cp_async_bulk_tensor_*call(s), so the transaction counter is set before TMA can decrement it.Call from one thread only — either
tx == 0or a single elected thread (nvvm.elect_sync()). For CTA_2 multicast: leader only; multiplytxcountby the number of CTAs in the cluster.- Returns:
Opaque 64-bit state token for shared/generic pointers;
Nonefor cluster-space pointers.
# Single producer: separate elect_sync for arrive vs TMA loads (performance) if nvvm.elect_sync(): nvvm.mbarrier_arrive_expect_tx(full_bar + s, A_bytes + B_bytes) if nvvm.elect_sync(): nvvm.cp_async_bulk_tensor_shared_cta_global(sA, tma_a, full_bar + s, coord_a) if nvvm.elect_sync(): nvvm.cp_async_bulk_tensor_shared_cta_global(sB, tma_b, full_bar + s, coord_b)
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_arrive_nocomplete( ) Int64#
Wrapper over
nvvm.mbarrier_arrive_nocomplete.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_complete_tx( ) None#
Manually report completed async-copy bytes to an mbarrier (SM90+).
Decrements the barrier’s expected transaction count by
txcountbytes without issuing a software arrive. Use this when the hardware does not deliver thecomplete_txsignal automatically — for example, after acp_async_shared_globalpipeline that does not use TMA.For TMA-based pipelines (
cp_async_bulk_tensor_*), the TMA hardware deliverscomplete_txautomatically when the copy finishes — do not callmbarrier_complete_txin that case (double-counting corrupts the barrier state).- Parameters:
addr – Pointer/Array to the 64-bit SMEM mbarrier.
txcount – Number of bytes to report as completed. Must equal the total bytes delivered by the corresponding async copies.
scope – Memory scope (default:
cta).
- Raises:
ValueError –
txcountis a negative Pythonint.TypeError –
addris not in shared, shared::cluster, or generic memory.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_expect_tx( ) None#
Wrapper over
nvvm.mbarrier_expect_tx.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_init( ) None#
Initialize a 64-bit mbarrier object in shared memory.
Sets the expected-arrival count and resets the barrier’s phase to 0. Lowers to
mbarrier.init{.shared{::cta}}.b64 [addr], count.- Parameters:
addr (cutlass.Array or cutlass.Pointer) – Pointer or Array addressing the 64-bit mbarrier object. Must resolve to shared memory at runtime.
cutlass.AddressSpace.smemis preferred (NVVM emits the.sharedvariant with a 32-bit address operand);cutlass.AddressSpace.genericis also accepted and emits the generic-formmbarrier.init.b64instruction, which requires the runtime address to fall within the.shared::ctawindow (behavior is undefined otherwise). Typically allocated viacutlass.Array(cutlass.Int64, N, space=cutlass.AddressSpace.smem).count (int or cutlass.Int32 or cutlass.Uint32) – Expected arrival count: how many
mbarrier_arrive(ormbarrier_arrive_expect_tx) calls must be satisfied before the barrier fires and flips its phase. Valid range is[1, 2**20 - 1]. Constexprintvalues are checked at trace time; dynamic values get a--enable-assertions-gated runtime check. For TMA pipelines use1: one elected thread callsarrive_expect_txand TMA hardware deliverscomplete_tx.
- Raises:
TypeError –
addris in an address space other thanSHAREDorGENERIC(e.g. global, local, tensor memory, orSHARED_CLUSTER: the PTX ISA does not definembarrier.initfor these).ValueError –
countis a Pythonintoutside[1, 2**20 - 1].
# Preferred: one warp initializes one disjoint mbarrier group. warp_idx = cute.arch.warp_idx() tidx, _, _ = cute.arch.thread_idx() lane_idx = tidx & 31 if warp_idx == 0: if lane_idx < NUM_AB_STAGES: nvvm.mbarrier_init(ab_full_bar + lane_idx, 1) nvvm.bar_warp_sync(cute.arch.FULL_MASK) elif warp_idx == 1: if lane_idx < NUM_AB_STAGES: nvvm.mbarrier_init(ab_empty_bar + lane_idx, ab_empty_count) nvvm.bar_warp_sync(cute.arch.FULL_MASK) elif warp_idx == 2: if lane_idx < NUM_SF_STAGES: nvvm.mbarrier_init(sf_full_bar + lane_idx, 1) nvvm.bar_warp_sync(cute.arch.FULL_MASK) elif warp_idx == 3: if lane_idx < NUM_SF_STAGES: nvvm.mbarrier_init(sf_empty_bar + lane_idx, sf_empty_count) nvvm.bar_warp_sync(cute.arch.FULL_MASK) elif warp_idx == 4: if lane_idx < 3: count = 1 if lane_idx == 0 else 8 if lane_idx == 1 else 32 nvvm.mbarrier_init(aux_bar + lane_idx, count) nvvm.bar_warp_sync(cute.arch.FULL_MASK) nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync()
# Fallback: simple elected-thread init for a few barriers. if warp_idx == 0: if nvvm.elect_sync(): for i in cutlass.range_constexpr(NUM_STAGES): nvvm.mbarrier_init(full_bar + i, 1) nvvm.mbarrier_init(empty_bar + i, 1) nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync()
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_inval(addr: Array | Pointer) None#
Invalidate an mbarrier object so its storage can be reused.
Emits
mbarrier.inval. Marks the mbarrier ataddrinvalid; the underlying shared-memory bytes may then be repurposed. Pair withmbarrier_init()to recreate a barrier in the same storage.- Parameters:
addr (Array | Pointer) – Pointer to the mbarrier. Must reside in shared memory (generic addressing into
.sharedis also accepted).- Raises:
TypeError –
addris a typed operand in an address space other than shared or generic.
if nvvm.elect_sync(): nvvm.mbarrier_inval(mbar)
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_test_wait(
- addr: Array | Pointer,
- state_or_phase: Int64 | Int32,
- *,
- scope: MemScope | None = None,
- relaxed: bool | None = None,
Test whether an mbarrier phase has completed.
Returns the PTX
waitCompletepredicate formbarrier.test_wait. The PTX ISA 9.3 primary-phase form can also producereportPredicateandreportValueoperands; this wrapper intentionally exposes only the completion boolean.- Parameters:
addr – Pointer/Array to the 64-bit mbarrier object.
state_or_phase – State token returned by
mbarrier_arriveor a parity value for parity-style waits.scope – Optional memory scope when using explicit acquire/relaxed semantics.
relaxed – Emit relaxed ordering when
True; omit for default acquire semantics.
- Returns:
Truewhen the requested phase has completed.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_try_wait(
- addr: Array | Pointer,
- state_or_phase: Int64 | Int32,
- *,
- ticks: int | Int32 | Uint32 | None = None,
- scope: MemScope | None = None,
- relaxed: bool | None = None,
Try to wait for an mbarrier phase, optionally with a time hint.
Returns the PTX
waitCompletepredicate formbarrier.try_wait. If the phase has not completed, the executing thread may suspend for up toticksnanoseconds, or for an implementation-defined time whenticksis omitted. PTX ISA 9.3 primary-phase report operands are not exposed by this wrapper.- Parameters:
addr – Pointer/Array to the 64-bit mbarrier object.
state_or_phase – State token returned by
mbarrier_arriveor a parity value for parity-style waits.ticks – Optional time hint in nanoseconds.
scope – Optional memory scope when using explicit acquire/relaxed semantics.
relaxed – Emit relaxed ordering when
True; omit for default acquire semantics.
- Returns:
Truewhen the requested phase has completed.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_try_wait_parity(
- addr: Array | Pointer,
- phase: int | Int32 | Uint32,
- *,
- time_limit: int | Int32 | Uint32 = 10000000,
- scope: MBarrierScope | None = None,
- order: MemOrder | None = None,
Attempt a parity try-wait on an mbarrier phase.
This issues one
mbarrier.try_wait.parityattempt. If the phase has not completed, the executing thread may be hardware-suspended for up totime_limitnanoseconds before returningFalse. The caller is responsible for retrying when a blocking wait is needed.- Parameters:
addr – Pointer/Array to the 64-bit SMEM mbarrier object.
phase –
Parity value to wait against. Returns
Trueonce the barrier’s internal parity differs fromphase(i.e. the barrier fired and advanced its phase). ReturnsFalseon timeout.Fresh barrier starts at parity 0.
Pass
phase=0to block until the first arrival: waits because current parity (0) equalsphase(0).Pass
phase=1on a fresh barrier to pass immediately: current parity (0) ≠phase(1), so no waiting needed.
time_limit – Hardware suspend timeout in nanoseconds. Defaults to
10_000_000(10 ms). Omit it unless you need a different suspend window. The warp may be hardware-suspended for up to this many nanoseconds on each call.
- Returns:
Truewhen the barrier phase has advanced pastphase;Falseiftime_limitexpired without completion.- Raises:
ValueError –
phaseis a Pythonintother than 0 or 1.ValueError –
time_limitis a negative Pythonint.TypeError –
addris not in shared or generic memory.
Always wrap in a while loop — the retry is the caller’s responsibility:
while not nvvm.mbarrier_try_wait_parity(bar + s, parity): pass
This lowers to
mbarrier.try_wait.parity.acquire.cta.shared::cta.b64with explicit.acquire.ctaordering (ensures TMA writes are visible after the wait). The same intrinsic without the.acquire.ctaqualifier has weaker ordering and benchmarks ~5% slower, so avoid it.Phase formula for a circular N-stage pipeline at iteration k:
cons_parity = (k // cutlass.Int32(NUM_STAGES)) & cutlass.Int32(1) while not nvvm.mbarrier_try_wait_parity(full_bar + s, cons_parity): pass
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_try_wait_timelimit(
- addr: Array | Pointer,
- state: int | Int64 | Uint64,
- time_limit: int | Int32 | Uint32,
- *,
- scope: MBarrierScope | None = None,
- order: MemOrder | None = None,
Try to wait for a state-token mbarrier phase with an explicit time limit.
This is the non-parity time-limited
mbarrier.try_waitform. It returns only the PTXwaitCompletepredicate; PTX ISA 9.3reportPredicate/reportValueoperands are not exposed here.- Parameters:
addr – Pointer/Array to the 64-bit mbarrier object.
state – State token returned by a previous mbarrier arrive operation.
time_limit – Time hint in nanoseconds before the suspended thread may resume and return
False.scope – Optional memory scope for the wait.
order – Optional memory-order qualifier.
- Returns:
Truewhen the state-token phase has completed;Falsewhen the time limit expires first.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_wait(
- addr: Array | Pointer,
- state: int | Int64 | Uint64,
- kind: MBarrierWait,
- *,
- scope: MBarrierScope | None = None,
- order: MemOrder | None = None,
Wrapper over
nvvm.mbarrier_wait.
- cutlass.experimental.primitives.nvvm_wrapper.mbarrier_wait_parity(
- addr: Array | Pointer,
- phase: int | Int32 | Uint32,
- kind: MBarrierWait,
- *,
- scope: MBarrierScope | None = None,
- order: MemOrder | None = None,
Test- or try-wait on an mbarrier’s parity (no time-limit form).
Returns the
waitCompletepredicate:Trueonce the barrier’s phase has advanced pastphase(its internal parity differs), elseFalse.- Parameters:
addr – Pointer/Array to the 64-bit mbarrier, in shared or generic memory.
phase – Parity to wait against (
0or1). Completes when the barrier’s internal parity differs fromphase; a fresh barrier starts at parity 0.kind –
MBarrierWait.TESTselectsmbarrier.test_wait.parity(a non-blocking check that never suspends);MBarrierWait.TRYselectsmbarrier.try_wait.parity(may hardware-suspend the thread until the phase completes or an implementation-defined limit). Both return thewaitCompletepredicate, so aTRYwait is wrapped in a retry loop.scope – Memory-ordering scope (default CTA).
order – Memory ordering; the underlying op supports only
acquire(the effective default) orrelaxed.
- Raises:
ValueError –
phaseis a Pythonintother than 0 or 1.- Returns:
Truewhen the barrier phase has advanced pastphase;Falseotherwise.
# Blocking consumer wait (TRY): retry until the phase flips. while not nvvm.mbarrier_wait_parity(bar, parity, nvvm.MBarrierWait.TRY): pass # Non-blocking probe (TEST): single check, never suspends. done = nvvm.mbarrier_wait_parity(bar, parity, nvvm.MBarrierWait.TEST)
- cutlass.experimental.primitives.nvvm_wrapper.memory_barrier(scope: MemScope) None#
Order this thread’s memory accesses at the given scope.
Emits
membar.{scope}. Guarantees that the issuing thread’s prior memory accesses are performed atscopebefore any of its subsequent accesses. This is the legacymembarordering primitive; preferfence_proxy()/ the acquire-release atomics for finer control.- Parameters:
scope (MemScope) – Scope at which the ordering is observed (
"cta"/"cluster"/"gpu"/"sys").
nvvm.memory_barrier("gpu")
- cutlass.experimental.primitives.nvvm_wrapper.mma_block_scale(
- res: Any,
- shape: Any,
- scale_vec_size: Any,
- block_scale_format: Any,
- kind: Any,
- *args: Any,
- **kwargs: Any,
Gated 1:1 wrapper over
nvvm.mma.block_scale.The
.scale_vec::4X+.ue8m0scale type +.kind::mxf4nvf4combination was introduced in PTX ISA 9.1 and is unavailable on CTK 12.9 (PTX ISA 8.8); every other combination predates it, so only that one is gated.
- cutlass.experimental.primitives.nvvm_wrapper.mma_smem_desc(
- pointer: int | Int32 | Uint32,
- ldm: int | Int32 | Uint32,
- stride: int | Int32 | Uint32,
- base_offset: int | Int8 | Uint8,
- swizzle: int | Int8 | Uint8,
- *,
- mma_desc_version: int | None = None,
Wrapper over
nvvm.mma_smem_desc.
- cutlass.experimental.primitives.nvvm_wrapper.mma_sp_block_scale(
- res: Any,
- shape: Any,
- scale_vec_size: Any,
- block_scale_format: Any,
- kind: Any,
- *args: Any,
- **kwargs: Any,
Gated 1:1 wrapper over
nvvm.mma.sp.block_scale.The
.scale_vec::4X+.ue8m0scale type +.kind::mxf4nvf4combination was introduced in PTX ISA 9.1 and is unavailable on CTK 12.9 (PTX ISA 8.8); every other combination predates it, so only that one is gated.
- cutlass.experimental.primitives.nvvm_wrapper.mma_sp_sync(
- res: _install_cutlass_mlir_autodoc_stub.<locals>._DocMlirType,
- shape: ~sphinx.ext.autodoc.mock._MockObject,
- operand_a: ir.Value,
- operand_b: ir.Value,
- operand_c: ir.Value,
- sparse_metadata: int | ~cutlass.Int32 | ~cutlass.Uint32,
- sparsity_selector: int | ~cutlass.Int32 | ~cutlass.Uint32,
- *,
- int_overflow_behavior: ~cutlass.experimental.primitives.nvvm_wrapper.MMAIntOverflow | None = None,
- multiplicand_a_ptx_type: ~cutlass.experimental.primitives.nvvm_wrapper.MMAType | None = None,
- multiplicand_b_ptx_type: ~cutlass.experimental.primitives.nvvm_wrapper.MMAType | None = None,
- ordered_metadata: bool | None = None,
- kind: ~cutlass.experimental.primitives.nvvm_wrapper.MMAKind | None = None,
Wrapper over
nvvm.mma_sp_sync.Returns an LLVM struct. Caller provides the raw MLIR result type as res.
- cutlass.experimental.primitives.nvvm_wrapper.mma_sync(res: _install_cutlass_mlir_autodoc_stub.<locals>._DocMlirType, shape: ~sphinx.ext.autodoc.mock._MockObject | tuple[int, int, int] | dict, layout_a: ~cutlass.experimental.primitives.nvvm_wrapper.MMALayout, layout_b: ~cutlass.experimental.primitives.nvvm_wrapper.MMALayout, operand_a: list[ir.Value], operand_b: list[ir.Value], operand_c: list[ir.Value], *, b1_op: ~cutlass.experimental.primitives.nvvm_wrapper.MMAB1Op | None = None, int_overflow_behavior: ~cutlass.experimental.primitives.nvvm_wrapper.MMAIntOverflow | None = None, multiplicand_a_ptx_type: ~cutlass.experimental.primitives.nvvm_wrapper.MMAType | None = None, multiplicand_b_ptx_type: ~cutlass.experimental.primitives.nvvm_wrapper.MMAType | None = None) Any#
Cooperative warp-wide matrix multiply-accumulate (
D = A*B + C).Emits
mma.sync.aligned.{shape}.{alayout}.{blayout}{.kind}.{dtype}.{atype}.{btype}.{ctype}. All 32 lanes of the issuing warp collectively compute one MMA on the fragments distributed across their registers; the four matrices are sliced across the warp per the PTX ISA’s per-shape fragment layout.The wrapper is a 1:1 mapping over the NVVM dialect
nvvm.mma.syncop. The dialect verifier validates the shape/type/operand-count combination (e.g.m16n8k16.f16requires 4f16x2a, 2f16x2b, and 2f16x2or 4f32c/dregisters). This wrapper adds string/StrEnum coercion for the qualifier attributes and trace-time guards for the qualifier-vs-multiplicand-type coupling that the dialect cannot infer.- Parameters:
res (ir.Type) – MLIR result type of the fragment
D— typically anllvm.struct<...>whose element layout matches the per-thread fragment for the givenshapex type combination.shape (ir.Attribute or tuple[int, int, int] or dict) – MMA shape attribute. Either a pre-built
ir.Attribute(#nvvm.shape<m = ..., n = ..., k = ...>), a 3-tuple(m, n, k), or a{"m": ..., "n": ..., "k": ...}dict.layout_a (MMALayout) – Layout of multiplicand A. Usually
MMALayout.ROW.MMALayout.COLis only legal formma.m8n8k4.f16.layout_b (MMALayout) – Layout of multiplicand B. Usually
MMALayout.COL.MMALayout.ROWis only legal formma.m8n8k4.f16.operand_a – Per-thread fragment of A as a sequence of
ir.Value.operand_b – Per-thread fragment of B as a sequence of
ir.Value.operand_c – Per-thread fragment of accumulator C as a sequence of
ir.Value.b1_op (MMAB1Op or None) – Bit-op selector for single-bit (
.b1) multiplicands;MMAB1Op.XOR_POPCformma.xor.popc(default for.b1),MMAB1Op.AND_POPCformma.and.popc. Only valid when both multiplicand types are.b1.int_overflow_behavior (MMAIntOverflow or None) – Accumulator overflow handling for integer multiplicands (
.u8/.s8/.u4/.s4).MMAIntOverflow.SATFINITEclamps to thes32range (PTX.satfinitemodifier);MMAIntOverflow.WRAPPEDwraps modulo2**32.multiplicand_a_ptx_type (MMAType or None) – Element type of multiplicand A as a
MMAType. Defaults to the type the dialect infers fromoperand_a’s MLIR type; pass it explicitly when the operand carrier type does not uniquely determine the PTX element type (e.g.i32carriers for packeds8/u8/f8data).multiplicand_b_ptx_type (MMAType or None) – Element type of multiplicand B as a
MMAType. Same inference rules asmultiplicand_a_ptx_type.
- Raises:
ValueError –
b1_opis given butmultiplicand_{a,b}_ptx_typeis set to a non-.b1type.ValueError –
int_overflow_behavioris given butmultiplicand_{a,b}_ptx_typeis set to a non-integer type.ValueError –
shapeis a tuple/list of the wrong arity or with non-intentries, or a dict missing the requiredm/n/kkeys.
# m16n8k16 f16 = f16 * f16 + f16 d = nvvm.mma_sync( T.struct([T.vector(2, T.f16())] * 2), shape=(16, 8, 16), layout_a=nvvm.MMALayout.ROW, layout_b=nvvm.MMALayout.COL, operand_a=[a0, a1, a2, a3], # 4 x f16x2 operand_b=[b0, b1], # 2 x f16x2 operand_c=[c0, c1], # 2 x f16x2 )
- cutlass.experimental.primitives.nvvm_wrapper.mov_b32(
- a: int | float | ~cutlass.Int32 | ~cutlass.Uint32 | ~cutlass.Float32,
- *,
- target_type: type = <class 'cutlass.Int32'>,
mov.b32— reinterpret a 32-bit value’s bits as target_type.Emits an
arith.bitcast(no value conversion), e.g. float bits → int for NaN-safe integer compares.
- cutlass.experimental.primitives.nvvm_wrapper.mul(
- a: Int16 | Int32 | Int64,
- b: Int16 | Int32 | Int64,
- mode: MulMode,
- *,
- is_signed: bool | None = None,
Wrapper over
nvvm.mul.
- cutlass.experimental.primitives.nvvm_wrapper.mul_bf16x2(a: Int32, b: Int32) Int32#
mul.bf16x2— packed bf16x2 multiply (two bf16 lanes packed in i32).
- cutlass.experimental.primitives.nvvm_wrapper.mul_packed_f32x2(
- src_a: tuple | Vector,
- src_b: tuple | Vector,
- *,
- rnd: FPRoundingMode | None = None,
- ftz: bool | None = None,
Wrapper over
nvvm.mul_packed_f32x2.Accepts a 2-tuple of f32 scalars or a
Vectorfor each operand and returns a tuple when called with tuples, else aVector.
- cutlass.experimental.primitives.nvvm_wrapper.nanosleep(duration: int | Int32 | Uint32) None#
Wrapper over
nvvm.nanosleep.
- cutlass.experimental.primitives.nvvm_wrapper.pmevent(
- event_id: int | None = None,
- *,
- mask: int | None = None,
Trigger one or more performance-monitor events.
Emits
pmevent(single event) orpmevent.mask(a set of events). Exactly one ofevent_id/maskmust be given.- Parameters:
event_id (int | None) – Single event index in
[0, 15](pmevent).mask (int | None) – 16-bit mask selecting a set of events (
pmevent.mask); bititriggers eventi.
- Raises:
ValueError – neither or both of
event_id/maskare given,event_idis outside[0, 15], ormaskis outside[0, 0xFFFF].
nvvm.pmevent(event_id=3) nvvm.pmevent(mask=0b1010)
- cutlass.experimental.primitives.nvvm_wrapper.prefetch_l1(addr: Array | Pointer) None#
Bring a cache line into L1.
Emits
prefetch.L1 [addr](orprefetch.global.L1/prefetch.local.L1when the address space is statically visible). Per-thread, non-collective; the cache line is warm but no register is loaded — subsequent reads ofaddrstill have to issue anld.For TMA descriptor warm-up use
prefetch_tensormap(); for the uniform-cache hint useprefetchu(); for an L2 warm-up (with optional eviction-priority hint) useprefetch_l2().For conditional execution wrap the call:
if pred: nvvm.prefetch_l1(addr)— the dialect does support a PTX@p prefetchguard, but that lowering is undocumented and equivalent to the explicitiffor every observable effect.- Parameters:
addr (Array or Pointer) – Pointer/Array in generic, global, or local space. Prefetch on SMEM is a no-op per the PTX ISA.
# Warm L1 ahead of a global load nvvm.prefetch_l1(gmem_ptr)
- cutlass.experimental.primitives.nvvm_wrapper.prefetch_l2(
- addr: Array | Pointer,
- *,
- evict_priority: EvictPriority | None = None,
Bring a cache line into L2, optionally with eviction-priority hint.
Emits
prefetch.L2 [addr](orprefetch.global.L2::<priority> [addr]whenevict_priorityis set). Use to warm L2 before a TMA descriptor read or to bias L2 replacement against re-fetched data with"last".For conditional execution wrap the call in
if; the same note onprefetch_l1()applies.- Parameters:
addr (Array or Pointer) – Pointer/Array in generic, global, or local space. Prefetch on SMEM is a no-op per the PTX ISA.
evict_priority (EvictPriority, optional) –
"normal"or"last"(the only two policies theprefetchinstruction supports), orNone(default) to leave the default policy. Maps to the PTX.L2::evict_normal/.L2::evict_lastmodifier with theevict_prefix dropped, to match the rest of the memory-model API. OtherEvictPrioritymembers are valid onld/st/cpbut rejected here.
- Raises:
ValueError – if
evict_priorityis neither"normal"nor"last".
# Warm L2 with an eviction hint nvvm.prefetch_l2(gmem_ptr, evict_priority="last")
- cutlass.experimental.primitives.nvvm_wrapper.prefetch_tensormap(
- addr: Array | Pointer,
- *,
- space: TensormapSpace = 'const',
Warm the TMA tensormap descriptor cache.
Emits
prefetch.tensormap [addr](defaultspace="const") orprefetch.param.tensormap [addr](space="param"). Issue once, on a single thread, before the firstcp.async.bulk.tensorthat consumes the descriptor — cuts the first TMA’s launch latency by hiding the descriptor fetch behind independent work.The canonical pattern is
if nvvm.elect_sync(): prefetch_tensormap(...)— the explicitifis the recommended way to make the prefetch conditional; the dialect’s undocumented@pguard is equivalent.Typical pattern (one lane warms each descriptor):
if nvvm.elect_sync(): nvvm.prefetch_tensormap(tma_desc_a.get_ptr()) nvvm.prefetch_tensormap(tma_desc_b.get_ptr())
- Parameters:
addr (Array or Pointer) – Pointer to the tensormap descriptor.
space (TensormapSpace) –
"const"(default) or"param"— the state space the descriptor lives in."param"is the kernel-argument case;"const"covers the typical__constant__/cutlass.GridConstantcase.
- cutlass.experimental.primitives.nvvm_wrapper.prefetchu(addr: Array | Pointer) None#
Prefetch into the uniform L1 cache.
Emits
prefetchu.L1 [addr]. The uniform cache backs addresses that all lanes in a warp agree on (e.g. constants, kernel parameters) and is separate from the per-thread L1 data cache. Use when a uniformly-addressed value is about to be read by many warps.For conditional execution wrap the call in
if; the same note onprefetch_l1()applies.- Parameters:
addr (Array or Pointer) – Pointer/Array to a uniformly-addressed location.
nvvm.prefetchu(param_ptr)
- cutlass.experimental.primitives.nvvm_wrapper.prmt(
- lo: int | Int32 | Uint32,
- selector: int | Int32 | Uint32,
- mode: PermuteMode,
- *,
- hi: int | Int32 | Uint32 | None = None,
Wrapper over
nvvm.prmt.
- cutlass.experimental.primitives.nvvm_wrapper.red(
- op: ~cutlass.experimental.primitives.nvvm_wrapper.ReductionOp,
- type_: ~cutlass.experimental.primitives.nvvm_wrapper.ReductionType,
- a: ~cutlass.Array | ~cutlass.Pointer,
- b: ~cutlass.Int32 | ~cutlass.Int64 | ~cutlass.Float64 | ~cutlass.BFloat16 | ~cutlass.Float16 | ~cutlass.Float32 | ~cutlass.Vector,
- *,
- mem_order: ~cutlass.experimental.primitives.nvvm_wrapper.MemOrder | None = None,
- mem_scope: ~cutlass.base_dsl.array.MemScope | None = None,
- shared_space: _install_cutlass_mlir_autodoc_stub.<locals>._DocDialectObject | None = None,
- cache_hint: int | ~cutlass.Int64 | ~cutlass.Uint64 | None = None,
Apply a non-returning atomic reduction to a global or shared memory cell.
Emits the PTX
redinstruction family. The value in memory atais combined with operandbusingopand the result is written back toa. Unlikeatomicrmw(), this operation does not return the old memory value.Scalar reductions may target global or shared memory. Vector reductions are global-memory only; the hardware guarantees atomicity independently for each scalar element, not for the whole vector as one transaction. When
mem_orderis omitted PTX assumes.relaxed; whenmem_scopeis omitted PTX assumes.gpu.Note
shared_spaceselects the explicit shared-memory PTX spelling. When omitted for amapa-produced shared-cluster pointer (addrspace 7), the wrapper selectsSharedSpace.shared_clusterautomatically.- Parameters:
op (ReductionOp) – Reduction operation:
AND,OR,XOR,ADD,INC,DEC,MIN, orMAX.type (ReductionType) – PTX reduction type such as
S32,U32,F32,F64,F16,F16X2,BF16, orBF16X2.a (Array or Pointer) – Pointer or Array naming the destination memory cell.
b (Int32, Int64, Float16, BFloat16, Float32, Float64, or Vector) – Value contributed to the reduction. For vector reductions, pass a vector matching the PTX vector/type combination.
mem_order (MemOrder, optional) – Optional memory ordering qualifier.
RELAXEDandRELEASEare the PTXredsemantics.mem_scope (MemScope, optional) – Optional memory scope:
CTA,CLUSTER,GPU, orSYS.cache_hint (int or Int64 or Uint64, optional) – Optional 64-bit L2 cache policy. PTX permits this only for global memory reductions with
.L2::cache_hint.
- Raises:
ValueError – if op is not a valid
ReductionOpor type_ is not a validReductionType.
# Every valid thread contributes one Int32 to a global sum. ptr = sum_out.iterator.raw_ptr() nvvm.red( "add", "s32", ptr, contribution, mem_order="relaxed", mem_scope="gpu", )
- cutlass.experimental.primitives.nvvm_wrapper.redux_sync(
- val: ~cutlass.Int32 | ~cutlass.Float32,
- kind: _install_cutlass_mlir_autodoc_stub.<locals>._DocDialectObject,
- mask_and_clamp: int | ~cutlass.Int32 | ~cutlass.Uint32,
- *,
- abs: bool | None = None,
- nan: bool | None = None,
Reduce
valacross warp lanes selected bymask_and_clamp(sm_80+).Low-level NVVM dialect wrapper for
redux.sync. All participating lanes receive the same result (implicit broadcast). Prefer this over a 5-step butterfly shuffle loop for simple reductions.- Parameters:
kind (ReductionKind) – Reduction operation (ADD, MIN, MAX, AND, OR, XOR, UMIN, UMAX, FMIN, FMAX).
mask_and_clamp (int or Int32 or Uint32) – 32-bit member mask (
0xFFFFFFFFfor all lanes).abs (bool, optional) – Apply
|val|before reducing;FMIN/FMAXonly, sm_100+, defaults to None (disabled).nan (bool, optional) – Propagate NaN to result;
FMIN/FMAXonly, sm_100+, defaults to None (NaN inputs are ignored).
- Returns:
Warp-reduced result broadcast to all participating lanes.
- Return type:
Float32 for FMIN/FMAX; Int32 for all other kinds.
- Raises:
ValueError – if a static mask_and_clamp does not fit in 32 bits, or if abs / nan is set for a non-FMIN/FMAX kind. A runtime (non-
int) mask_and_clamp is not checked at trace time.
# Per-block abs-max for MXFP8 quantization (sm_100+): amax = nvvm.redux_sync(gv, ReductionKind.FMAX, 0xFFFFFFFF, abs=True)
- cutlass.experimental.primitives.nvvm_wrapper.setmaxregister(
- reg_count: int,
- action: SetMaxRegisterAction,
Adjust the per-thread register budget for the issuing warp.
Emits
setmaxnreg.inc/setmaxnreg.dec. The instruction provides a hint that changes the maximum number of per-thread registers owned by the executing warp, claiming registers from or releasing registers to the CTA register pool. PTX requires every warp in a warpgroup to execute the samesetmaxnreginstruction; branch on a warpgroup-uniform role, not on an individual warp role.- Parameters:
reg_count (int) – Target per-thread register count. Must be a multiple of 8, within
[24, 256].action (SetMaxRegisterAction) –
"increase"(claim more from the pool) or"decrease"(release to the pool).
- Raises:
TypeError – if
reg_countis not anintliteral.ValueError – if
reg_countis outside[24, 256]or not a multiple of 8.
warpgroup = cute.arch.warp_idx() // 4 if warpgroup == PROD_WARPGROUP: nvvm.setmaxregister(40, "decrease") # ... TMA issue + mbarrier arrive ... else: nvvm.setmaxregister(232, "increase") # ... MMA + epilogue ... nvvm.barrier_cta_sync()
- cutlass.experimental.primitives.nvvm_wrapper.shfl_sync(
- thread_mask: int | Int32 | Uint32,
- val: Int32 | Float32,
- offset: int | Int32 | Uint32,
- mask_and_clamp: int | Int32 | Uint32,
- kind: Shfl,
- *,
- return_value_and_is_valid: bool | None = None,
Synchronise participating lanes and shuffle a 32-bit value within a warp.
Emits
shfl.sync.{idx|up|down|bfly}.b32— the PTX warp-shuffle family. All lanes named in thread_mask must execute the same instruction before any lane receives a shuffled value from another lane. This synchronizes the register exchange itself, but it does not provide the memory-ordering guarantee ofbar.sync/bar_warp_sync().ShflKind variants:
"idx"— each lane reads from absolute source lane offset;result[lane] = val[offset]. Used for broadcast (offset=0)."up"— each lane reads from lanemax(lane - offset, lower_bound); result is the value offset lanes earlier in the warp."down"— each lane reads from lanemin(lane + offset, upper_bound); result is the value offset lanes ahead. Used in butterfly reductions."bfly"— each lane reads from lanelane XOR offset; enables butterfly reduction trees without out-of-range clamping.
mask_and_clamp encoding:
mask_and_clampis a packed 32-bit integer that controls sub-warp segmentation and the out-of-range clamp boundary:Bits
[12:8]— segmask:(WARP_SIZE - 1) XOR (width - 1). Lanes that differ only in the lowlog2(width)bits form one shuffle segment. Use(31 << 8) | clampfor a full 32-lane warp.Bits
[4:0]— clamp: upper boundary (width - 1) foridx/down/bfly; lower boundary (0) forup. When a source lane would fall outside the segment, the clamped boundary lane’s value is returned instead.
For a full 32-lane warp these precomputed values are correct:
"idx"/down/bfly:mask_and_clamp = 0x1F(segmask = 0, clamp = 31)"up":mask_and_clamp = 0x00(segmask = 0, clamp = 0)
The following higher-level helpers compute
mask_and_clampautomatically from awidthargument; prefer them unless you need explicit control over the packed field:shuffle_syncshuffle_sync_upshuffle_sync_downshuffle_sync_xor
- Parameters:
thread_mask (int or cutlass.Int32 or cutlass.Uint32) – 32-bit participation mask; bit i = 1 means lane i takes part. All participating lanes must execute the instruction together. Pass
0xFFFFFFFFfor a full-warp shuffle.val (cutlass.Int32 or cutlass.Float32) – The 32-bit value this lane contributes to the shuffle.
offset (int or cutlass.Int32 or cutlass.Uint32) – Interpretation depends on kind:
idx→ absolute source lane ID[0, 31];up/down→ relative lane delta[0, 31];bfly→ XOR lane mask[0, 31].mask_and_clamp (int or cutlass.Int32 or cutlass.Uint32) – Packed sub-warp segmentation and clamp boundary. See encoding description above. For full-warp shuffles use
0x1F(idx/down/bfly) or0x00(up).kind (Shfl) – Shuffle direction. One of
"idx","up","down","bfly".return_value_and_is_valid (bool, optional) – If
True, return a(value, is_valid)tuple where is_valid isTruewhen the source lane was within the active segment (i.e. the result is not the clamped fallback). Defaults toNone(return value only).
- Returns:
Shuffled value from the source lane, or a
(value, is_valid)tuple when return_value_and_is_valid isTrue.- Return type:
cutlass.Int32 or cutlass.Float32, or tuple[cutlass.Int32 or cutlass.Float32, cutlass.Boolean]
- Raises:
ValueError – if a static thread_mask or mask_and_clamp does not fit in 32 bits, if a static offset is outside
[0, 31], or if kind is not one of"idx","up","down","bfly". Runtime (non-int) values are not checked at trace time.
Constraints:
All lanes in thread_mask must reach the instruction; a subset that diverges before
shfl.synccauses undefined behaviour.offset and the low bits of mask_and_clamp must be in
[0, 31].Available on SM30+ (Kepler); the synchronisation guarantee requires SM70+ (Volta) for correctness in independently-scheduled warps.
Prefer the higher-level helpers (
shuffle_sync,shuffle_sync_down,shuffle_sync_xor) for common patterns; usenvvm.shfl_syncdirectly only when you need fine-grained control overmask_and_clamporkind.
# Broadcast lane 0's value to all lanes (full warp) val = cutlass.Float32(cute.arch.lane_idx) broadcast = nvvm.shfl_sync(0xFFFFFFFF, val, 0, 0x1F, "idx") # Butterfly reduction: warp sum acc = cutlass.Int32(cute.arch.lane_idx) for delta in [16, 8, 4, 2, 1]: other = nvvm.shfl_sync(0xFFFFFFFF, acc, delta, 0x1F, "bfly") acc = acc + other # acc now holds the warp sum on all lanes
- cutlass.experimental.primitives.nvvm_wrapper.st_bulk( ) None#
Bulk-initialize a shared-memory byte range to a constant value.
Emits
st.bulk.shared::cta [addr], size, init_val;. Writes a contiguous run ofsizebytes ataddr(SMEM) toinit_val(currently the only legal value is0). Useful for zero-initializing tiles without a per-thread loop.- Parameters:
addr (cutlass.Array or cutlass.Pointer) – SMEM destination pointer/array; must be 16-byte aligned.
size (int or cutlass.Int64 or cutlass.Uint64) – Number of bytes to write; positive multiple of 16.
init_val (int, optional) – Constant byte pattern; PTX currently mandates 0.
- Raises:
ValueError – if a statically known
sizeis not a positive multiple of 16, orinit_valis not0/None.
- cutlass.experimental.primitives.nvvm_wrapper.stmatrix(
- ptr: Array | Pointer,
- sources: int | Int32 | Uint32 | Vector | list | tuple,
- layout: MMALayout,
- *,
- shape: StoreShape | None = None,
Warp-cooperative store of one to four 8x8 (or 16x8) matrix tiles to SMEM.
Emits
stmatrix.sync.aligned.{shape}.{num}{.trans}{.ss}.{type} [a], d. All 32 lanes of the issuing warp collectively storenumtiles whose fragment registers they hold; the per-lane row-start address goes throughptrexactly as forldmatrix().sourcesaccepts three shapes:a single scalar (
int/Int32/Uint32) – storesnum=1,a
Vector[N x Int32]– storesnum=N(Nin{1, 2, 4}); decomposed viavector.extractbefore forwarding to the dialect,a Python
list/tupleof scalars – storesnum=len(sources); each element is coerced toInt32.
- Parameters:
ptr (cutlass.Array or cutlass.Pointer) – Pointer/Array into shared memory; per the PTX ISA the address space must be
.shared{::cta}.sources (int or Int32 or Uint32 or Vector or list or tuple) – Per-lane source fragment. Length (or
Vectorshape) must be 1, 2, or 4 – the PTX.x1/.x2/.x4qualifiers.layout (MMALayout) –
MMALayout.ROWfor the default store.MMALayout.COLselects.trans(in-register transpose before committing to SMEM).shape (StoreShape or None) – Tile shape selector. Defaults to
m8n8;m16n8is the new PTX 9.3.b8store variant.
- Returns:
None–stmatrixwrites to SMEM and has no SSA result.- Raises:
ValueError –
sourcescount is statically known and is not in{1, 2, 4}.
# Store a Vector[4 x Int32] accumulator fragment to a 4 x 8x8 SMEM tile. smem = cutlass.Array(cutlass.Int16, 4 * 8 * 8, space=cutlass.AddressSpace.smem) nvvm.stmatrix(smem, frag, nvvm.MMALayout.ROW) # frag : Vector[4 x Int32]
- cutlass.experimental.primitives.nvvm_wrapper.store_ext(
- value: ir.Value,
- addr: ~cutlass.Array | ~cutlass.Pointer,
- *,
- l2_cache_hint: int | ~cutlass.Int64 | ~cutlass.Uint64 | None = None,
- order: ~cutlass.experimental.primitives.nvvm_wrapper.MemOrder | None = None,
- scope: ~cutlass.base_dsl.array.MemScope | None = None,
- evict: ~cutlass.base_dsl.array.L1EvictKind | None = None,
- cache_modifier: ~cutlass.base_dsl.array.StoreCacheModifier | None = None,
- shared_space: _install_cutlass_mlir_autodoc_stub.<locals>._DocDialectObject | None = None,
Store a scalar to generic, global, or shared memory with explicit cache, eviction, and memory-ordering qualifiers (
nvvm.store.ext/ PTXst).The underlying op supports only
b8/b16/b32/b64/b128integer widths andf32/f64floats: store a 16-bit float by bitcasting it toInt16first.- Parameters:
value (ir.Value | Vector) – Register value to store; its type selects the store width. May be a scalar or a
Vector(PTX.v2/.v4/.v8).addr (Array | Pointer) – Destination address (generic, global, or shared pointer).
l2_cache_hint (int | Int64 | Uint64 | None) – 64-bit L2 cache-eviction policy handle (generic / global space only).
order (MemOrder | None) – Memory ordering (
weakdefault,relaxed,release,volatile,mmio).relaxed/releaserequirescope.scope (MemScope | None) – Memory scope (
cta,cluster,gpu,sys) for an ordered store.evict (L1EvictKind | None) – L1 eviction-priority hint; mutually exclusive with
cache_modifier.cache_modifier (StoreCacheModifier | None) – Cache operator (
wb/cg/cs/wt); only valid on the defaultweakordering.shared_space (SharedSpace | None) – Shared sub-space (
ctadefault,clusterfor distributed shared memory); for shared-space pointers only.
- Raises:
ValueError – if the qualifier combination is illegal, e.g.
cache_modifierwithevictor with non-weakordering;relaxed/releasewithoutscope;volatilewith any cache op/hint;mmiowithoutscope=sys; orshared_spacecombined withl2_cache_hint/mmio.TypeError – if
valueis a 16-bit float (Float16/BFloat16); bitcast it toInt16before callingstore_ext.
ptr = arr.data_ptr() + tx # Streaming store (likely written once): bypass L1 reuse tracking. nvvm.store_ext(val.ir_value(), ptr, cache_modifier=StoreCacheModifier.CS)
- cutlass.experimental.primitives.nvvm_wrapper.sub_packed_f32x2(
- src_a: Vector,
- src_b: Vector,
- *,
- rnd: FPRoundingMode | None = None,
- ftz: bool | None = None,
Wrapper over
nvvm.sub_packed_f32x2.
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_alloc( ) None#
Allocate TMEM columns for tcgen05 operations.
- Parameters:
addr –
Pointer/Array to a 32-bit SMEM cell that receives the TMEM token (start column address). Pass
tmem_ptr(anArrayorPointerof dtypeInt32) and read it back after alloc.Validated at trace time: must reside in shared memory when the passed object exposes
.space; opaque pointers are deferred to the MLIR verifier.n_cols –
Number of TMEM columns to allocate. The allocation unit is 32 columns and all lanes per column. Statically known
intvalues must be a power of 2 in[32, 512](validated at trace time); dynamic IR values are forwarded as-is and may fault at runtime if out of range. PTX also requires the number of columns allocated not to increase between any two allocations in CTA execution order. Standard values:CTA_1:
n_cols = (N_TILE // 8) * 32whereN_TILEis the per-CTA accumulator N.N_TILE=128→512columns (fills entire TMEM);N_TILE=64→256columns (two accumulators fit: 256 + 256 = 512).CTA_2: each CTA in the 2-SM group still allocates from its own 512-column TMEM bank, but the accumulator is split M-wise (Layout A: leader holds top M-half × full pair-N, peer holds bot M-half × full pair-N). The CTA_1 formula applied to the per-CTA half (
N_TILE = N_PER_GROUP / 2) gives the minimum, but the simpler safe default is to always allocate 512 for any CTA_2 GEMM (over-allocation is harmless). Applying the CTA_1 formula withN_TILE = N_PER_GROUP(collective N) instead of per-CTA N over-allocates beyond 512 and faults withcudaErrorIllegalInstruction.
group –
'CTA_1'(default) or'CTA_2'. SeeCTAGroup. Alltcgen05instructions within a kernel must use the same group.
# Correct — both alloc and relinquish are warp-collective, neither # is inside elect_sync. mbarrier_init IS elect-safe so it stays # inside the elect_sync block. if warp == 0: if nvvm.elect_sync(): for s in cutlass.range_constexpr(S): nvvm.mbarrier_init(full_bar + s, 1) nvvm.tcgen05_alloc(tmem_ptr, num_cols, group="cta_1") nvvm.tcgen05_relinquish_alloc_permit(group="cta_1")
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_commit(
- addr: Array | Pointer,
- *,
- multicast_mask: Int16 | Int32 | None = None,
- smem_a_read: bool | None = None,
- group: CTAGroup | None = None,
Track prior async tcgen05 operations with an mbarrier.
Used in pipelined kernels to release the SMEM staging buffer back to the producer after prior async
tcgen05operations have completed. Emitstcgen05.commit.cta_group::N.mbarrier::arrive::oneand makes addr track all prior asynctcgen05operations of the same CTA group that were initiated by the executing thread.- Parameters:
addr – Pointer/Array to the 64-bit mbarrier to signal (typically
empty_bar + stage).smem_a_read – When
True, tells the hardware that this commit releases the A-operand SMEM buffer (defaultNone= both A and B).multicast_mask –
CTA participation mask for CTA_2 multicast — a per-bit mask over cluster ranks (not pair-internal). Bit
iset means the arrive lands on the mbar copy at the same SMEM offset in cluster ranki. Two regimes:Single 2-CTA cluster (
cluster_shape=(2,1,1)) — only one 2-SM group, leader at rank 0. Canonical value is3(=0b11covering ranks 0 and 1 = both pair members); this is the value you see in2cta_mma_basic.pyand is right whenever there is exactly one 2-SM group in the cluster.Multi-group clusters (
cluster_shape=(2, n_groups, 1)withn_groups > 1) — each 2-SM groupGhas its leader at cluster rank2*G. Usemulticast_mask = 3 << cluster_rankso the issuing group leader signals ranks2Gand2G+1(its own pair), NOT ranks 0,1. Hard-codingmask=3is a frequent deadlock cause: groups 1, 2, … never receive the arrive and stall at the nexttry_wait_parity.
multicast_mask=1 << cta_rank(commit only to the issuer) is also valid when the follower CTA never waits on this mbar; otherwise the follower deadlocks.Note that the value semantics differ from cluster TMA loads, where commits count arrives on named mbar copies but TMA counts bytes delivered per CTA. Do not cross-apply multicast masks between TMA loads and
tcgen05.commitwithout rechecking which mbarrier copies are signaled. DefaultNoneselects the CTA_1 path with no multicast. Staticintvalues are validated at trace time to fit in 16 bits (the dialect mask isi16); larger literals would be silently truncated.group –
"cta_1"(default) or"cta_2".
# Per-k-tile: signal empty_bar when prior async tcgen05 ops complete. if warp == MMA_WARP: if nvvm.elect_sync(): nvvm.tcgen05_commit(empty_bar + s, group="cta_1") # After K-loop: signal acc_mbar to release the TMEM accumulator if warp == MMA_WARP: if nvvm.elect_sync(): nvvm.tcgen05_commit(acc_mbar, group="cta_1")
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_cp(
- shape: Tcgen05CpShape,
- taddr: Array | Pointer,
- smem_desc: int | Int64 | Uint64,
- *,
- group: CTAGroup | None = None,
- multicast: Tcgen05CpMulticast | None = None,
- src_format: Tcgen05CpSrcFormat | None = None,
Asynchronous SMEM → TMEM copy with optional decompression / multicast.
Emits
tcgen05.cp— staged into the TC issue queue alongsidetcgen05.mmaandtcgen05.shift(the PTX ISA lists the legal ordering pairs). Used to seed an A-from-TMEM operand, to gather narrow-format data (FP6/FP4) into TMEM with on-the-fly widening, or to place block-scaled SFA/SFB metadata in TMEM before the matchingtcgen05_mma_block_scale()-style wrapper consumesscale_a/scale_b.- Parameters:
shape –
Tcgen05CpShapeselecting the data dimensions. Per the PTX ISA these come with.multicastconstraints:.64x128brequireswarpx2::02_13orwarpx2::01_23;.32x128brequireswarpx4; the wider shapes (.128x256b/.4x256b/.128x128b) take no multicast. The.64x128b.warpx2::01_23form is the direct SFB scale metadata copy used by the 2x2 CTA_2 block-scaled N=256 example. The shape<->multicast coupling is validated at trace time.taddr – TMEM destination address (
Array/Pointer). Validated at trace time: must reside in tensor memory when the passed object exposes.space(dialect operand typeLLVM_PointerTensor).smem_desc – 64-bit SMEM matrix descriptor (built via
cutlass.experimental.primitives.Tcgen05SmemDesc.build()).group –
"cta_1"(default) — destination is the issuing CTA’s TMEM."cta_2"— also writes the peer CTA’s TMEM (cluster-collective; both CTAs’ issuing warp must reach this op cooperatively). Alltcgen05.*ops in a kernel must agree ongroup.multicast – Warp-pair / quad-warp multicast policy for the narrow shapes; see
shapeconstraints above.src_format –
B6x16_P32/B4x16_P64to enable on-the-fly decompression tob8x16in TMEM.Nonemeans same-format copy.
# Stage A from SMEM to TMEM ahead of an A-from-TMEM MMA if warp == TMA_WARP: if nvvm.elect_sync(): nvvm.tcgen05_cp( "shape_128x256b", a_tmem_addr, smem_desc_a, group="cta_1", ) nvvm.tcgen05_commit(empty_bar) # Stage 2x2 SFB scale metadata before the block-scaled MMA reads scale_b. sfb_shape, sfb_multicast = ... # S2T copy mode, e.g. 64x128b WARPX2_01_23 if warp == MMA_WARP: if nvvm.elect_sync(): nvvm.tcgen05_cp( sfb_shape, sfb_tmem_addr, sfb_smem_desc, group=nvvm.CTAGroup.CTA_2, multicast=sfb_multicast, )
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_dealloc( ) None#
Free TMEM columns previously allocated by
tcgen05_alloc().- Parameters:
taddr –
TMEM base pointer (addrspace 6) — typically the value read back from the SMEM slot
tcgen05_alloc()wrote to, converted with an addrspace-6 (TMEM) pointer.Validated at trace time: must reside in tensor memory when the passed object exposes
.space; opaque pointers are deferred to the MLIR verifier (the dialect operand type isLLVM_PointerTensor).n_cols – Number of TMEM columns to free; must equal the value passed to the paired
tcgen05_alloc(). Statically knownintvalues are validated against the same whitelist astcgen05_alloc()({32, 64, 128, 256, 512}) at trace time.group –
'CTA_1'(default) or'CTA_2'— must match the group used at alloc time. SeeCTAGroup.
# After the epilogue's TMEM reads and their sync: nvvm.barrier_cta_sync() if warp_idx == 0: nvvm.tcgen05_dealloc(tmem_ptr, NUM_TMEM_COLS, group="cta_1")
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_fence(
- kind: Tcgen05Fence,
Order async tcgen05 operations around thread synchronization.
Emits
tcgen05.fence.{before_thread_sync|after_thread_sync}.- Parameters:
kind (Tcgen05Fence) –
"before_thread_sync"— place before an execution-ordering operation such asnvvm.barrier_cta_sync()or a flag store to order all prior asynctcgen05operations before that synchronization point."after_thread_sync"— place after such a synchronization point and before subsequent asynctcgen05operations.
# After reading TMEM, fence before syncing: c_vec = nvvm.tcgen05_ld(shape, tmem_addr, num=n) nvvm.tcgen05_wait("load") nvvm.tcgen05_fence("before_thread_sync") nvvm.barrier_cta_sync()
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_ld(
- shape: Tcgen05LdStShape,
- tmem_addr: Array | Pointer,
- *,
- num: int = 1,
- pack: bool | None = None,
- offset: int | Int64 | Uint64 | None = None,
Load data from TMEM into registers.
Emits
tcgen05.ld.sync.aligned.{shape}{.xN}{.pack}.b32. The issuing warp collectively loads from TMEM; each lane receives its own register slice from the warp’s accessible TMEM sub-partition. Pair withtcgen05_wait()(LOAD) before reading the result.TMEM access restriction. TMEM is divided into 4 lane-chunks of 32 lanes each. A given lane is accessible by exactly one warp in each warpgroup, determined by the warp’s position within its warpgroup (i.e.
warp_idx % 4):Warp position in warpgroup
Accessible lanes
0
0–31
1
32–63
2
64–95
3
96–127
All warps see all 512 columns; the lane (row) restriction is the only cross-warp partitioning. A warp cannot read another warp’s lanes by changing the address — the row field in
tmem_addr[31:16]is a local index within the warp’s own chunk (5 bits; values ≥ 32 wrap asrow mod 32).Implication for ``M_TILE=128`` GEMM epilogues. Covering all 128 accumulator rows requires 4 issuing warps whose warpgroup positions cover 0..3. For a 4-warp CTA those are warps 0..3. A shifted TMEM-load range such as warps 2..5 is also valid because
warp_idx % 4covers 2, 3, 0, 1. In that case, derive the TMEM row and output row from the physical SP position, not from the logical TMEM-load rank:tmem_sp = warp_idx % 4androw = tmem_sp * 32 + lane. Usingwarp_idx - tmem_ld_warp_startfor the row offset makes warp 4 try to read SP2 rows while the hardware routes it to SP0, causing wrong data or an illegal instruction. Valid organizations include:Time-multiplex — 4-warp CTA; warp 0 does TMA + MMA in the K-loop, then warps 0..3 do epilogue.
Split with shifted epilogue — producer / MMA use earlier warps, and four epilogue warps such as 2..5 drain rows via
warp_idx % 4.
Also note:
tcgen05.mmaonly fills rows for them_dimit was given. Ifm_dim < 8(i.e.M_TILE < 128), rows beyondm_dim * 4per chunk contain stale data. UseM_TILE=128(m_dim=8) to populate all rows.- Parameters:
shape (Tcgen05LdStShape) – Load shape (PTX
.shape1/.shape2). One of"16x32bx2","16x64b","32x32b","16x128b","16x256b". Most common is"32x32b"(one 32-bit register per thread, reads 32 rows × 32 bits = 128 bytes per warp).num > 1stacks multiple shapes into a contiguousVector.tmem_addr (cutlass.Array or cutlass.Pointer) – Pointer/Array to the TMEM location. The address encodes
(row << 16) | colwhere row is the local row within the sub-partition and col is the TMEM column fromtcgen05_alloc. Validated at trace time: must reside in tensor memory when the passed object exposes.space(the dialect operand type isLLVM_PointerTensor); opaque pointers are deferred to the verifier.num (int) – Number of shape repetitions; result has
regs_per_elem * numregisters.regs_per_elemis determined by theshape:"16x32bx2"/"16x64b"/"32x32b"→ 1,"16x128b"→ 2,"16x256b"→ 4. Must be a power of 2 in [1, 128]; total registers (regs_per_elem * num) must not exceed 128.pack (bool, optional) – Enable element packing (reduces register count for sub-32-bit dtypes).
offset (int or cutlass.Int64 or cutlass.Uint64, optional) – Required (and only valid) for
"16x32bx2"— column offset added totmem_addrat runtime. Must beNonefor all other shapes.
- Returns:
A
Vectorof the loaded registers (always a vector, even for a single register). Iftmem_addr.dtypeis notInt32, the result is bitcast to that dtype automatically (e.g.Float32).- Return type:
- Raises:
ValueError –
shapeis not a recognized literal;numis not a power of 2 in [1, 128]; total registers exceed 128;offsetis set for a shape other than"16x32bx2"(or missing for"16x32bx2").TypeError –
tmem_addrexposes a.spacethat is not tensor memory (TMEM).
# Each issuing warp loads its own 32 lanes from the accumulator tile. # Use the physical TMEM/SP owner, not the logical TMEM-load rank. # Build the encoded (row << 16) | col integer address, then convert # to a TMEM pointer in one shot -- pointer arithmetic on an already- # constructed TMEM pointer does not interpret the row/col layout the # way callers expect. tmem_sp = warp_idx % 4 base_row = tmem_sp * 32 tmem_addr = (base_row << 16) | base_col tmem_ptr = ... # addrspace-6 TMEM pointer built from tmem_addr result = nvvm.tcgen05_ld("32x32b", tmem_ptr)
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_mma(
- mma_kind: Tcgen05MMAKind,
- cta_group: CTAGroup,
- d: Array | Pointer,
- a: Array | Pointer | Int64,
- b: int | Int64 | Uint64,
- idesc: int | Int32 | Uint32,
- enable_input_d: int | Boolean,
- *,
- collector_op: Tcgen05MMACollectorOp | None = None,
- a_shift: bool | None = None,
- scale_input_d: int | Int64 | Uint64 | None = None,
- write_disable_mask: Vector | None = None,
Issue a 5th-generation Blackwell tensor-core multiply-accumulate (
tcgen05.mma).Emits
tcgen05.mma.cta_group::{1|2}.kind::<kind>with optional.collector::a::*,.ashift, andscale-input-dmodifiers. The accumulatorD = A * B [+ D]lives in Tensor Memory (TMEM); A is read from SMEM (default) or TMEM; B is always an SMEM descriptor. Requiressm_100aor a supported family target.- Parameters:
mma_kind (Tcgen05MMAKind) –
Data-type/kind selector. Dense kinds are issued by this wrapper; block-scaled kinds are listed for descriptor compatibility and should be issued with
tcgen05_mma_block_scale()."f16"— f16/bf16 operands, f32 accumulator. K=16."tf32"— tf32 operands, f32 accumulator. K=8."int8"— signed/unsigned 8-bit, i32 accumulator. K=32."f8f6f4"— mixed {E4M3, E5M2, E2M3, E3M2, E2M1} inputs, f32 accumulator. K=32."mxf8f6f4"— block-scaled F8F6F4 (usetcgen05_mma_block_scale()). K=32."mxf4"— block-scaled E2M1. K=64."mxf4nvf4"— block-scaled E2M1+NVFP4. K=64.
cta_group (CTAGroup) –
"cta_1"for single-CTA scope (M ∈ {32, 64, 128});"cta_2"for 2-CTA cooperative scope (collective M ∈ {128, 256} across peer CTAs; N effectively doubled via peer SMEM). Everytcgen05.*op in a kernel must use the same group.d (cutlass.Array or cutlass.Pointer) – TMEM accumulator pointer returned by
tcgen05_alloc()(address-space 6). Rows fill according tom_dim; hardware routes writes to sub-partitions by row. Validated at trace time: must reside in tensor memory when the passed object exposes.space(dialect operand typeLLVM_PointerTensor).a (cutlass.Int64, cutlass.Array, or cutlass.Pointer) –
Either a 64-bit SMEM descriptor from
cutlass.experimental.primitives.Tcgen05SmemDesc.build()(A-from-SMEMpath, common) or a TMEM pointer (A-from-TMEMpath — used in BMM2 of FMHA where the previous MMA’s output is reused as the next A).For the TMEM path:
Build an addrspace-6 pointer from
tmem_addrwheretmem_addr = (row << 16) | colis the packed 32-bit TMEM address. Pointer arithmetic on TMEM pointers applies raw packed-token offsets; it does not interpret row/col fields.A’s columns advance by the K-granule per K-step (e.g. 64 TMEM columns per BF16 K-step, not per K element).
Load A into TMEM beforehand via
tcgen05_cp()(SMEM→TMEM copy) ortcgen05_st()(register→TMEM store), or produce it as the output of a priortcgen05_mma()accumulator.
b (int, cutlass.Int64, or cutlass.Uint64) – 64-bit SMEM descriptor for B from
cutlass.experimental.primitives.Tcgen05SmemDesc.build().idesc (int, cutlass.Int32, or cutlass.Uint32) –
Packed 32-bit instruction descriptor encoding
c_format/a_format/b_format/a_major/b_major/m_dim/n_dim. Build withcutlass.experimental.primitives.Tcgen05InstrDesc.build():idesc = cutlass.experimental.primitives.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, # 1 = f32 accumulator a_dtype=cutlass.Float16, # MUST match A dtype — see table below b_dtype=cutlass.Float16, # MUST match B dtype — see table below n_dim=N_TILE, # logical N (multiple of 8); 3 LSBs not encoded m_dim=M_TILE, # logical M (multiple of 16); 4 LSBs not encoded )
Common
a_format/b_formatvalues (must match operand dtype):Value
Dtype
Applicable
mma_kind0
FP16
F161
BF16
F162
TF32
TF320..4
E4M3..E2M1
F8F6F4/MX*variants0
U8
INT81
S8
INT8See
cutlass.experimental.primitives.Tcgen05InstrDescfor the full bit layout. ForCTA_2them_dim/n_dimdescribe the collective tile across both CTAs (e.g. M=256 →m_dim=16).Warning
a_format/b_formatmust match the operand dtype. A mismatch (e.g. BF16 operands witha_format=0) produces silently wrong results (~40–60 max_err on random data); no compile or runtime error is raised.enable_input_d (int or cutlass.Boolean) – Controls first-tile behavior.
True/ non-zero computesD = A*B + D(accumulate into existing TMEM content);False/ zero computesD = A*B(overwrite). Typically passed ask > 0in a k-tile loop so the first k-tile clears the accumulator and subsequent tiles accumulate.collector_op (Tcgen05MMACollectorOp, optional) –
Collector cache usage for operand A. The collector is a small per-MMA-issuer cache that lets consecutive MMAs reuse the same A operand without re-reading SMEM. Values:
"fill"— load A into the collector and use it (seeds the cache on the first MMA of a chain)."use"— read A from the collector (later MMAs in a chain that reuse the same A)."lastuse"— read from the collector, then invalidate the entry (final MMA of the chain)."discard"(default) — do not cache A.
Reuse is opportunistic; hardware may reload despite the permission. Treat the collector strictly as a performance hint. The source memory for A must not be modified while any MMA using that matrix has not completed, regardless of collector state.
a_shift (bool, optional) –
When
True, emits the.ashiftmodifier. In the.ashiftMMA pipeline, the shift is a post-MMA operation: the current MMA reads unshifted A and producesA @ B, then the A TMEM region is shifted for subsequent reads. To observe the shift you need a follow-up MMA (ortcgen05_ld()) that targets the same A TMEM region.A MUST be in TMEM (the ``[a-tmem]`` form). The PTX ISA defines
.ashiftonly fortcgen05.mma.kind.ashift [d-tmem], [a-tmem], b-desc, ...; there is no.ashiftform that takes an SMEM descriptor for A. Passinga_shift=Truewith an SMEM-descriptorais silently ignored by HW.Shift semantics (per-SP, not global): PTX wording “shifts rows down by 1 except the last row” refers to the last row of each TMEM sub-partition, not the global last row. For M=128 (4 SPs × 32 rows) global M-rows {31, 63, 95, 127} retain their original values:
shifted[r] = A[r + 1] when r % 32 != 31 shifted[r] = A[r] when r % 32 == 31
Other constraints:
M ∈ {128, 256}only; mutually exclusive withcollector_op in {FILL, USE}and with.ws(warp- specialized) MMA variants.idesc.max_shiftdoes NOT control.ashift— that field is for the.wsvariant (values 0/1/2/3 → max shifts 0/8/16/32 rows). For plain.ashiftthe shift is always exactly 1 row per MMA regardless ofmax_shift; leave it at 0.scale_input_d (int, cutlass.Int64, or cutlass.Uint64, optional) – Immediate in
[0, 15]. When set, scales the input accumulator asD *= 2**(-scale_input_d)before the MAC. Valid only for"f16"and"tf32"(PTX ISA 9.3). The kind restriction and the static[0, 15]range are validated at trace time.write_disable_mask (cutlass.Vector, optional) –
Per-row TMEM write-disable vector. 4-element
Vector[Int32]forCTA_1(128 mask bits / M-rows), 8-element forCTA_2(256 mask bits across the collective M=256 tile).Bit mapping (Layout D, CTA_1): bit
iof elementvmasks output M-rowv * 32 + i. A set bit suppresses the TMEM write for that row, leaving it at the TMEM-initial value (zero aftertcgen05_alloc()).Typical use: row-wise partial-tile handling — suppress writes to M-rows beyond the live fraction of the tile so they don’t clobber accumulator rows still needed by the adjacent tile. The element count (4 for
CTA_1, 8 forCTA_2) is validated at trace time.
Constraints:
Single-thread issue: call inside
if nvvm.elect_sync():. All 32 warp threads issuing the intrinsic would emit 32 duplicate MMAs (undefined behavior / observed hangs). ForCTA_2, only the leader CTA’s elected thread issues.A/B SMEM visibility: producer writes (TMA,
cp.async) must be ordered before the MMA via an mbarrier wait;tcgen05.mmaitself is async with respect to generic memory but synchronous with respect to an arrive-on-completion mbarrier.Completion tracking: follow the last MMA of a chain with
tcgen05_commit()to release the A/B SMEM buffers, and with a commit on an accumulator mbarrier (ortcgen05_wait) before the epilogue reads TMEM.TMEM read ordering: after the accumulator mbarrier fires, use
tcgen05_fence()(BEFORE_THREAD_SYNC) betweentcgen05_ld()and the followingbarrier()/ next MMA.Predication: not supported.
tcgen05.mmahas variable-length operand lists (write_disable_mask, collector ops,scale_input_d) that require NVVM-level register allocation — gate the intrinsic with a surroundingifinstead.
# Per k-tile of a pipelined GEMM (warp 0 issues MMA): if warp_idx == 0: if nvvm.elect_sync(): nvvm.tcgen05_mma( nvvm.Tcgen05MMAKind.F16, nvvm.CTAGroup.CTA_1, tmem_ptr, # D: accumulator in TMEM desc_a, # A: SMEM descriptor desc_b, # B: SMEM descriptor idesc, # Packed instr descriptor k > 0, # enable_input_d: k==0 clears, k>0 accumulates ) # Multi-MMA-warp shared TMEM: the first MMA that starts a fresh # accumulator region must clear (enable_input_d=False); every later # MMA targeting the same accumulator region accumulates. # A-reuse across 2 back-to-back MMAs (same A, different B): if nvvm.elect_sync(): nvvm.tcgen05_mma(..., collector_op=nvvm.Tcgen05MMACollectorOp.FILL) nvvm.tcgen05_mma(..., collector_op=nvvm.Tcgen05MMACollectorOp.LASTUSE)
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_mma_block_scale(
- mma_kind: ~cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMAKind,
- cta_group: ~cutlass.experimental.primitives.nvvm_wrapper.CTAGroup,
- d: ~cutlass.Array | ~cutlass.Pointer,
- a: ~cutlass.Array | ~cutlass.Pointer | ~cutlass.Int64,
- b: int | ~cutlass.Int64 | ~cutlass.Uint64,
- idesc: int | ~cutlass.Int32 | ~cutlass.Uint32,
- enable_input_d: int | ~cutlass.Boolean,
- scale_a: ~cutlass.Array | ~cutlass.Pointer,
- scale_b: ~cutlass.Array | ~cutlass.Pointer,
- *,
- scale_vec_size: ~cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMAScaleVecSize | _install_cutlass_mlir_autodoc_stub.<locals>._DocDialectObject | None = None,
- collector_op: ~cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMACollectorOp | None = None,
- a_shift: bool | None = None,
MMA with per-block scale factors (MXFP / NVFP block scaling).
Emits
tcgen05.mma.cta_group::N.{kind}.block_scale[.scale_vectorsize]. Liketcgen05_mma(), but each MMA additionally multiplies blocks of A and B by per-block scale factors before accumulating, enabling MXFP8 / MXFP6 / MXFP4 / NVFP4 block-scaled formats whose dynamic range is otherwise too narrow for direct GEMM. See the PTX ISA “Block Scaling for tcgen05.mma” section for the full scale-factor layout spec — the layout depends onscale_vec_sizeand the K-dim, and is dense (different from per-row or per-channel scaling).- Parameters:
mma_kind – Top-level block-scale kind (for example
MXF8F6F4for MXFP8/6/4 narrow formats; selects which block-scaling variants are legal). Validated at trace time: must be a block-scaled kind (mxf8f6f4/mxf4/mxf4nvf4); usetcgen05_mma()for non-block-scaled kinds.cta_group –
CTA_1orCTA_2(cluster shape (2,1,1)).d – TMEM accumulator destination. Validated at trace time to reside in tensor memory (dialect operand type
LLVM_PointerTensor).a – A operand — SMEM descriptor (Int64) or TMEM address.
b – B operand SMEM descriptor (Int64).
idesc – Packed 32-bit instruction descriptor. Build FP8/FP6/MX descriptors via
build(), and FP4/NVFP4 descriptors viabuild().enable_input_d – Boolean — when False, ignores prior D contents (D = A·B·scale instead of D += A·B·scale).
scale_b (scale_a,) – TMEM addresses of the scale-factor tiles for A and B respectively. Layout depends on
scale_vec_size— see PTX ISA. Both are validated at trace time to reside in tensor memory (dialect operand typeLLVM_PointerTensor).scale_vec_size –
Tcgen05MMAScaleVecSize— selects 1X / 2X / 4X scale-vector packing within each block. Different sizes have different K-dim compatibility tables.collector_op – A operand reuse policy (see
tcgen05_mma()).a_shift – Same semantics as
tcgen05_mma().
if warp == MMA_WARP: if nvvm.elect_sync(): nvvm.tcgen05_mma_block_scale( nvvm.Tcgen05MMAKind.MXF8F6F4, nvvm.CTAGroup.CTA_1, d_tmem, a_smem_desc, b_smem_desc, idesc, enable_input_d=k > 0, scale_a=scale_a_tmem, scale_b=scale_b_tmem, scale_vec_size=nvvm.Tcgen05MMAScaleVecSize.X2, ) nvvm.tcgen05_commit(empty_bar)
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_mma_sp(
- mma_kind: Tcgen05MMAKind,
- cta_group: CTAGroup,
- d: Array | Pointer,
- a: Array | Pointer | Int64,
- b: int | Int64 | Uint64,
- idesc: int | Int32 | Uint32,
- enable_input_d: int | Boolean,
- sparse_metadata: Array | Pointer,
- *,
- collector_op: Tcgen05MMACollectorOp | None = None,
- a_shift: bool | None = None,
- scale_input_d: int | Int64 | Uint64 | None = None,
- write_disable_mask: Vector | None = None,
Issue a structured-sparse 5th-gen tensor-core MMA (
tcgen05.mma.sp).Like
tcgen05_mma(), but operand A is a structured-sparseM x (K/2)matrix andsparse_metadata(in TMEM) maps the compressed columns back to the logical K dimension.D = A * B [+ D]accumulates into TMEM. Seetcgen05_mma()for the shared single-thread-issue / elect-safe / commit semantics, A/B descriptor construction, and thecollector_op/a_shift/scale_input_d/write_disable_maskparameters.- Parameters:
mma_kind (Tcgen05MMAKind) – Data-type/kind selector (see
tcgen05_mma()).cta_group (CTAGroup) –
'CTA_1'or'CTA_2'(seetcgen05_mma()).d (cutlass.Array or cutlass.Pointer) – TMEM accumulator pointer. Validated at trace time to reside in tensor memory (dialect operand type
LLVM_PointerTensor).a (cutlass.Array, cutlass.Pointer, or cutlass.Int64) – A operand – SMEM descriptor (Int64) or TMEM pointer.
b (int, cutlass.Int64, or cutlass.Uint64) – 64-bit SMEM descriptor for B.
idesc (int, cutlass.Int32, or cutlass.Uint32) – Packed 32-bit instruction descriptor (sparse bit set).
enable_input_d (int or cutlass.Boolean) –
D = A*B + Dwhen true, elseD = A*B.sparse_metadata (cutlass.Array or cutlass.Pointer) – TMEM pointer to the sparsity metadata mapping the K/2 packed columns to the logical K dimension. Validated at trace time to reside in tensor memory.
collector_op – A-operand collector policy (see
tcgen05_mma()).scale_input_d –
[0, 15]; valid only for"f16"/"tf32"(validated at trace time). Seetcgen05_mma().write_disable_mask – 4-element (
CTA_1) / 8-element (CTA_2) per-row TMEM write-disable vector (validated at trace time).
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_mma_sp_block_scale(
- mma_kind: ~cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMAKind,
- cta_group: ~cutlass.experimental.primitives.nvvm_wrapper.CTAGroup,
- d: ~cutlass.Array | ~cutlass.Pointer,
- a: ~cutlass.Array | ~cutlass.Pointer | ~cutlass.Int64,
- b: int | ~cutlass.Int64 | ~cutlass.Uint64,
- idesc: int | ~cutlass.Int32 | ~cutlass.Uint32,
- enable_input_d: int | ~cutlass.Boolean,
- sparse_metadata: ~cutlass.Array | ~cutlass.Pointer,
- scale_a: ~cutlass.Array | ~cutlass.Pointer,
- scale_b: ~cutlass.Array | ~cutlass.Pointer,
- *,
- scale_vec_size: ~cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMAScaleVecSize | _install_cutlass_mlir_autodoc_stub.<locals>._DocDialectObject | None = None,
- collector_op: ~cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMACollectorOp | None = None,
- a_shift: bool | None = None,
Structured-sparse MMA with per-block scale factors (
tcgen05.mma.sp.block_scale).Combines the structured-sparse A path of
tcgen05_mma_sp()with the per-block scaling oftcgen05_mma_block_scale(): A is sparse withsparse_metadatain TMEM, and blocks of A/B are scaled by thescale_a/scale_bfactor tiles (also in TMEM) before accumulation. Seetcgen05_mma_block_scale()for the block-scale descriptor / scale-vector details andtcgen05_mma()for the shared issue/commit semantics.- Parameters:
mma_kind (Tcgen05MMAKind) – Block-scale kind; validated at trace time to be a block-scaled kind (
mxf8f6f4/mxf4/mxf4nvf4).cta_group (CTAGroup) –
'CTA_1'or'CTA_2'.d – TMEM accumulator pointer; validated to reside in tensor memory.
a – A operand – SMEM descriptor or TMEM pointer.
b – 64-bit SMEM descriptor for B.
idesc – Packed 32-bit instruction descriptor.
enable_input_d –
D = A*B*scale + Dwhen true, elseD = A*B*scale.sparse_metadata – TMEM pointer to the sparsity metadata; validated to reside in tensor memory.
scale_b (scale_a,) – TMEM addresses of the A/B scale-factor tiles; both validated at trace time to reside in tensor memory.
scale_vec_size – 1X / 2X / 4X scale-vector packing (see
tcgen05_mma_block_scale()).collector_op – A-operand collector policy (see
tcgen05_mma()).
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_mma_ws(
- mma_kind: Tcgen05MMAKind,
- d: Array | Pointer,
- a: Array | Pointer | Int64,
- b: int | Int64 | Uint64,
- idesc: int | Int32 | Uint32,
- enable_input_d: int | Boolean,
- *,
- collector_b_buffer: Tcgen05MMACollectorBBuffer | None = None,
- collector_op: Tcgen05MMACollectorOp | None = None,
- col_b_zero_mask: int | Int64 | Uint64 | None = None,
Issue a weight-stationary tcgen05 MMA.
Emits
tcgen05.mma.ws.cta_group::1. The instruction initiatesD = A*B+Dfor a dense A matrix and uses a B-matrix collector buffer for weight-stationary convolution-style reuse. A may be an SMEM descriptor or a TMEM pointer; B is an SMEM descriptor. Whenenable_input_dis false, the operation computesD = A*B.- Parameters:
mma_kind (Tcgen05MMAKind) – Data-type/kind selector. Public PTX ISA 9.3 forms support
F16,TF32,F8F6F4, andINT8fortcgen05.mma.ws. (The dialect’s non-block-scale kind attribute rejects block-scaled kinds; usetcgen05_mma_block_scale()for those.)d (cutlass.Array or cutlass.Pointer) – TMEM accumulator destination. Validated at trace time to reside in tensor memory (dialect operand type
LLVM_PointerTensor).a (cutlass.Array, cutlass.Pointer, or cutlass.Int64) – A operand as an SMEM descriptor or TMEM pointer.
b (int, cutlass.Int64, or cutlass.Uint64) – 64-bit SMEM descriptor for B.
idesc (int, cutlass.Int32, or cutlass.Uint32) – Packed 32-bit instruction descriptor.
enable_input_d (int or cutlass.Boolean) – When true, accumulate into existing D; when false, compute
D = A*B.collector_b_buffer (Tcgen05MMACollectorBBuffer, optional) – Optional B collector buffer selector (
B0throughB3). PTX defaults toB0withDISCARDwhen no collector usage is specified.collector_op (Tcgen05MMACollectorOp, optional) – Optional B collector operation:
FILL,USE,LASTUSE, orDISCARD.col_b_zero_mask (int, cutlass.Int64, or cutlass.Uint64, optional) – Optional zero-column-mask descriptor for B columns that should be treated as zero regardless of SMEM contents.
if nvvm.elect_sync(): nvvm.tcgen05_mma_ws( nvvm.Tcgen05MMAKind.INT8, d_tmem, a_tmem, b_desc, idesc, k > 0, collector_b_buffer=nvvm.Tcgen05MMACollectorBBuffer.B2, collector_op=nvvm.Tcgen05MMACollectorOp.USE, )
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_mma_ws_sp(
- mma_kind: Tcgen05MMAKind,
- d: Array | Pointer,
- a: Array | Pointer | Int64,
- b: int | Int64 | Uint64,
- idesc: int | Int32 | Uint32,
- enable_input_d: int | Boolean,
- sparse_metadata: Array | Pointer,
- *,
- collector_b_buffer: Tcgen05MMACollectorBBuffer | None = None,
- collector_op: Tcgen05MMACollectorOp | None = None,
- col_b_zero_mask: int | Int64 | Uint64 | None = None,
Issue a sparse weight-stationary tcgen05 MMA.
Emits
tcgen05.mma.ws.sp.cta_group::1. The instruction initiatesD = A*B+Dwhere A is a structured sparse matrix packed asM x (K/2)and accompanied by sparse metadata in TMEM. A may be an SMEM descriptor or a TMEM pointer; B is an SMEM descriptor. Whenenable_input_dis false, the operation computesD = A*B.- Parameters:
mma_kind (Tcgen05MMAKind) – Data-type/kind selector. Public PTX ISA 9.3 forms support
F16,TF32,F8F6F4, andINT8fortcgen05.mma.ws.sp.d (cutlass.Array or cutlass.Pointer) – TMEM accumulator destination.
a (cutlass.Array, cutlass.Pointer, or cutlass.Int64) – A operand as an SMEM descriptor or TMEM pointer.
b (int, cutlass.Int64, or cutlass.Uint64) – 64-bit SMEM descriptor for B.
idesc (int, cutlass.Int32, or cutlass.Uint32) – Packed 32-bit instruction descriptor.
enable_input_d (int or cutlass.Boolean) – When true, accumulate into existing D; when false, compute
D = A*B.sparse_metadata (cutlass.Array or cutlass.Pointer) – TMEM address or pointer for sparse-A metadata.
collector_b_buffer (Tcgen05MMACollectorBBuffer, optional) – Optional B collector buffer selector (
B0throughB3). PTX defaults toB0withDISCARDwhen no collector usage is specified.collector_op (Tcgen05MMACollectorOp, optional) – Optional B collector operation:
FILL,USE,LASTUSE, orDISCARD.col_b_zero_mask (int, cutlass.Int64, or cutlass.Uint64, optional) – Optional zero-column-mask descriptor for B columns that should be treated as zero regardless of SMEM contents.
if nvvm.elect_sync(): nvvm.tcgen05_mma_ws_sp( nvvm.Tcgen05MMAKind.TF32, d_tmem, a_tmem, b_desc, idesc, k > 0, sparse_metadata_tmem, collector_b_buffer=nvvm.Tcgen05MMACollectorBBuffer.B1, collector_op=nvvm.Tcgen05MMACollectorOp.FILL, )
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_relinquish_alloc_permit(
- *,
- group: CTAGroup | None = None,
Release the SM-level TMEM allocation permit after
tcgen05_alloc.Emits
tcgen05.relinquish_alloc_permit.sync.aligned. The.sync.alignedqualifier means this is a warp-collective instruction: all 32 threads of the warp must execute it simultaneously.CTA_2 placement: call from a warp that runs on both CTAs at a convergence point (e.g. warp 0 right after
barrier_cluster_waitand before any warp-role branch that may diverge between CTAs).- Parameters:
group –
'CTA_1'(default) or'CTA_2'. SeeCTAGroup.
# CTA_1: inside the MMA warp, outside elect_sync (both alloc and # relinquish are .sync.aligned — all 32 threads must participate). if warp == MMA_WARP: nvvm.tcgen05_alloc(tmem_ptr, num_cols, group="cta_1") nvvm.tcgen05_relinquish_alloc_permit(group="cta_1") # CTA_2: from warp 0 BEFORE warp-role branches — ensures both CTAs converge if warp == 0: nvvm.tcgen05_alloc(tmem_ptr, num_cols, group="cta_2") nvvm.barrier_cluster_wait() nvvm.barrier_cta_sync() # Both CTAs' warp 0 reach here simultaneously → safe collective call: if warp == 0: nvvm.tcgen05_relinquish_alloc_permit(group="cta_2") # Only NOW diverge into warp-specialized roles (TMA / MMA / epilogue)
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_shift( ) None#
Shift TMEM rows down by one within each sub-partition (
tcgen05.shift).Asynchronous instruction that shifts the 32-byte elements of the matrix at
taddrdownwards by one row across all rows except the last of each sub-partition. Used to advance a sliding-window operand in TMEM (e.g. the A operand of a chainedtcgen05_mma()). Staged into the tensor-core issue queue alongsidetcgen05.mma/tcgen05.cp; pair withtcgen05_commit()(or a downstream consumer) to observe completion.- Parameters:
taddr (cutlass.Array or cutlass.Pointer) – TMEM base pointer of the matrix whose rows are shifted. Validated at trace time: must reside in tensor memory when the passed object exposes
.space(dialect operand typeLLVM_PointerTensor).group (CTAGroup, optional) –
'CTA_1'(default) or'CTA_2'. Selects the single-CTA vs 2-CTA shift; must match the group of the othertcgen05ops in the kernel. SeeCTAGroup.
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_st(
- shape: Tcgen05LdStShape,
- tmem_addr: Array | Pointer,
- val: Int32 | Vector,
- *,
- unpack: bool | None = None,
- offset: int | Int64 | Uint64 | None = None,
Store register data into TMEM.
Emits
tcgen05.st.sync.aligned.{shape}{.xN}{.unpack}.b32. The issuing warp collectively stores into TMEM; each lane supplies its own register slice for the warp’s accessible TMEM sub-partition. Mirror oftcgen05_ld(): same shape literals, same address encoding, same TMEM access restriction (seetcgen05_ld()for the canonical warp-position-to-lane-range table and theM_TILE=128epilogue / writer implications). Read both docstrings together —tcgen05_ld()is the source of truth for the access table; this docstring covers only the store-direction differences.Accepts any
Vectoror scalar for val; non-Int32values are bitcast toInt32internally (the hardware requires i32 register words) — no manual.bitcast(Int32)wrapping needed at the call site.- Parameters:
shape (Tcgen05LdStShape) – Store shape (PTX
.shape1/.shape2). One of"16x32bx2","16x64b","32x32b","16x128b","16x256b". Same access restriction astcgen05_ld().tmem_addr (cutlass.Array or cutlass.Pointer) – Pointer/Array to the TMEM location;
(row << 16) | colencoding (seetcgen05_ld()for full encoding details). Validated at trace time: must reside in tensor memory when the passed object exposes.space(the dialect operand type isLLVM_PointerTensor); opaque pointers are deferred to the verifier.val (cutlass.Int32 or cutlass.Vector) – Value to store.
Int32scalar orVector[any dtype]; non-Int32 dtypes are auto-bitcast toInt32internally.unpack (bool, optional) – Enable element unpacking — mirror of
packontcgen05_ld().offset (int or cutlass.Int64 or cutlass.Uint64, optional) – Required (and only valid) for
"16x32bx2"— column offset added totmem_addrat runtime. Must beNonefor all other shapes.
- Raises:
ValueError –
shapeis not a recognized literal;offsetis set for a shape other than"16x32bx2"(or missing for"16x32bx2").TypeError –
tmem_addrexposes a.spacethat is not tensor memory (TMEM).
# Store a vector of accumulator data back into TMEM. nvvm.tcgen05_st("32x32b", tmem_ptr, data_vec)
- cutlass.experimental.primitives.nvvm_wrapper.tcgen05_wait(
- kind: Tcgen05Wait,
Wait for pending TMEM load or store operations to complete.
Emits
tcgen05.wait::{ld|st}.sync.aligned. Place aftertcgen05_ld()(LOAD) to ensure the TMEM→register transfer has finished before any thread reads the returned register values, or aftertcgen05_st()(STORE) to ensure register→TMEM writes are visible before a downstreamtcgen05_mma()reads them.- Parameters:
kind –
"load"or"store".
# Epilogue: read accumulator from TMEM into registers c_vec = nvvm.tcgen05_ld(shape, tmem_addr, num=num_cols) nvvm.tcgen05_wait("load") # c_vec is now safe to use
- cutlass.experimental.primitives.nvvm_wrapper.tensormap_cp_fenceproxy( ) None#
Copy a TMA descriptor with a proxy fence at the requested scope.
Emits
tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.<scope>.sync.aligned [dst], [src], size;. Copiessizebytes of TMA descriptor data fromsrctodstand acts as a release fence between the generic and tensormap proxies, so subsequent TMA ops see the new descriptor at the chosen visibility scope (cta/cluster/gpu/sys).- Parameters:
dst (cutlass.Array or cutlass.Pointer) – TMA-descriptor destination pointer / array (typically a constant-banked descriptor slot).
src (cutlass.Array or cutlass.Pointer) – TMA-descriptor source pointer / array.
size (int or cutlass.Int32 or cutlass.Uint32) – Number of bytes to copy; the descriptor size (128 for standard TMA descriptors).
scope (MemScope) – Visibility scope of the release fence.
- Raises:
ValueError – if a statically known
sizeis not 128.
- cutlass.experimental.primitives.nvvm_wrapper.tensormap_replace(
- field: TensormapField,
- addr: Array | Pointer,
- *,
- new_value: int | Int32 | Int64 | None = None,
- ord: int | None = None,
- new_value_attr: TensormapElemtype | TensormapInterleaveLayout | TensormapSwizzleMode | TensormapSwizzleAtomicity | TensormapFillMode | None = None,
Replace one field of an in-memory TMA tensor-map descriptor.
Emits
tensormap.replaceagainst the descriptor ataddr(global or shared memory). Used to patch a TMA descriptor at runtime, e.g. swap theglobal_addressor resize aglobal_dimbetween launches without rebuilding the whole descriptor on the host.Fields split into two groups:
integer-valued (pass
new_value):global_address,rank,box_dim,global_dim,global_stride,element_stride.global_address/global_strideare 64-bit; the rest are 32-bit (coerced for you).ranktakes one less than the desired tensor rank (zero-based).enum-valued (pass
new_value_attr):elemtype,interleave_layout,swizzle_mode,swizzle_atomicity,fill_mode, each taking the matching enum.
The
ord(dimension ordinal) is required forbox_dim,global_dim,global_stride, andelement_stride, and rejected for all other fields.- Parameters:
field (TensormapField) – Which descriptor field to replace.
addr – Pointer/Array to the tensor-map descriptor (global or shared).
new_value – New value for an integer-valued field (coerced to
Int64orInt32by field width). Mutually exclusive withnew_value_attr.ord – Dimension ordinal of the field across the tensor; required for
box_dim/global_dim/global_stride/element_strideand rejected otherwise. The valid range is enforced by the dialect verifier for the target build.new_value_attr – New value for an enum-valued field, as the matching enum (or its string). Mutually exclusive with
new_value.
- Raises:
ValueError – if
fieldis not aTensormapField(raw NVVM dialect enums are rejected); if the wrong one ofnew_value/new_value_attris supplied for the field; ifordis supplied for a field that forbids it or omitted for one that requires it, or is outside the dialect’s valid range; or ifnew_value_attris not the enum type the field expects.TypeError – if
ordis supplied but is not a Pythonint.
# Patch the source pointer of a copied TMA descriptor at runtime. if nvvm.elect_sync(): nvvm.tensormap_replace( nvvm.TensormapField.GLOBAL_ADDRESS, desc_ptr, new_value=cutlass.Int64(new_base), ) nvvm.tensormap_replace( nvvm.TensormapField.GLOBAL_DIM, desc_ptr, new_value=cutlass.Int32(new_dim), ord=0, )
- cutlass.experimental.primitives.nvvm_wrapper.trace_mark(
- event_type: int,
- domain: str,
- event: str,
- *,
- payload: Int32 | Int64 | None = None,
- payload_descriptor: str | None = None,
Wrapper over
nvvm.trace_mark.
- cutlass.experimental.primitives.nvvm_wrapper.vote_sync( ) Int32 | Boolean#
Perform a collective warp vote across masked lanes.
Maps to
vote.sync.{all|any|uni|ballot}— each participating lane contributes a boolean predicate and all receive a shared result. All lanes named in mask must be actively executing the instruction (convergence requirement).VoteSyncKind variants:
"all"→Boolean:Trueiff every masked lane setpred = True."any"→Boolean:Trueiff at least one masked lane setpred = True."uni"→Boolean:Trueiff all masked lanes cast the same vote (all-True or all-False). Use to detect uniform control flow without requiring all lanes to be true."ballot"→Int32: 32-bit bitmask where bit i is set when lane i setpred = True. Bit i is 0 for lanes not in mask.
- Parameters:
mask (int or cutlass.Int32 or cutlass.Uint32) – 32-bit member mask; bit i set means lane i participates. All named lanes must execute the instruction. Pass
0xFFFFFFFFfor a full-warp vote.pred (int or cutlass.Boolean) – Per-lane boolean vote input — what each lane contributes to the collective result. Pass
~predto vote on the negation.kind (VoteSync) – Vote mode — determines the semantics and return type.
- Returns:
Booleanforall/any/uni;Int32bitmask forballot.- Return type:
- Raises:
ValueError – if a static mask does not fit in 32 bits, or if kind is not one of
"all","any","uni","ballot". A runtime (non-int) mask is not checked at trace time.
Constraints:
All lanes named in mask must reach the instruction; a divergent lane causes the remaining lanes to stall indefinitely.
Available on SM30+;
ballotrequires SM35+.This wrapper has no
exec_predargument.
tx, _, _ = cute.arch.thread_idx() is_even = (tx % cutlass.Int32(2)) == cutlass.Int32(0) # all: True only when every lane voted True all_even = nvvm.vote_sync(0xFFFFFFFF, is_even, "all") # any: True when at least one lane voted True any_even = nvvm.vote_sync(0xFFFFFFFF, is_even, "any") # uni: True when all lanes agree (all-True OR all-False) uniform = nvvm.vote_sync(0xFFFFFFFF, is_even, "uni") # ballot: bitmask of lanes that voted True (even lanes → 0x55555555) even_mask = nvvm.vote_sync(0xFFFFFFFF, is_even, "ballot")
- cutlass.experimental.primitives.nvvm_wrapper.wgmma_commit_group_sync_aligned() None#
Commit outstanding async warpgroup MMAs into a group.
Emits
wgmma.commit_group.sync.aligned. Bundles allwgmma.mma_asyncoperations issued by the warpgroup since the last commit into a new group, whichwgmma_wait_group_sync_aligned()can later wait on.nvvm.wgmma_commit_group_sync_aligned() nvvm.wgmma_wait_group_sync_aligned(0)
- cutlass.experimental.primitives.nvvm_wrapper.wgmma_fence_aligned() None#
Fence register accesses around async warpgroup MMA.
Emits
wgmma.fence.sync.aligned. Orders the executing warpgroup’s accesses to the registers/shared memory that feedwgmma.mma_async, so the async MMA observes a consistent view. Issue once before the firstwgmma.mma_asyncof a sequence (and after writing its inputs).nvvm.wgmma_fence_aligned() # ... wgmma.mma_async issues ... nvvm.wgmma_commit_group_sync_aligned()
- cutlass.experimental.primitives.nvvm_wrapper.wgmma_mma_async(
- results_: _install_cutlass_mlir_autodoc_stub.<locals>._DocMlirType,
- inouts: ir.Value,
- descriptor_a: int | ~cutlass.Int64 | ~cutlass.Uint64,
- descriptor_b: int | ~cutlass.Int64 | ~cutlass.Uint64,
- shape: ~sphinx.ext.autodoc.mock._MockObject,
- type_a: ~cutlass.experimental.primitives.nvvm_wrapper.WGMMAType,
- type_b: ~cutlass.experimental.primitives.nvvm_wrapper.WGMMAType,
- type_d: ~cutlass.experimental.primitives.nvvm_wrapper.WGMMAType,
- scale_d: ~cutlass.experimental.primitives.nvvm_wrapper.WGMMAScaleOut,
- scale_a: ~cutlass.experimental.primitives.nvvm_wrapper.WGMMAScaleIn,
- scale_b: ~cutlass.experimental.primitives.nvvm_wrapper.WGMMAScaleIn,
- layout_a: ~cutlass.experimental.primitives.nvvm_wrapper.MMALayout,
- layout_b: ~cutlass.experimental.primitives.nvvm_wrapper.MMALayout,
- *,
- satfinite: ~cutlass.experimental.primitives.nvvm_wrapper.MMAIntOverflow | None = None,
Wrapper over
nvvm.wgmma_mma_async.Returns an LLVM struct. Caller provides the raw MLIR result type as results_.
- cutlass.experimental.primitives.nvvm_wrapper.wgmma_wait_group_sync_aligned(group: int) None#
Wait until at most
groupasync warpgroup-MMA groups are pending.Emits
wgmma.wait_group.sync.aligned. Blocks the warpgroup until no more thangrouppreviously-committed MMA groups remain in flight;group=0waits for all of them. Registers written by completed MMAs are safe to read afterwards.- Parameters:
group (int) – Maximum number of committed groups allowed to remain pending. Must be a non-negative
int.- Raises:
ValueError –
groupis a negativeint.
nvvm.wgmma_wait_group_sync_aligned(0) # drain all pending MMAs
- cutlass.experimental.primitives.nvvm_wrapper.wmma_load(
- res: _install_cutlass_mlir_autodoc_stub.<locals>._DocMlirType,
- ptr: ~cutlass.Array | ~cutlass.Pointer,
- stride: int | ~cutlass.Int32 | ~cutlass.Uint32,
- m: int,
- n: int,
- k: int,
- layout: ~cutlass.experimental.primitives.nvvm_wrapper.MMALayout,
- eltype: ~cutlass.experimental.primitives.nvvm_wrapper.MMAType,
- frag: ~cutlass.experimental.primitives.nvvm_wrapper.MMAFrag,
Wrapper over
nvvm.wmma_load.Returns an LLVM struct. Caller provides the raw MLIR result type as res.
- cutlass.experimental.primitives.nvvm_wrapper.wmma_store(
- ptr: Array | Pointer,
- m: int,
- n: int,
- k: int,
- layout: MMALayout,
- eltype: MMAType,
- args: ir.Value,
- stride: int | Int32 | Uint32,
Wrapper over
nvvm.wmma_store.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05LdStShape(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMAScaleVecSize#
Bases:
objectCompatibility namespace for tcgen05 block-scale selectors.
Older call sites spell the selector as
Tcgen05MMAScaleVecSizewhile newer NVVM bindings split 1X/2X/4X and block16/block32 into separate enums. Keep both spellings available at the wrapper boundary.- X1#
alias of
_MockObject
- X2#
alias of
_MockObject
- X4#
alias of
_MockObject
- DEFAULT#
alias of
_MockObject
- Default#
alias of
_MockObject
- BLOCK16#
alias of
_MockObject
- BLOCK32#
alias of
_MockObject
- class cutlass.experimental.primitives.nvvm_wrapper.TensormapElemtype(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.TensormapField(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.TensormapFillMode(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.TensormapInterleaveLayout(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.TensormapSwizzleAtomicity(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.TensormapSwizzleMode(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.AtomicOp(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.BarrierRedux(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.CTAGroup(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.CVTPackFloat(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.CacheLevel(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.ClusterLaunchControlQueryType(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.CmpOp(value)#
Bases:
StrEnumPTX comparison operator for
setp(PTX ISA §9.7.6).
- class cutlass.experimental.primitives.nvvm_wrapper.ConvertScale(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.CpReduceOp(value)#
Bases:
StrEnumReduction op for
cp.reduce.async.bulk(non-TMA, PTX §9.7.9.25.4.2).
- class cutlass.experimental.primitives.nvvm_wrapper.CpReduceType(value)#
Bases:
StrEnumElement type for
cp.reduce.async.bulk(non-TMA, PTX §9.7.9.25.4.2).
- class cutlass.experimental.primitives.nvvm_wrapper.CvtaSize(value)#
Bases:
StrEnumAddress-width qualifier for
cvta:.u32or.u64.
- class cutlass.experimental.primitives.nvvm_wrapper.CvtaSpace(value)#
Bases:
StrEnumTarget address space for
cvta(PTX ISA §9.7.9.20).
- class cutlass.experimental.primitives.nvvm_wrapper.DotAccumulateType(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.EvictPriority(value)#
Bases:
StrEnumEviction priority hint for cache lines.
Controls the priority of cache line eviction.
Reference: PTX ISA - Cache Eviction Priority
- class cutlass.experimental.primitives.nvvm_wrapper.FPRoundingMode(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.GridDepAction(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.IntRoundingMode(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.L1EvictKind(value)#
Bases:
StrEnumL1 cache eviction priority hint.
Controls the eviction policy for L1 cache lines. Members match the 5-member subset of
EvictPrioritythat the L1 path supports (the L2 path adds the sm_90+*_demote/*_nearvariants).Reference: PTX ISA - Cache Eviction Priority https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#cache-eviction-priority
- class cutlass.experimental.primitives.nvvm_wrapper.L2PrefetchSize(value)#
Bases:
StrEnumL2 cache prefetch size hint.
Specifies the prefetch granularity for L2 cache operations.
Reference: PTX ISA - Data Movement Instructions https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-prefetch-prefetchu
- class cutlass.experimental.primitives.nvvm_wrapper.LoadCacheModifier(value)#
Bases:
StrEnumCache operation modifier for load instructions.
Controls L1/L2 cache behavior for memory loads.
Reference: PTX ISA - Cache Operators https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#cache-operators
- PTX cache operators for ld:
.ca - Cache at all levels (L1 and L2), default .cg - Cache at global level (L2 only, bypass L1) .cs - Cache streaming (likely accessed once, evict first) .lu - Last use (hint data won’t be needed again) .cv - Cache volatile (don’t cache, always fetch from memory)
- class cutlass.experimental.primitives.nvvm_wrapper.LoadShape(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.LoadSrcFormat(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MBarrierScope(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MBarrierWait(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MMAB1Op(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MMAFrag(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MMAIntOverflow(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MMAKind(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MMALayout(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MMAType(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MatchSync(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MemOrder(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.MemScope(value)#
Bases:
StrEnumMemory scope for memory operations.
Controls which threads observe the memory operation effects. Used for loads, stores, atomics, and memory barriers.
Reference: PTX ISA - Scope https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#scope
- PTX scope qualifiers:
.cta - Threads within the same CTA (thread block) .cluster - Threads within the same cluster .gpu - All threads on the same GPU device .sys - All threads in the system (including host CPU)
- class cutlass.experimental.primitives.nvvm_wrapper.MulMode(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.PermuteMode(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Proxy(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.ReductionOp(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.ReductionType(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.SaturationMode(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.SaturationModeKind(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.SetMaxRegisterAction(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.SetpType(value)#
Bases:
StrEnumSource type suffix for
setp/selp(PTX ISA §9.7.6).
- class cutlass.experimental.primitives.nvvm_wrapper.Shfl(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.StoreCacheModifier(value)#
Bases:
StrEnumCache operation modifier for store instructions.
Controls L1/L2 cache behavior for memory stores.
Reference: PTX ISA - Cache Operators https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#cache-operators
- PTX cache operators for st:
.wb - Write-back (cache at all levels), default .cg - Cache at global level (L2 only, bypass L1) .cs - Cache streaming (likely accessed once) .wt - Write-through (write to memory immediately)
- class cutlass.experimental.primitives.nvvm_wrapper.StoreShape(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.TMALoadMode(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.TMARedux(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.TMAStoreMode(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05CpMulticast(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05CpShape(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05CpSrcFormat(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05Fence(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMACollectorBBuffer(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMACollectorOp(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05MMAKind(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.Tcgen05Wait(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.TensormapSpace(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.VoteSync(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.WGMMAScaleIn(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.WGMMAScaleOut(value)#
Bases:
StrEnumAn enumeration.
- class cutlass.experimental.primitives.nvvm_wrapper.WGMMAType(value)#
Bases:
StrEnumAn enumeration.
- cutlass.experimental.primitives.nvvm_wrapper.barrier(
- *,
- barrier_id: cutlass.cute.typing.Int | None = None,
- number_of_threads: cutlass.cute.typing.Int | None = None,
Creates a barrier, optionally named.
- class cutlass.experimental.primitives.nvvm_wrapper.S2TCopyMode#
Bases:
objectS2T (SMEM->TMEM) copy mode enumeration for
prims.tcgen05_cp.Combines shape and multicast into valid configurations for SMEM-to-TMEM copy. Each mode specifies both the data shape and the required warp broadcast pattern.
Available modes: - S2T_128x256b: 128 rows x 256 bits, no multicast - S2T_128x128b: 128 rows x 128 bits, no multicast - S2T_4x256b: 4 rows x 256 bits, no multicast - S2T_32x128b_WARPX4: 32 rows x 128 bits, broadcast to all 4 warps - S2T_64x128b_WARPX2_01_23: 64 rows x 128 bits, broadcast to warp pairs (0,1)(2,3) - S2T_64x128b_WARPX2_02_13: 64 rows x 128 bits, broadcast to warp pairs (0,2)(1,3)
- cutlass.experimental.primitives.nvvm_wrapper.make_tmem_ptr( ) Array#
Convert a TMEM address to a typed TMEM
Arrayview (address space 6).- Parameters:
tmem_addr – The TMEM address value (
intor a DSL integer).dtype – The element type for the returned view.
- Returns:
An
Arrayover TMEM (address space 6).
- cutlass.experimental.primitives.nvvm_wrapper.make_tmem_ptr_from_warp_row_col( ) Array#
Build a typed TMEM
Arrayview for TMEM/SPwarpat rowwarp*32.Each warp in a tcgen05 MMA group owns one sub-partition of the TMEM accumulator. The canonical epilogue formula
tmem_sp = warp_idx % 4 tmem_addr = (tmem_base_row + tmem_sp * 32) << 16 | base_col
is bundled here so callers don’t reassemble the bitfield by hand. For shifted epilogue ranges such as warps 2..5, pass
warp_idx % 4rather than the logical epilogue rank.- Parameters:
tmem_base – TMEM address of the accumulator base (row 0, col 0); bits [0:16) hold the starting column, [16:32) hold the starting row.
warp – TMEM/SP index (0..3), usually
warp_idx % 4.base_col – Column offset within the accumulator.
dtype – Element type of the returned TMEM view.
- Returns:
An
Arrayover TMEM at(base_row + warp*32, base_col).