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 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_CANDIDATES   4
 
#define NEKO_MFMA_NWF(C)   (1 << (C))
 
#define NEKO_MFMA_NTHRDS(C)   dim3(64, NEKO_MFMA_NWF(C), 1)
 
#define NEKO_MFMA_NGROUPS(LX)   (((LX) * (LX) + 15) / 16)
 
#define NEKO_MFMA_EB_N(NWF, LX)
 
#define NEKO_MFMA_EB(LX, C)   NEKO_MFMA_EB_N(NEKO_MFMA_NWF(C), LX)
 
#define NEKO_MFMA_WPE(LX, C)   (NEKO_MFMA_NWF(C) / NEKO_MFMA_EB(LX, C))
 
#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)
 

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 int neko_mfma_sweep ()
 
static int neko_mfma_env ()
 

Macro Definition Documentation

◆ NEKO_MFMA_CANDIDATES

#define NEKO_MFMA_CANDIDATES   4

Definition at line 199 of file mfma_kernel.h.

◆ NEKO_MFMA_EB

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

Definition at line 242 of file mfma_kernel.h.

◆ NEKO_MFMA_EB_N

#define NEKO_MFMA_EB_N (   NWF,
  LX 
)
Value:
((NWF) / NEKO_MFMA_NGROUPS(LX) >= 8 ? 8 : \
(NWF) / NEKO_MFMA_NGROUPS(LX) >= 4 ? 4 : \
(NWF) / NEKO_MFMA_NGROUPS(LX) >= 2 ? 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_NGROUPS(LX)

Definition at line 238 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 245 of file mfma_kernel.h.

◆ NEKO_MFMA_NGROUPS

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

Definition at line 235 of file mfma_kernel.h.

◆ NEKO_MFMA_NTHRDS

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

Definition at line 201 of file mfma_kernel.h.

◆ NEKO_MFMA_NWF

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

Definition at line 200 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 244 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 %dwf %-2de: %9.2f us/call", \
} \
} while (0)
#define NEKO_TUNE_US(T, ITERS)
#define NEKO_TUNE_INIT
#define NEKO_MFMA_CANDIDATES
#define NEKO_MFMA_NWF(C)
#define NEKO_MFMA_EB(LX, C)

Definition at line 285 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 133 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.

Recovered onto develop 2026-08-22 from the feature/mfma branch, scoped down to Ax-helm: the other operators it covered (dudxyz, cdtp, conv1, convect_scalar) are the four worst placed in the arithmetic intensity ranking and are not worth carrying without a measurement. The contraction primitives below are unchanged from the hardware-validated versions – see the layout note on mfma_contract_4x4 for why they must not be re-derived from host simulation alone.

Maps a reference derivative-matrix * field contraction onto the AMD matrix cores on CDNA2 (MI250X / gfx90a) and CDNA3 (MI300A / MI300X / gfx942). Double precision defaults to the batched 4x4x4 tile (v_mfma_f64_4x4x4f64, full M-utilisation for M = LX < 16; see mfma_contract_4x4), single precision uses the 16x16x4 tile (v_mfma_f32_16x16x4f32; no f32 4x4x4 equivalent). Double precision can be forced back onto 16x16x4 with -DMFMA_F64_USE_16X16. One thread block is launched as a single 64-lane wavefront (blockDim = (64,1,1), or (64, NWF, 1) for the multi-wavefront Ax-helm) and processes one element; each contraction is a D * U GEMM with M = LX, N = LX^2, K = LX, the partial tiles masked off.

Supported for single and double precision and 4 <= LX <= 12; the upper bound is set by the LDS needed to keep the cubes resident (operators stage up to 4*LX^3 + 3*LX^2 elements, ~57 KB of f64 at LX = 12).

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 117 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. Single or double precision and 4 <= LX <= 12 (bounded above by the LDS needed to keep one element resident). MUST match the dispatch specialisations in each operator's kernel header.

Definition at line 182 of file mfma_kernel.h.

Here is the call graph for this function:

◆ neko_mfma_env()

static int neko_mfma_env ( )
static

Definition at line 272 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 261 of file mfma_kernel.h.

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