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