nvgpu Dialect Overview

Abstract

nvgpu is the bridge dialect between MLIR's generic gpu dialect and NVPTX-specific nvvm. It names the NVIDIA kernel patterns that gpu cannot express — warp shuffle, MMA and WGMMA, cp.async, mbarrier, TMA — without committing yet to a concrete NVVM intrinsic. Tileiras links the upstream dialect unchanged. cute_nvgpu feeds it from above; convert-nvgpu-to-nvvm drains it from below.

About thirty ops live here. The conversion pass installs one OpConversionPattern per op and rewrites the module in a single sweep, each pattern emitting a small fixed body of nvvm.* ops — or, for a few exception cases, expanded memref / llvm / llvm.inline_asm. The pass mnemonic is convert-nvgpu-to-nvvm; in the O3 pipeline it runs immediately after the broad convert-to-llvm step and before convert-vector-to-llvm, arith-expand, and convert-memref-to-llvm (see Pass List by Optimization Level — O3), so by the time it fires every operand is already in LLVM-dialect or memref form.

Position in the Cascade

cute_nvgpu
    |
    | lower architecture atoms into stock GPU operations
    v
nvgpu
    |
    | convert-nvgpu-to-nvvm: ~30 patterns, one sweep
    v
nvvm
    |
    | translate to LLVM IR and the NVPTX backend
    v
PTX

cute_nvgpu ops still speak SM-tier vocabulary — TMA atoms, WGMMA atoms, Blackwell tensor-memory operations. nvgpu strips the source-level atom naming and re-presents the same behaviour over MLIR memrefs, vectors, descriptors, barrier groups, and async tokens. That makes the NVVM conversion mechanical: every nvgpu op below has a fixed nvvm (or llvm.inline_asm) lowering.

Op Roster

populateNVGPUToNVVMConversionPatterns installs one OpConversionPattern per op. Tileiras links the upstream populator unchanged. The rewriter callbacks branch on source memory space to pick the generic or .shared form of the mbarrier and cp.async intrinsics — address space 3 always selects .shared.

The "Status" column distinguishes ops whose mnemonic string appears verbatim in this binary's string table from upstream-MLIR patterns that the linked populator carries but whose mnemonic was never interned (either because the op was renamed, dropped, or only reached through gpu-dialect routing).

Patterns are registered at benefit = 1. The 64-bit values consumed by nvgpu.warpgroup.mma's descriptorA / descriptorB operands are the same SMEM descriptor words that flow through cute_nvgpu MMA atoms — the canonical bitfield decode is on the cute_nvgpu MMA atoms page.

Operand and Attribute Tables

The tables below pin every op family in the dialect to its operand list, attribute list, result list, and the NVVM rewrite the conversion pattern emits. SM gating lives in the per-arch availability table further down.

nvgpu.device_async_copy

Rewriter emits nvvm.cp.async.shared.global with cp_size = dstElements * eltbytes and cp_modifier = bypassL1 ? cg : ca.

nvgpu.device_async_create_group

Rewriter emits a single nvvm.cp.async.commit.group. Input tokens are erased; the SSA edge survives only as a happens-before constraint.

nvgpu.device_async_wait

Rewriter emits nvvm.cp.async.wait.group N where N = numGroups (default 0).

nvgpu.mbarrier.create

Rewriter emits no nvvm.* op. It generates a memref.global "private" @__mbarrier ... : memref<NxNumBarrier x i64, 3> and a memref.get_global returning a base pointer.

nvgpu.mbarrier.init.shared (alias of mbarrier.init over addr-space 3)

Rewriter emits nvvm.mbarrier.init.shared against the GEP-resolved slot address.

nvgpu.mbarrier.arrive

Rewriter emits nvvm.mbarrier.arrive or nvvm.mbarrier.arrive.shared based on the slot's address space.

nvgpu.mbarrier.arrive.expect_tx

Rewriter emits nvvm.mbarrier.arrive.expect_tx[.shared]. No SSA result; the side effect is on the shared-memory mbarrier slot.

nvgpu.mbarrier.try_wait.parity

Rewriter emits nvvm.mbarrier.try_wait.parity.shared returning an i1 polled in a loop.

nvgpu.mbarrier.inval (absent in this binary)

The mnemonic nvgpu.mbarrier.inval is not interned in this tileiras string table; the inval-side intrinsic (nvvm.mbarrier.inval.shared) is still present and is reached through the lower-level NVVM lowering. The operand list below is documented for upstream-MLIR parity and as a reference for reimplementers that choose to surface the wrapper.

Upstream rewriter shape: emit nvvm.mbarrier.inval[.shared].

nvgpu.tma.async.load

Rewriter emits a single nvvm.cp.async.bulk.tensor.shared.global. See Lowering: TMA Async Load — Operand Mapping for the operand-slot mapping.

nvgpu.tma.async.store

Rewriter emits nvvm.cp.async.bulk.tensor.global.shared. No barrier — the producer side does not wait.

nvgpu.tma.prefetch.descriptor

Rewriter emits nvvm.prefetch.tensormap [%tmap].

nvgpu.tma.fence.descriptor

Rewriter emits nvvm.fence.proxy.acquire.sync.cluster — the proxy-acquire fence that the WGMMA descriptor consumer needs.

nvgpu.tma.async.reduce (absent in this binary)

The nvgpu.tma.async.reduce mnemonic is not interned in this tileiras build. The underlying NVVM op (nvvm.cp.async.bulk.tensor.reduce) is present and consumed by cute_nvgpu lowerings directly; no nvgpu wrapper surfaces the reduce variant. The operand layout below documents the upstream wrapper for parity.

Upstream rewriter shape: emit nvvm.cp.async.bulk.tensor.reduce with red_op decoded from the attribute.

nvgpu.tma.create.descriptor

Rewriter emits no nvvm.* op. It allocates a 128-byte CUtensorMap on the host stack via llvm.alloca, fills it through llvm.getelementptr + llvm.store, and calls cuTensorMapEncodeTiled.

nvgpu.tensormap.create.descriptor (device-side replace path; absent in this binary)

This op family is not interned in this tileiras build. The device-side descriptor replace path is reached directly through cute_nvgpu -> nvvm.tensormap.* without going through an nvgpu.tensormap.create.descriptor wrapper. Operand layout documented below for upstream-MLIR parity.

Upstream rewriter shape: emit nvvm.tensormap.cp.async.shared followed by a sequence of nvvm.tensormap.replace.* ops.

nvgpu.tensormap.update.global_address / box_dim / element_stride (absent in this binary)

These per-field update wrappers are also not interned in this build. Field-level descriptor updates lower directly through the nvvm.tensormap.replace.* family.

Upstream rewriter shape: each maps to the matching nvvm.tensormap.replace.* op against the SMEM-resident descriptor.

nvgpu.warpgroup.descriptor (also spelled warpgroup.generate.descriptor)

Rewriter packs the descriptor bits inline. The result is a i64 LLVM value built by an shl/or chain; the bit layout (start_addr[14] | lbo[16] | sbo[16] | base_offset[3] | reserved[3] | swizzle_mode[2] | pad[10]) is documented on the cute_nvgpu MMA atoms page.

nvgpu.warpgroup.mma

Rewriter expands to the canonical four-op WGMMA sequence: nvvm.wgmma.fence.aligned, one nvvm.wgmma.mma_async per accumulator tile, nvvm.wgmma.commit.group.sync.aligned, then nvvm.wgmma.wait.group.sync.aligned waitGroup. See WGMMA Emission Protocol — The Four-Op Sequence for the timing rules and accumulator lifetime. It validates GMMA layout up front with the canonical "Not a canonical GMMA_MN Layout" wording lifted from CUTLASS's gmma.hpp.

nvgpu.warpgroup.mma.store

Rewriter decomposes the accumulator into per-thread llvm.store operations against the destination memref. No nvvm.* op is emitted.

nvgpu.warpgroup.mma.init.accumulator

Rewriter emits llvm.mlir.zero (or llvm.mlir.undef followed by per-field zero stores) producing the accumulator aggregate.

nvgpu.mma.sync

Rewriter emits nvvm.wmma.mma.sync.aligned (Ampere/Ada) or routes through nvvm.wgmma.mma_async (Hopper) based on the active SM.

nvgpu.mma.sp.sync

Rewriter emits llvm.inline_asm with the mma.sp.sync.aligned.m... template; upstream NVVM exposes no sparse-MMA op in the snapshot tileiras tracks.

nvgpu.ldmatrix

Rewriter emits nvvm.ldmatrix.sync.aligned returning an llvm.struct<(i32, i32, ...)>, then a pack-struct-into-vector repack to match the result type.

nvgpu.stmatrix (absent in this binary)

There is no nvgpu.stmatrix mnemonic in this tileiras build's string table. The stmatrix store path is reached from the upstream MLIR vector / nvvm populators directly into nvvm.stmatrix. The operand layout below mirrors the upstream wrapper.

Upstream rewriter shape: emit nvvm.stmatrix.sync.aligned on sm_90+ targets, or llvm.inline_asm with the matching stmatrix... template otherwise.

gpu.warp_execute_on_lane_0 (routed through upstream gpu dialect)

There is no nvgpu.warp.execute_on_lane_0 mnemonic; the corresponding op is the upstream gpu dialect's gpu.warp_execute_on_lane_0, which convert-nvgpu-to-nvvm rewrites in passing.

Rewriter emits a region predicate against lane == 0, runs the body, and broadcasts each result with nvvm.shfl.sync (idx, 0, 0xffffffff).

Packed conversion and arithmetic helpers

Each packed op carries a rnd enum (rn, rz, rm, rp) and, where applicable, a packed_kind enum that selects between MXFP / NVFP packing modes.

Lowering-Target Table

What each rewriter emits. The middle column gives the concrete NVVM op (or the expanded form when the pattern bypasses NVVM on purpose); the right column is what the NVPTX backend ultimately prints, not anything nvgpu itself emits.

The sparse-MMA path reaches PTX through llvm.inline_asm because the snapshot's upstream NVVM does not yet expose a sparse-MMA op. The template, constraint string, and result type live in the pattern body and drop verbatim into the LLVM module — see Inline-PTX templates and constraint strings on the NVVM overview for the constraint-string form.

Per-Arch Availability

convert-nvgpu-to-nvvm runs unconditionally on every target — the gates live inside the patterns and in NVVM verification, not in pass scheduling. The first column gives the lowest SM that accepts each pattern, the second the form it emits at that floor, the third the lowest PTX ISA version that defines the resulting instruction.

sm_90a is the architecture-qualified variant wgmma and TMA require; plain sm_90 rejects them at NVVM verification. The dialect has no sm_100 op of its own — the Blackwell tcgen05 surface lives entirely in nvvm, accessed through cute_nvgpu atoms that lower past nvgpu. See Per-SM Emission Templates for the per-tier capability matrix.

Pattern-Set Construction

populateNVGPUToNVVMConversionPatterns is a flat populator: one OpConversionPattern per nvgpu.* op, each registered with benefit = 1. The patterns are stateless — they read operands and attributes through their OpAdaptor, emit a fixed sequence of nvvm.* (or llvm.* / memref.*) ops, and replace the root.

Tileiras consumes this populator unchanged from upstream MLIR. Reimplementations should match the same one-pattern-per-op shape; the rewriter's branch on source memory space is the only piece of policy the patterns carry. See Lowering: nvgpu / gpu to NVVM — Pattern Shapes for the rewrite primitives the patterns share.

Lowering Contract

The conversion never reinfers layout intent. By the time IR reaches nvgpu, descriptor shape, memory space, vector shape, MMA tile shape, sparse metadata, and barrier identity already live in operands and attributes. Pattern bodies stay small as a result.

The mbarrier family branches on memory space and emits one nvvm.mbarrier.*[.shared] intrinsic per op. See mbarrier State Machine for the slot transitions and NVVM mbarrier Ops for the per-op intrinsic mapping. TMA load and store each emit a single nvvm.cp.async.bulk.tensor.* intrinsic, threading the variadic coordinates, multicast mask, and L2 cache hint through unchanged. The largest pattern is nvgpu.warpgroup.mma: it emits the four-stage Hopper WGMMA sequence — fence, async MMA, commit, wait — and validates GMMA layout up front with the canonical "Not a canonical GMMA_MN Layout" wording lifted from CUTLASS's gmma.hpp.

A handful of patterns emit no nvvm.* op at all. nvgpu.mbarrier.create emits a memref.global with "private" visibility plus a memref.get_global, allocating the __mbarrier slot in shared memory. nvgpu.tma.create.descriptor emits an llvm.alloca for a 128-byte CUtensorMap, fills it via llvm.getelementptr+llvm.store sequences, then calls the CUDA driver's cuTensorMapEncodeTiled. nvgpu.warpgroup.descriptor is a pure shl/or chain over the WGMMA descriptor bitfield. nvgpu.mma.sp.sync emits an llvm.inline_asm with the verbatim "mma.sp.sync.aligned.m..." PTX template; at the snapshot revision tileiras tracks, upstream NVVM has no sparse-MMA op yet, and inline-asm is the upstream design.

Verification Invariants

The interesting nvgpu verifier checks are semantic, not lexical. TMA ops demand valid descriptor types, compatible source or destination memrefs, supported tensor-map ranks, and a legal shared-memory layout. WGMMA demands rank-2 matrix fragments, compatible M/N/K, a supported tile shape, matching accumulator and result types, and legal transpose flags. MMA and sparse MMA add element-type checks, sparse-selector bounds, and a guard that tf32 only pairs with valid floating-point operands. Device async copy requires matching element types, unit-stride minor dimensions, supported transfer sizes, and correct alignment when L1 bypass is requested.

The boundary matters because NVVM conversion assumes the op is already legal for the selected target. Invalid shapes slipping through here resurface later as much less useful intrinsic-selection or backend diagnostics.

Reimplementation Checklist

A practical reimplementation needs the operation families above, typed descriptor and barrier values, shape-aware verifiers, and a deterministic conversion table to NVVM. Keep the layer transient. Independent scheduling, high-level layout algebra, and CUDA Tile semantics all belong above nvgpu. The dialect's job is to normalise hardware operations, verify their low-level shape contracts, and hand them to NVVM with as little policy as possible.

The minimum useful surface: tensor-map descriptor creation and async TMA load/store/reduce; shared-memory barrier groups and barrier tokens with expect_tx and try_wait.parity; WGMMA accumulator init / mma / store; the WGMMA descriptor packer; MMA ops with explicit shape attributes; sparse MMA metadata and selector validation; ldmatrix and stmatrix; SM80 cp.async device-async copy; packed conversion and arithmetic helpers; a complete nvgpu-to-nvvm conversion table; and target-aware verification before conversion.