#include <stdlib.h>#include <stdio.h>#include <cuda_runtime.h>#include <mma.h>#include <device/device_config.h>
Go to the source code of this file.
Macros | |
| #define | NEKO_DMMA_CANDIDATES 3 |
| #define | NEKO_DMMA_NW(C) (1 << ((C) + 1)) |
| #define | NEKO_DMMA_NTHRDS(C) dim3(32 * NEKO_DMMA_NW(C), 1, 1) |
| #define | NEKO_DMMA_PPA(LX) ((DMMA_P % (LX) == 0) ? (DMMA_P / (LX)) : 1) |
| #define | NEKO_DMMA_PACK(LX) (NEKO_DMMA_PPA(LX) * NEKO_DMMA_PPA(LX) * NEKO_DMMA_PPA(LX)) |
| #define | NEKO_DMMA_NBLCKS(NELV, LX) dim3(((NELV) + NEKO_DMMA_PACK(LX) - 1) / NEKO_DMMA_PACK(LX), 1, 1) |
| #define | NEKO_TUNE_LOG_DMMA(LX, T3) |
Enumerations | |
| enum | { DMMA_P = 8 , DMMA_KS = 4 , DMMA_SI = DMMA_P , DMMA_SJ = DMMA_P * DMMA_P , DMMA_CUBE = DMMA_P * DMMA_P * DMMA_P , DMMA_MAT = DMMA_P * DMMA_P , DMMA_NTILES = DMMA_P , DMMA_KSTEPS = DMMA_P / DMMA_KS } |
Functions | |
| static int | neko_dmma_env () |
| static bool | dmma_arch_compiled () |
| static bool | cuda_have_dmma () |
| template<const int LX> | |
| static bool | dmma_lx_supported () |
| template<const int LX> | |
| static bool | dmma_vector_lx_supported () |
Macro Definition Documentation
◆ NEKO_DMMA_CANDIDATES
| #define NEKO_DMMA_CANDIDATES 3 |
Definition at line 123 of file dmma_kernel.h.
◆ NEKO_DMMA_NBLCKS
| #define NEKO_DMMA_NBLCKS | ( | NELV, | |
| LX | |||
| ) | dim3(((NELV) + NEKO_DMMA_PACK(LX) - 1) / NEKO_DMMA_PACK(LX), 1, 1) |
Definition at line 150 of file dmma_kernel.h.
◆ NEKO_DMMA_NTHRDS
| #define NEKO_DMMA_NTHRDS | ( | C | ) | dim3(32 * NEKO_DMMA_NW(C), 1, 1) |
Definition at line 125 of file dmma_kernel.h.
◆ NEKO_DMMA_NW
Definition at line 124 of file dmma_kernel.h.
◆ NEKO_DMMA_PACK
| #define NEKO_DMMA_PACK | ( | LX | ) | (NEKO_DMMA_PPA(LX) * NEKO_DMMA_PPA(LX) * NEKO_DMMA_PPA(LX)) |
Definition at line 148 of file dmma_kernel.h.
◆ NEKO_DMMA_PPA
Definition at line 147 of file dmma_kernel.h.
◆ NEKO_TUNE_LOG_DMMA
Definition at line 169 of file dmma_kernel.h.
Enumeration Type Documentation
◆ anonymous enum
Double precision matrix-core (DMMA) primitives for the spectral element tensor contractions, the CUDA counterpart of the AMD MFMA strategies.
Maps a reference derivative matrix * field contraction onto the fp64 tensor cores, which are exposed as a single fixed 8x8x4 WMMA tile lowering to mma.sync.aligned.m8n8k4.f64. Each contraction is a D * U GEMM with M = LX, N = LX^2, K = LX, and rather than mask the partial M, N and K tiles that a general LX would leave, every cube is staged into shared memory padded to DMMA_P^3 and every derivative matrix to DMMA_P^2, with the padding zero filled. The GEMM is then always 8 x 64 x 8 and every tile is full:
- zero rows of D contribute nothing to a padded output row,
- zero rows of the staged cube contribute nothing to any output,
- padded output entries are never read back out of shared memory,
so no lane ever has to test an index. At LX = 8 the padding is empty and the staging is a straight copy, which is the case the strategy is aimed at.
That fixed tile is also what bounds the strategy: LX > DMMA_P would need the cubes padded to 16^3, i.e. 4 * 16^3 * 8 = 128 kB of shared memory for the four cubes, well past the 48 kB a block gets without the opt-in dynamic path. Supported for double precision and 2 <= LX <= 8 only – the lower bound is 2 rather than 4 because of the packing below, see dmma_lx_supported(); note that fp32 has no tensor core equivalent, only TF32 with an 11 bit significand, which is not usable for the operator inside a Krylov solve, so a single precision build has no DMMA strategy.
The fp64 tensor cores themselves only exist on the data centre parts: sm_80 (A100) and sm_90 (H100 / GH200), where they run at twice the fp64 FMA rate. Consumer sm_86 / sm_89 assemble the instruction but have no fp64 tensor hardware at all, and Blackwell moves the fp64 rates again, so both the compile time guard and cuda_have_dmma() below allow-list [sm_80, sm_90] rather than testing for "at least sm_80". Widen both together if a later part turns out to be worth measuring.
The result is bit-for-bit an fp64 computation – mma.m8n8k4.f64 is IEEE double throughout, unlike the reduced precision tensor tiles – but the summation order differs from the scalar variants, so the Ax output differs in the last bits, exactly as it already does between the 1d and kstep variants.
Cube and matrix layout. A staged cube is
cube[i + DMMA_SI * j + DMMA_SJ * k], i, j, k < DMMA_P
and a staged derivative matrix is dmat[m + DMMA_P * l] = D(m, l), matching the dx[i + l*LX] = D(i,l) convention of the scalar kernels. Contracting axis AXIS turns the cube into an M x N matrix whose (m, n) element sits at a constant stride, which is all load_matrix_sync() needs:
AXIS = 0: m = i, n = (j,k), element at m + DMMA_SI * n -> column major, ldm = DMMA_SI, one tile per group of 8 (j,k) pairs AXIS = 1: m = j, n = i, element at n + DMMA_SI * m + DMMA_SJ * k -> row major, ldm = DMMA_SI, one tile per k slab AXIS = 2: m = k, n = i, element at n + DMMA_SI * j + DMMA_SJ * m -> row major, ldm = DMMA_SJ, one tile per j slab
Each view has DMMA_P tiles and the same view describes the input and the output of a contraction, so one traits struct covers both.
| Enumerator | |
|---|---|
| DMMA_P | |
| DMMA_KS | |
| DMMA_SI | |
| DMMA_SJ | |
| DMMA_CUBE | |
| DMMA_MAT | |
| DMMA_NTILES | |
| DMMA_KSTEPS | |
Definition at line 105 of file dmma_kernel.h.
Function Documentation
◆ cuda_have_dmma()
Returns true if the current device exposes the fp64 tensor cores used by the DMMA strategies and this build can actually reach them. Result is cached after the first query.
Definition at line 224 of file dmma_kernel.h.
◆ dmma_arch_compiled()
◆ dmma_lx_supported()
Compile-time predicate for the LX values that the DMMA strategy supports. Double precision and 2 <= LX <= DMMA_P, see the note on the padded staging above. The lower bound is 2 rather than 4 because packing makes LX = 2 useful rather than absurd – sixty-four elements to a cube, and it is a p-multigrid level. LX = 3 is supported but packs one element and wastes 95% of every contraction; the tuner will reject it.
MUST match the dispatch specialisations in each operator's kernel header.
The double precision requirement is deliberate and is NOT an oversight to be brought in line with the HIP side, which does offer both precisions. That difference is hardware: AMD's v_mfma_f32_16x16x4f32 is true IEEE fp32, so an sp build there loses nothing, while NVIDIA has no fp32 tensor path at all – only TF32, whose 11 bit significand (unit roundoff 4.9e-4, against fp32's 6e-8) is not usable for the operator inside a Krylov solve. The accuracy preserving alternative, 3xTF32 splitting, costs three MMAs to buy back precision in a kernel that is bandwidth bound with flops to spare.
The empirical argument is stronger than the roofline one. AMD ran the experiment NVIDIA cannot: with a true fp32 matrix core and therefore no accuracy compromise whatsoever, MFMA still measured 4.5% BEHIND a plain 1d kernel at lx = 8, and the whole fp32 speedup over fp64 (2.05x) came from moving half the bytes at identical achieved bandwidth. A path that loses where it is free will not win where it is lossy.
Definition at line 273 of file dmma_kernel.h.
◆ dmma_vector_lx_supported()
The same predicate for the vector operator, whose supported range is not the same and must not be assumed to be.
The vector kernel stages one element per block and does not pack, so the lower bound stays at 4: below it the padding waste that packing removes for the scalar kernel is still there, and there is nothing to gain. That is a performance argument, but the predicate is a correctness one – the tuner launches whatever it is told is supported, and an LX the vector dispatch does not specialise resolves to the no-op primary template, which times as free, wins the comparison and leaves stale values in au/av/aw. Lowering dmma_lx_supported() to 2 for the packed scalar kernel is exactly how that happened.
MUST match NEKO_AX_HELM_DMMA_VECTOR_DISPATCH in ax_helm_kernel.h.
Definition at line 295 of file dmma_kernel.h.
◆ neko_dmma_env()
Definition at line 154 of file dmma_kernel.h.
