wmma GPU Virtual ISA NVIDIA

Warp Matrix Multiply-Accumulate Warp Level Matrix Multiply-Accumulate Instructions

wmma.load.a.sync.aligned.layout.shape.type r, [p];

Higher-level warp matrix-multiply-accumulate built from explicit load/mma/store steps.

Encoding

PTX is a virtual instruction set. It has no single, stable native binary encoding - the compiler lowers this instruction to different native machine code depending on the selected NVIDIA target architecture (compute capability). This page intentionally shows no bit-diagram; see the target/version requirements below for what governs how this instruction compiles.
PTX ISA Version Introduced PTX ISA 6.0
Minimum Target sm_70

Syntax Forms

One mnemonic covers many type / state-space / scope / modifier combinations - each row below is an independently valid form.

Syntax Data Types State Space(s) Modifiers Min. Target Description
wmma.load.a.sync.aligned.layout.shape.type r, [p]; sm_70 Load a matrix fragment cooperatively across the warp.
wmma.mma.sync.aligned.layout.shape.dtype.ctype d, a, b, c; sm_70 Perform the accumulate step on already-loaded fragments.
wmma.store.d.sync.aligned.layout.shape.type [p], r; sm_70 Store an accumulator fragment cooperatively across the warp.

Operands

  • r/d
    Fragment register(s)
  • a
    Matrix A fragment
  • b
    Matrix B fragment
  • c
    Accumulator fragment (input)
  • p
    Memory address for load/store forms

At a Glance

Data Types -

Related AMDGPU Concepts

Related

More in Warp Level Matrix Multiply-Accumulate Instructions

Reference

NVIDIA PTX ISA

Description

Perform a warp-level matrix multiply-and-accumulate computation D = A * B + C using matrices A, B and C loaded in registers a, b and c respectively, and store the result matrix in register d. The register arguments a, b, c and d hold unspecified fragments of the corresponding matrices as described in Matrix Fragments for WMMA The qualifiers.dtype,.atype,.btype and.ctype indicate the data-type of the elements in the matrices D, A, B and C respectively. For wmma.mma without explicit.atype and.btype:.atype and.btype are implicitly set to.f16. For integer wmma,.ctype and.dtype must be specified as.s32. (see the official PTX ISA docs for the full description)

Semantics

Together, the load/mma/store triplet compute D = A * B + C for a fixed tile shape, cooperatively across the warp.

Examples

// Load elements from f16 row-major matrix B
.reg .b32 x<8>;

wmma.load.b.sync.aligned.m16n16k16.row.f16 {x0,x1,x2,x3,x4,x5,x,x7}, [ptr];
// Now use {x0, ..., x7} for the actual wmma.mma
// (truncated - see the official PTX ISA docs for the full example)

.global .align 32 .f16 A[256], B[256];
.global .align 32 .f32 C[256], D[256];
.reg .b32 a<8> b<8> c<8> d<8>;

wmma.load.a.sync.aligned.m16n16k16.global.row.f16
        {a0, a1, a2, a3, a4, a5, a6, a7}, [A];
// (truncated - see the official PTX ISA docs for the full example)

// Storing f32 elements computed by a wmma.mma
.reg .b32 x<8>;

wmma.mma.sync.m16n16k16.row.col.f32.f32
              {d0, d1, d2, d3, d4, d5, d6, d7}, ...;
wmma.store.d.sync.m16n16k16.row.f32
// (truncated - see the official PTX ISA docs for the full example)

Reproduced from NVIDIA's official PTX ISA documentation for technical accuracy.

Sources