Neko 1.99.9
A portable framework for high-order spectral element flow simulations
Loading...
Searching...
No Matches
mfma_kernel.h File Reference
#include <stdlib.h>
#include <string.h>
#include <hip/hip_runtime.h>
#include <device/device_config.h>
#include <device/hip/check.h>
#include "elem_block.h"
Include dependency graph for mfma_kernel.h:
This graph shows which files directly or indirectly include this file:

Go to the source code of this file.

Macros

#define NEKO_MFMA_NWF_CANDIDATES   5
 
#define NEKO_MFMA_TILE_CANDIDATES   2
 
#define NEKO_MFMA_CANDIDATES    (NEKO_MFMA_NWF_CANDIDATES * NEKO_MFMA_TILE_CANDIDATES)
 
#define NEKO_MFMA_NWF(C)   (1 << ((C) % NEKO_MFMA_NWF_CANDIDATES))
 
#define NEKO_MFMA_TILE(C)   ((C) / NEKO_MFMA_NWF_CANDIDATES)
 
#define NEKO_MFMA_NTHRDS(C)   dim3(64, NEKO_MFMA_NWF(C), 1)
 
#define NEKO_MFMA_PAD_N(LX, SZ)    ((((LX) % 4 == 0) && !((LX) == 4 && (SZ) == 4)) ? 1 : 0)
 
#define NEKO_MFMA_SJ_N(LX, SZ)   ((LX) + NEKO_MFMA_PAD_N(LX, SZ))
 
#define NEKO_MFMA_SK_N(LX, SZ)   (NEKO_MFMA_SJ_N(LX, SZ) * (LX))
 
#define NEKO_MFMA_CUBE_N(LX, SZ)   (NEKO_MFMA_SK_N(LX, SZ) * (LX))
 
#define NEKO_MFMA_SD_N(LX, SZ)   NEKO_MFMA_SJ_N(LX, SZ)
 
#define NEKO_MFMA_DMAT_N(LX, SZ)   (NEKO_MFMA_SD_N(LX, SZ) * (LX))
 
#define NEKO_MFMA_NGROUPS(LX)   (((LX) * (LX) + 15) / 16)
 
#define NEKO_MFMA_LDS_N(EB, LX, SZ)    ((3 * NEKO_MFMA_DMAT_N(LX, SZ) + 4 * (EB) * NEKO_MFMA_CUBE_N(LX, SZ)) * (SZ))
 
#define NEKO_MFMA_LDS_FITS(EB, LX, SZ)    (NEKO_MFMA_LDS_N(EB, LX, SZ) <= NEKO_EB_MAX_LDS)
 
#define NEKO_MFMA_EB_N(NWF, LX, SZ)
 
#define NEKO_MFMA_EB(LX, C)    NEKO_MFMA_EB_N(NEKO_MFMA_NWF(C), LX, sizeof(real))
 
#define NEKO_MFMA_WPE(LX, C)   (NEKO_MFMA_NWF(C) / NEKO_MFMA_EB(LX, C))
 
#define NEKO_MFMA_POW2_FLOOR(X)    ((X) >= 16 ? 16 : (X) >= 8 ? 8 : (X) >= 4 ? 4 : (X) >= 2 ? 2 : 1)
 
#define NEKO_MFMA_NEL_N(EB, LX, SZ, TILE)
 
#define NEKO_MFMA_NEL(LX, C)    NEKO_MFMA_NEL_N(NEKO_MFMA_EB(LX, C), LX, sizeof(real), NEKO_MFMA_TILE(C))
 
#define NEKO_MFMA_SPT_N(WPE, LX, SZ)    ((NEKO_MFMA_CUBE_N(LX, SZ) + (WPE) * 64 - 1) / ((WPE) * 64))
 
#define NEKO_MFMA_SPT(LX, C)    NEKO_MFMA_SPT_N(NEKO_MFMA_WPE(LX, C), LX, sizeof(real))
 
#define NEKO_MFMA_VECTOR_GREG_VGPRS   64
 
#define NEKO_MFMA_VECTOR_GREG_N(SPT, SZ)    ((7 * (SPT) * ((SZ) / 4)) <= NEKO_MFMA_VECTOR_GREG_VGPRS)
 
#define NEKO_MFMA_VECTOR_GREG(LX, C)    NEKO_MFMA_VECTOR_GREG_N(NEKO_MFMA_SPT(LX, C), sizeof(real))
 
#define NEKO_MFMA_NBLCKS(NELV, LX, C)    dim3(((NELV) + NEKO_MFMA_EB(LX, C) - 1) / NEKO_MFMA_EB(LX, C), 1, 1)
 
#define NEKO_TUNE_LOG_MFMA(LX, T3)
 
#define NEKO_TUNE_LOG_MFMA_VEC(LX, T3)
 

Functions

static __global__ void hip_mfma_arch_probe (int *flag)
 
static bool hip_have_mfma ()
 
template<const int LX>
static bool mfma_lx_supported ()
 
static bool mfma_tile_offered (const int c)
 
static const char * mfma_tile_name (const int c)
 
static int neko_mfma_sweep ()
 
static int neko_mfma_pin ()
 
static int neko_mfma_candidates ()
 

Macro Definition Documentation

◆ NEKO_MFMA_CANDIDATES

Definition at line 278 of file mfma_kernel.h.

◆ NEKO_MFMA_CUBE_N

#define NEKO_MFMA_CUBE_N (   LX,
  SZ 
)    (NEKO_MFMA_SK_N(LX, SZ) * (LX))

Definition at line 347 of file mfma_kernel.h.

◆ NEKO_MFMA_DMAT_N

#define NEKO_MFMA_DMAT_N (   LX,
  SZ 
)    (NEKO_MFMA_SD_N(LX, SZ) * (LX))

Definition at line 368 of file mfma_kernel.h.

◆ NEKO_MFMA_EB

#define NEKO_MFMA_EB (   LX,
  C 
)     NEKO_MFMA_EB_N(NEKO_MFMA_NWF(C), LX, sizeof(real))

Definition at line 472 of file mfma_kernel.h.

◆ NEKO_MFMA_EB_N

#define NEKO_MFMA_EB_N (   NWF,
  LX,
  SZ 
)
Value:
((NWF) / NEKO_MFMA_NGROUPS(LX) >= 16 && NEKO_MFMA_LDS_FITS(16, LX, SZ) ? 16 :\
(NWF) / NEKO_MFMA_NGROUPS(LX) >= 8 && NEKO_MFMA_LDS_FITS(8, LX, SZ) ? 8 : \
(NWF) / NEKO_MFMA_NGROUPS(LX) >= 4 && NEKO_MFMA_LDS_FITS(4, LX, SZ) ? 4 : \
(NWF) / NEKO_MFMA_NGROUPS(LX) >= 2 && NEKO_MFMA_LDS_FITS(2, LX, SZ) ? 2 : 1)
__global__ void ale_add_kinematics_kernel(const int n, T *__restrict__ wx, T *__restrict__ wy, T *__restrict__ wz, const T *__restrict__ x_ref, const T *__restrict__ y_ref, const T *__restrict__ z_ref, const T *__restrict__ phi, const T *__restrict__ x, const T *__restrict__ y, const T *__restrict__ z, const kinematics_params_t kin_params)
#define NEKO_MFMA_LDS_FITS(EB, LX, SZ)
#define NEKO_MFMA_NGROUPS(LX)

Definition at line 467 of file mfma_kernel.h.

◆ NEKO_MFMA_LDS_FITS

#define NEKO_MFMA_LDS_FITS (   EB,
  LX,
  SZ 
)     (NEKO_MFMA_LDS_N(EB, LX, SZ) <= NEKO_EB_MAX_LDS)

Definition at line 458 of file mfma_kernel.h.

◆ NEKO_MFMA_LDS_N

#define NEKO_MFMA_LDS_N (   EB,
  LX,
  SZ 
)     ((3 * NEKO_MFMA_DMAT_N(LX, SZ) + 4 * (EB) * NEKO_MFMA_CUBE_N(LX, SZ)) * (SZ))

Definition at line 456 of file mfma_kernel.h.

◆ NEKO_MFMA_NBLCKS

#define NEKO_MFMA_NBLCKS (   NELV,
  LX,
  C 
)     dim3(((NELV) + NEKO_MFMA_EB(LX, C) - 1) / NEKO_MFMA_EB(LX, C), 1, 1)

Definition at line 554 of file mfma_kernel.h.

◆ NEKO_MFMA_NEL

#define NEKO_MFMA_NEL (   LX,
  C 
)     NEKO_MFMA_NEL_N(NEKO_MFMA_EB(LX, C), LX, sizeof(real), NEKO_MFMA_TILE(C))

Definition at line 503 of file mfma_kernel.h.

◆ NEKO_MFMA_NEL_N

#define NEKO_MFMA_NEL_N (   EB,
  LX,
  SZ,
  TILE 
)
Value:
(((SZ) == 4 && (TILE) == 0) ? \
((EB) < NEKO_MFMA_POW2_FLOOR(64 / ((LX) * (LX))) ? \
(EB) : NEKO_MFMA_POW2_FLOOR(64 / ((LX) * (LX)))) : 1)
#define NEKO_MFMA_POW2_FLOOR(X)

Definition at line 498 of file mfma_kernel.h.

◆ NEKO_MFMA_NGROUPS

#define NEKO_MFMA_NGROUPS (   LX)    (((LX) * (LX) + 15) / 16)

Definition at line 448 of file mfma_kernel.h.

◆ NEKO_MFMA_NTHRDS

#define NEKO_MFMA_NTHRDS (   C)    dim3(64, NEKO_MFMA_NWF(C), 1)

Definition at line 282 of file mfma_kernel.h.

◆ NEKO_MFMA_NWF

#define NEKO_MFMA_NWF (   C)    (1 << ((C) % NEKO_MFMA_NWF_CANDIDATES))

Definition at line 280 of file mfma_kernel.h.

◆ NEKO_MFMA_NWF_CANDIDATES

#define NEKO_MFMA_NWF_CANDIDATES   5

Definition at line 276 of file mfma_kernel.h.

◆ NEKO_MFMA_PAD_N

#define NEKO_MFMA_PAD_N (   LX,
  SZ 
)     ((((LX) % 4 == 0) && !((LX) == 4 && (SZ) == 4)) ? 1 : 0)

Definition at line 341 of file mfma_kernel.h.

◆ NEKO_MFMA_POW2_FLOOR

#define NEKO_MFMA_POW2_FLOOR (   X)     ((X) >= 16 ? 16 : (X) >= 8 ? 8 : (X) >= 4 ? 4 : (X) >= 2 ? 2 : 1)

Definition at line 493 of file mfma_kernel.h.

◆ NEKO_MFMA_SD_N

#define NEKO_MFMA_SD_N (   LX,
  SZ 
)    NEKO_MFMA_SJ_N(LX, SZ)

Definition at line 367 of file mfma_kernel.h.

◆ NEKO_MFMA_SJ_N

#define NEKO_MFMA_SJ_N (   LX,
  SZ 
)    ((LX) + NEKO_MFMA_PAD_N(LX, SZ))

Definition at line 345 of file mfma_kernel.h.

◆ NEKO_MFMA_SK_N

#define NEKO_MFMA_SK_N (   LX,
  SZ 
)    (NEKO_MFMA_SJ_N(LX, SZ) * (LX))

Definition at line 346 of file mfma_kernel.h.

◆ NEKO_MFMA_SPT

#define NEKO_MFMA_SPT (   LX,
  C 
)     NEKO_MFMA_SPT_N(NEKO_MFMA_WPE(LX, C), LX, sizeof(real))

Definition at line 520 of file mfma_kernel.h.

◆ NEKO_MFMA_SPT_N

#define NEKO_MFMA_SPT_N (   WPE,
  LX,
  SZ 
)     ((NEKO_MFMA_CUBE_N(LX, SZ) + (WPE) * 64 - 1) / ((WPE) * 64))

Definition at line 518 of file mfma_kernel.h.

◆ NEKO_MFMA_TILE

#define NEKO_MFMA_TILE (   C)    ((C) / NEKO_MFMA_NWF_CANDIDATES)

Definition at line 281 of file mfma_kernel.h.

◆ NEKO_MFMA_TILE_CANDIDATES

#define NEKO_MFMA_TILE_CANDIDATES   2

Definition at line 277 of file mfma_kernel.h.

◆ NEKO_MFMA_VECTOR_GREG

#define NEKO_MFMA_VECTOR_GREG (   LX,
  C 
)     NEKO_MFMA_VECTOR_GREG_N(NEKO_MFMA_SPT(LX, C), sizeof(real))

Definition at line 551 of file mfma_kernel.h.

◆ NEKO_MFMA_VECTOR_GREG_N

#define NEKO_MFMA_VECTOR_GREG_N (   SPT,
  SZ 
)     ((7 * (SPT) * ((SZ) / 4)) <= NEKO_MFMA_VECTOR_GREG_VGPRS)

Definition at line 549 of file mfma_kernel.h.

◆ NEKO_MFMA_VECTOR_GREG_VGPRS

#define NEKO_MFMA_VECTOR_GREG_VGPRS   64

Definition at line 547 of file mfma_kernel.h.

◆ NEKO_MFMA_WPE

#define NEKO_MFMA_WPE (   LX,
  C 
)    (NEKO_MFMA_NWF(C) / NEKO_MFMA_EB(LX, C))

Definition at line 475 of file mfma_kernel.h.

◆ NEKO_TUNE_LOG_MFMA

#define NEKO_TUNE_LOG_MFMA (   LX,
  T3 
)
Value:
do { \
for (int c = 0; c < NEKO_MFMA_CANDIDATES; c++) { \
if ((T3)[c] >= NEKO_TUNE_INIT) { continue; } \
sprintf(neko_log_buf, "MFMA %s %2dwf %-2de x%-2d %-5s: %9.2f us/call", \
(int) NEKO_MFMA_NEL(LX, c), \
NEKO_MFMA_PAD_N(LX, sizeof(real)) ? "pad" : "plain", \
} \
} while (0)
#define NEKO_TUNE_US(T, ITERS)
#define NEKO_TUNE_INIT
double real
static const char * mfma_tile_name(const int c)
#define NEKO_MFMA_CANDIDATES
#define NEKO_MFMA_NWF(C)
#define NEKO_MFMA_EB(LX, C)
#define NEKO_MFMA_NEL(LX, C)
#define NEKO_MFMA_PAD_N(LX, SZ)

Definition at line 633 of file mfma_kernel.h.

◆ NEKO_TUNE_LOG_MFMA_VEC

#define NEKO_TUNE_LOG_MFMA_VEC (   LX,
  T3 
)
Value:
do { \
for (int c = 0; c < NEKO_MFMA_CANDIDATES; c++) { \
if ((T3)[c] >= NEKO_TUNE_INIT) { continue; } \
"MFMA %s %2dwf %-2de x%-2d %-5s %-4s: %9.2f us/call", \
(int) NEKO_MFMA_NEL(LX, c), \
NEKO_MFMA_PAD_N(LX, sizeof(real)) ? "pad" : "plain", \
NEKO_MFMA_VECTOR_GREG(LX, c) ? "reg" : "glob", \
} \
} while (0)
#define NEKO_MFMA_VECTOR_GREG(LX, C)

Definition at line 655 of file mfma_kernel.h.

Function Documentation

◆ hip_have_mfma()

static bool hip_have_mfma ( )
inlinestatic

Returns true if the current device exposes the matrix cores used by the MFMA strategies (currently gfx90a / MI250X and gfx942 / MI300A / MI300X) and this build actually compiled the device code for them. The f64 4x4x4 / 16x16x4 and f32 16x16x4 instructions are all available on these arches, so a single check gates either precision. Result is cached after the first query.

Definition at line 144 of file mfma_kernel.h.

Here is the call graph for this function:
Here is the caller graph for this function:

◆ hip_mfma_arch_probe()

static __global__ void hip_mfma_arch_probe ( int *  flag)
static

Shared matrix-core (MFMA) primitives for the spectral-element tensor contractions, used by the Ax-helm autotuner strategy.

Shared by every operator that offers a matrix core strategy: the scalar, vector and compressible Ax-helm, dudxyz, opgrad, conv1 and cdtp. The contraction primitives below are the hardware-validated versions – see the layout note on mfma_contract_4x4 for why they must not be re-derived from host simulation alone – plus one, mfma_contract_4x4x1, whose layout is documented rather than measured and which the tuner verifies at runtime before it is allowed to win, see the note there.

Maps a reference derivative-matrix * field contraction onto the AMD matrix cores on CDNA2 (MI250X / gfx90a) and CDNA3 (MI300A / MI300X / gfx942). Each precision has two tiles, which every operator's autotuner measures against each other (TILE, see NEKO_MFMA_TILE below):

f64 TILE 0 v_mfma_f64_4x4x4f64 batched, four blocks, mfma_contract_4x4 f64 TILE 1 v_mfma_f64_16x16x4f64 mfma_contract f32 TILE 0 v_mfma_f32_4x4x1f32 batched, sixteen blocks, elements batched too, mfma_contract_4x4x1 f32 TILE 1 v_mfma_f32_16x16x4f32 mfma_contract

-DMFMA_F64_USE_16X16 and -DMFMA_F32_USE_16X16 collapse a precision onto its 16x16x4 tile. One thread block is launched as NWF 64-lane wavefronts (blockDim = (64, NWF, 1)) covering EB elements; each contraction is a D * U GEMM with M = LX, N = LX^2, K = LX, the partial tiles masked off, and the f32 4x4x1 tile batches NEL elements into one instruction stream at low order, see NEKO_MFMA_NEL_N.

Supported for 4 <= LX <= 12 in double precision and 2 <= LX <= 12 in single; the upper bound is set by the LDS needed to keep the cubes resident (operators stage four cubes plus three LX^2 derivative matrices, 63360 B of f64 at LX = 12 with the padded Ax-helm layout, 58752 B without it), the lower one by what a tile can still fill, see mfma_lx_supported().

16x16x4 register/lane layout (used by the f32 tile and the f64 16x16x4 fallback; the 4x4x4 layout is documented at mfma_contract_4x4).

Verification status, since this tile came from the same source that got the 4x4x4 block selector and contraction index the wrong way round: CONFIRMED on gfx90a 2026-08-22, both accumulator packings. The f32 sweep of mfma_probe covers the A and B layouts (shared between precisions) and the i = 4*g + r packing at ~1e-7, fp32 epsilon; a second run with -DMFMA_F64_USE_16X16 covers the i = g + 4*r packing at ~1e-16. A layout error reads as O(1) here, so neither pass rests on a loose tolerance.

Note the f64 results are bit-identical between this tile and the 4x4x4 one. That is expected rather than suspicious: both decompose K into chunks of four with one MFMA per chunk accumulated in sequence, so the summation order is the same. It is not the signature of dead code, which is what identical results across NWF would be – NWF changes which wavefront takes which column group and so must perturb the schedule. D = A*B + C, wave of 64 lanes, g = lane/16 in 0..3, c = lane%16 in 0..15. A and B share the same layout for both precisions; only the accumulator packing of the 16x16 result differs: A[i][k] : lane holds A[i = c][k = g] B[k][j] : lane holds B[k = g][j = c] D[i][j] : lane holds D[j = c] in accumulator slot r = 0..3, with row i = g + 4*r for f64 (rows spread with stride 4) and i = 4*g + r for f32 (four contiguous rows per lane group)

Definition at line 128 of file mfma_kernel.h.

Here is the caller graph for this function:

◆ mfma_lx_supported()

template<const int LX>
static bool mfma_lx_supported ( )
inlinestatic

Compile-time predicate for the LX values that the MFMA strategy supports. Double precision 4 <= LX <= 12, single precision 2 <= LX <= 12, bounded above by the LDS needed to keep one element resident. MUST match the dispatch specialisations in each operator's kernel header.

Both precisions are verified against a reference on gfx90a (mfma_probe, 144 configurations, 2026-08-22) on the 16x16x4 tile and, in double precision, on the 4x4x4 one. The single precision lower bound is 2 because the batched 4x4x1 tile makes those orders worth having – sixteen lx = 2 elements or four lx = 4 elements to one instruction, and p-multigrid smooths at both – where the 16x16x4 tile would be 2% and 25% utilised. Double precision keeps 4: its 4x4x4 tile already fills M there and batches nothing across elements, so lx = 2 would be a 6% utilised instruction. The tuner's verification pass (NEKO_TUNE_VERIFY, elem_block_tune.h) rejects any candidate that does not reproduce the reference formulation, so an order or tile that is wrong on some part fails loudly at tuning rather than silently in the solve.

Definition at line 199 of file mfma_kernel.h.

Here is the call graph for this function:

◆ mfma_tile_name()

static const char * mfma_tile_name ( const int  c)
inlinestatic

Definition at line 398 of file mfma_kernel.h.

Here is the caller graph for this function:

◆ mfma_tile_offered()

static bool mfma_tile_offered ( const int  c)
inlinestatic

Definition at line 379 of file mfma_kernel.h.

Here is the caller graph for this function:

◆ neko_mfma_candidates()

static int neko_mfma_candidates ( )
inlinestatic

Definition at line 624 of file mfma_kernel.h.

Here is the call graph for this function:
Here is the caller graph for this function:

◆ neko_mfma_pin()

static int neko_mfma_pin ( )
static

Definition at line 592 of file mfma_kernel.h.

Here is the call graph for this function:
Here is the caller graph for this function:

◆ neko_mfma_sweep()

static int neko_mfma_sweep ( )
static

Definition at line 570 of file mfma_kernel.h.

Here is the call graph for this function:
Here is the caller graph for this function: