RDNA4 WMMA transpose load builtins#

RDNA4 (gfx1200, gfx1201) provides hardware instructions that load a matrix sub-tile into per-lane VGPRs while transposing row-major and column-major storage order in the same operation. When a matrix’s storage order in memory doesn’t match the fragment layout the RDNA4 dense WMMA builtins expect, these builtins remove the need for a scalar transpose loop.

Note

Like the WMMA builtins they feed, these load-transpose instructions are supported only for wave32 on RDNA4. RDNA4 also has wavefrontsize64 transpose-load variants, but since RDNA4 WMMA itself is wave32-only, those aren’t tied to feeding WMMA and aren’t covered on this page.

Architecture availability#

The builtins on this page target RDNA4 GPUs. To automatically enable them, pass the Low Level Virtual Machine (LLVM) target architecture flag at compile time:

amdclang++ --offload-arch=gfx1200 ...
amdclang++ --offload-arch=gfx1201 ...

Naming convention#

All transpose-load builtins on this page follow the pattern:

__builtin_amdgcn_global_load_tr_b<M>_v<K><type>

Unlike CDNA5’s transpose-load builtins, RDNA4 names carry no bit-width digit (no tr16/tr8) – RDNA4 has only one transpose-load instruction family per load size, so no extra digit is needed to disambiguate it.

M

Bits loaded per lane (128 or 64).

K

Number of elements in the return vector.

type

Element type of the return vector: i32, i16, f16, or bf16.

Register types used in this reference#

The signatures below use the following type aliases, which you can declare with Clang vector attributes in any HIP translation unit:

using v2int   = int      [[clang::ext_vector_type(2)]];
using v8short = short    [[clang::ext_vector_type(8)]];
using v8half  = _Float16 [[clang::ext_vector_type(8)]];
using v8bf16  = __bf16   [[clang::ext_vector_type(8)]];

Note

The f16 builtin requires the __fp16 type (GCC-style half-precision), not _Float16 (ISO C half-precision). Cast pointers and vector types to __fp16-based variants at the call site only – the standard WMMA builtins accept _Float16 as usual. The worked example below does exactly this with its own half8 alias.

Hardware constraints#

  • Sub-tile per call, not a full 16x16 tile. global_load_tr_b128 covers an 8-row x 32-column sub-tile per call (see Per-lane fragment layout below) – two calls, the second offset by 8 rows, are needed to fill a full 16-row K-strip. This differs from CDNA5’s tr16_b128/tr8_b64, which each load a complete \(16 \times 16\) tile in a single call.

  • Wave32 only.

  • All builtins are const and have no side effects.

Per-lane fragment layout#

8x32 grid of lane numbers and return-vector element indices for GLOBAL_LOAD_TR_B128 (v8f16), colored by even/odd K row.

Per-lane assignment for global_load_tr_b128_v8f16. Each cell shows the lane index (bold) and the return-vector element index \(e\) (small) that produced it.#

Given a call whose base K-row is k_tile and base N-column is block_n, lane tx (0–31) reads from:

ptr(tx) = &input[(k_tile + tx/4) * cols + block_n + (tx%4)*8]

and the returned vector’s element e (0–7) lands at:

\[\begin{split}\text{row} &= k\_tile + \lfloor tx/8 \rfloor \cdot 2 + \lfloor e/4 \rfloor \\ \text{col} &= block\_n + (e \bmod 4) \cdot 8 + (tx \bmod 8)\end{split}\]

A second call with ptr shifted by 8 rows (k_tile + 8) fills the upper half of a 16-row K-strip; together the two calls cover \(16 \times 32\) elements – exactly the B operand a wmma_f32_16x16x16_f16_w32_gfx12 K-step needs from two adjacent wave32 blocks. global_load_tr_b128_v8i16 and _v8bf16 use the same 128-bit data path with a different element-type interpretation, so they follow this same mapping. global_load_tr_b64_v2i32 uses a different (64-bit) data path; no per-lane mapping is documented for it below.

Builtin reference#

__builtin_amdgcn_global_load_tr_b64_v2i32#

v2int __builtin_amdgcn_global_load_tr_b64_v2i32(v2int* ptr);

Loads and transposes 64 bits per lane of 8-bit data (INT8, FP8, or BF8), returning two packed 32-bit words rather than eight discrete 8-bit elements. This builtin uses a different data path from global_load_tr_b128; see Per-lane fragment layout.

Parameter

Type

Description

ptr

v2int*

Pointer to global address space.

Returns v2int – two 32-bit words holding the lane’s 8-bit fragment after transposition.

__builtin_amdgcn_global_load_tr_b128_v8i16#

v8short __builtin_amdgcn_global_load_tr_b128_v8i16(v8short* ptr);

Loads and transposes 128 bits per lane of INT16 data. Expected to share v8f16’s permutation – see Per-lane fragment layout.

Parameter

Type

Description

ptr

v8short*

Pointer to global address space.

Returns v8short – eight INT16 elements holding the lane’s fragment after transposition.

__builtin_amdgcn_global_load_tr_b128_v8f16#

v8half __builtin_amdgcn_global_load_tr_b128_v8f16(v8half* ptr);

Loads and transposes 128 bits per lane of FP16 data. See the __fp16 note above – cast ptr and the return value to a __fp16-based v8half at the call site. Its per-lane mapping is detailed in Per-lane fragment layout.

Parameter

Type

Description

ptr

v8half*

Pointer to global address space.

Returns v8half – eight FP16 elements holding the lane’s fragment after transposition.

__builtin_amdgcn_global_load_tr_b128_v8bf16#

v8bf16 __builtin_amdgcn_global_load_tr_b128_v8bf16(v8bf16* ptr);

Loads and transposes 128 bits per lane of BF16 data. Expected to share v8f16’s permutation – see Per-lane fragment layout.

Parameter

Type

Description

ptr

v8bf16*

Pointer to global address space.

Returns v8bf16 – eight BF16 elements holding the lane’s fragment after transposition.

Using a transpose load as a compute policy#

Using MFMA builtins as a compute policy explains the ComputePolicy/TilePolicy pattern used to separate tile staging from the compute loop. This section extends the RDNA4 WMMA FP16 example from Using WMMA builtins as a compute policy with a transpose-load tile policy.

The complete source file is available for download:

WmmaRdna4F16TrPolicy replaces the scalar B transpose with a hardware- accelerated GLOBAL_LOAD_TR_B128 instruction (__builtin_amdgcn_global_load_tr_b128_v8f16). This intrinsic loads 128 bits per lane from global memory and transposes the data into LDS in a single operation, avoiding the per-element scalar transpose loop. The per-lane addressing follows exactly the formula in Per-lane fragment layout.

SingleBufferTilePolicyF16TrB fills B’s shared-memory tile using this instruction instead of a scalar loop:

SingleBufferTilePolicyF16TrB (matrix_multiply_rdna4_wmma.hip)#
template<int BlockTileM_, int BlockTileN_, int KTileSize_>
struct SingleBufferTilePolicyF16TrB
{
    static constexpr int num_buffers  = 1;
    static constexpr int block_tile_m = BlockTileM_;
    static constexpr int block_tile_n = BlockTileN_;
    static constexpr int k_tile_size  = KTileSize_;

    struct SharedStorage
    {
        _Float16 tile_a[BlockTileM_][KTileSize_];
        // B stored transposed: tile_b_T[n][k], same layout as baseline.
        // Filled by global_load_tr hardware scatter from row-major global B.
        _Float16 tile_b_T[BlockTileN_][KTileSize_];
    };

    __device__ static void prologue(SharedStorage& /*smem*/,
                                    const _Float16* /*A*/, const _Float16* /*B*/,
                                    int /*block_row_start*/, int /*block_col_start*/,
                                    int /*m*/, int /*n*/, int /*k*/,
                                    int /*tid*/, int /*num_threads*/) {}

    __device__ static void prefetch(SharedStorage&  smem,
                                    const _Float16* A,
                                    const _Float16* B_row,
                                    int             t_next,
                                    int block_row_start, int block_col_start,
                                    int m, int n, int k,
                                    int tid, int num_threads)
    {
        const int t = t_next;

        cooperative_load_tile_a_f16(
            &smem.tile_a[0][0], A, block_row_start,
            block_tile_m, t, k_tile_size, m, k, tid, num_threads);

#if defined(__gfx1200__) || defined(__gfx1201__)
        using half8 = __attribute__((ext_vector_type(8))) __fp16;

        constexpr int wave_size = 32;
        const int lane_id       = tid % wave_size;
        const int k_tile        = t * k_tile_size;
        const int col_offset    = lane_id % 8;
        const int row_pair_base = (lane_id / 8) * 2;

        // Call 1: lower K-half (k_local = 0..7).
        {
            const __fp16* ptr_b =
                reinterpret_cast<const __fp16*>(
                    &B_row[(k_tile + lane_id / 4) * n
                           + block_col_start + (lane_id % 4) * 8]);
            const half8 tr_lo =
                __builtin_amdgcn_global_load_tr_b128_v8f16(
                    (__attribute__((address_space(1))) half8*)ptr_b);
            #pragma unroll
            for(int e = 0; e < 8; ++e)
            {
                const int n_local = (e % 4) * 8 + col_offset;
                const int k_local = row_pair_base + e / 4;
                smem.tile_b_T[n_local][k_local] = tr_lo[e];
            }
        }

        // Call 2: upper K-half (k_local = 8..15).
        {
            const __fp16* ptr_b =
                reinterpret_cast<const __fp16*>(
                    &B_row[(k_tile + (k_tile_size / 2) + lane_id / 4) * n
                           + block_col_start + (lane_id % 4) * 8]);
            const half8 tr_hi =
                __builtin_amdgcn_global_load_tr_b128_v8f16(
                    (__attribute__((address_space(1))) half8*)ptr_b);
            #pragma unroll
            for(int e = 0; e < 8; ++e)
            {
                const int n_local = (e % 4) * 8 + col_offset;
                const int k_local = (k_tile_size / 2) + row_pair_base + e / 4;
                smem.tile_b_T[n_local][k_local] = tr_hi[e];
            }
        }
#endif
    }

    __device__ static void acquire(int /*t*/) { __syncthreads(); }
    __device__ static void release(int /*t*/) { __syncthreads(); }
    __device__ static int  buf_idx(int /*t*/) { return 0; }

    __device__ static _Float16* get_tile_a(SharedStorage& smem, int /*buf*/)
    {
        return &smem.tile_a[0][0];
    }
    __device__ static _Float16* get_tile_b_T(SharedStorage& smem, int /*buf*/)
    {
        return &smem.tile_b_T[0][0];
    }
};

The compute path (load_a, load_b, mma, store_c) is identical to the baseline WmmaRdna4F16Policy – only the tile policy’s prefetch() stage changes:

WmmaRdna4F16TrPolicy (matrix_multiply_rdna4_wmma.hip)#
struct WmmaRdna4F16TrPolicy
{
    // -- ComputePolicy constants ----------------------------------------------
    static constexpr int thread_tile_m   = 16;
    static constexpr int thread_tile_n   = 16;
    static constexpr int frag_size_m     = 8;
    static constexpr int frag_size_n     = 8;
    // RDNA4: wave32, 32 lanes per wavefront.
    static constexpr int effective_lanes = 32;
    // k_step=16: one wmma_f32_16x16x16_f16_w32 call consumes 16 K-positions,
    // filled by two GLOBAL_LOAD_TR_B128 calls (one per 8-K half).
    static constexpr int k_step          = 16;

    using elem_a = _Float16;
    using elem_b = _Float16;

    // Accumulator and zero() are identical to WmmaRdna4F16Policy.
    using wmma_f8 = float [[clang::ext_vector_type(8)]];

    struct Accumulator
    {
        wmma_f8 regs;
    };

    __device__ static void zero(Accumulator& acc)
    {
        acc.regs = wmma_f8{};
    }

    // thread_tile_offset is identical to WmmaRdna4F16Policy.
    __device__ static void thread_tile_offset(int  tid,
                                              int  /*lane_id*/,
                                              int* thread_row,
                                              int* thread_col)
    {
        constexpr int waves_n = 2;
        const int wid   = tid / 32; // wave32
        const int wid_x = wid % waves_n;
        const int wid_y = wid / waves_n;
        *thread_row = wid_y * 16;
        *thread_col = wid_x * 16;
    }

    // -- load_a ---------------------------------------------------------------
    // Scalar load, identical to WmmaRdna4F16Policy::load_a.
    __device__ static void load_a(const _Float16* tile_a_ptr,
                                  int             tile_a_row,
                                  int             ki,
                                  int             k_tile_size,
                                  int             lane_id,
                                  elem_a        (&frag)[8])
    {
        const int row   = tile_a_row + (lane_id % 16);
        const int k_off = ki + 8 * (lane_id / 16);
        #pragma unroll
        for(int e = 0; e < 8; ++e)
            frag[e] = tile_a_ptr[row * k_tile_size + k_off + e];
    }

    // -- load_b ---------------------------------------------------------------
    // Scalar read from tile_b_T[n][k] in LDS - identical to baseline.
    // The global_load_tr hardware transpose already happened in
    // SingleBufferTilePolicyF16TrB::prefetch().
    __device__ static void load_b(const _Float16* tile_b_T_ptr,
                                  int             tile_b_col,
                                  int             ki,
                                  int             k_tile_size,
                                  int             lane_id,
                                  elem_b        (&frag)[8])
    {
        const int col   = tile_b_col + (lane_id % 16);
        const int k_off = ki + 8 * (lane_id / 16);
        #pragma unroll
        for(int e = 0; e < 8; ++e)
            frag[e] = tile_b_T_ptr[col * k_tile_size + k_off + e];
    }

    // mma and store_c are identical to WmmaRdna4F16Policy.
    __device__ static void mma(Accumulator&       acc,
                               const elem_a     (&a_frag)[8],
                               const elem_b     (&b_frag)[8])
    {
#if defined(__gfx1200__) || defined(__gfx1201__)
        using half8 = __attribute__((ext_vector_type(8))) __fp16;
        const half8 a_vec = {a_frag[0], a_frag[1], a_frag[2], a_frag[3],
                             a_frag[4], a_frag[5], a_frag[6], a_frag[7]};
        const half8 b_vec = {b_frag[0], b_frag[1], b_frag[2], b_frag[3],
                             b_frag[4], b_frag[5], b_frag[6], b_frag[7]};
        acc.regs = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(
            a_vec, b_vec, acc.regs);
#endif
    }

    __device__ static void store_c(const Accumulator& acc,
                                   float*             C,
                                   int                out_row_base,
                                   int                out_col_base,
                                   int                m,
                                   int                n,
                                   int                lane_id)
    {
        const int laneWrapped = lane_id % 16;
        const int laneGroup   = lane_id / 16;

        #pragma unroll
        for(int G = 0; G < 8; ++G)
        {
            const int i = laneGroup * 8 + G;
            const int j = laneWrapped;
            const int r = out_row_base + i;
            const int c = out_col_base + j;
            if(r < m && c < n)
                C[r * n + c] = acc.regs[G];
        }
    }
};

Instantiating the kernel

With either policy, plug it into the generic kernel alongside any TilePolicy whose block_tile_m and block_tile_n are multiples of 16 and whose k_tile_size is a multiple of k_step = 16. See Using WMMA builtins as a compute policy for the shared policy aliases, launch configuration, and kernel launch code – the same matrix_multiply_rdna4_wmma.hip file wires up both WmmaTilePolicy (baseline) and WmmaTilePolicyTrB (this page) and runs both back to back for comparison.

Compiling and running#

amdclang++ -O3 --offload-arch=gfx1200 \
    matrix_multiply_rdna4_wmma.hip -o mm_rdna4_wmma
./mm_rdna4_wmma

amdclang++ -O3 --offload-arch=gfx1201 \
    matrix_multiply_rdna4_wmma.hip -o mm_rdna4_wmma
./mm_rdna4_wmma

Note

WmmaRdna4F16TrPolicy requires an RDNA4 GPU (gfx1200 or gfx1201). The #if defined(__gfx1200__) || defined(__gfx1201__) guard in the example file falls back to ScalarFMAPolicy on other targets, so the file compiles without modification.