dmma_tma_kernel.h File Reference

#include <stdlib.h>
#include <stdint.h>
#include <stddef.h>
#include <cuda_runtime.h>
#include <device/device_config.h>
#include "dmma_kernel.h"

Include dependency graph for dmma_tma_kernel.h:

This graph shows which files directly or indirectly include this file:

Go to the source code of this file.

Classes
struct	dmma_tma_batch_smem
struct	opgrad_tma_smem
struct	conv1_tma_smem

Macros
#define	DMMA_NG   7
#define	DMMA_NG_OPGRAD   9
#define	DMMA_NG_CONV1   13
#define	NEKO_TMA_TOOLKIT   0
#define	NEKO_DMMA_TMA_BATCH_SMEM   ((int) sizeof(dmma_tma_batch_smem))
#define	NEKO_OPGRAD_TMA_SMEM   ((int) sizeof(opgrad_tma_smem))
#define	NEKO_CONV1_TMA_SMEM   ((int) sizeof(conv1_tma_smem))
#define	NEKO_TUNE_LOG_DMMA_TMA_BATCH(LX, T5)
#define	NEKO_TUNE_LOG_DMMA_TMA(LX, T4)

Functions
static bool	tma_arch_compiled ()
static bool	cuda_have_tma ()
template<const int LX>
static bool	dmma_tma_lx_supported ()
template<const int LX>
static bool	dmma_tma_vector_lx_supported ()
template<const int LX>
static bool	dmma_tma_dudxyz_lx_supported ()
template<const int LX>
static bool	dmma_tma_opgrad_lx_supported ()
template<const int LX>
static bool	dmma_tma_conv1_lx_supported ()
template<const int LX>
static bool	dmma_tma_cdtp_lx_supported ()
static int	cuda_tma_smem_optin ()
static bool	cuda_have_tma_batch ()
static bool	cuda_have_tma_opgrad ()
static bool	cuda_have_tma_conv1 ()
template<const int LX>
static bool	dmma_tma_batch_lx_supported ()
static bool	dmma_tma_ptr_aligned (const void *p)
static bool	dmma_tma_aligned (const void *w, const void *u, const void *h1, const void *g11, const void *g22, const void *g33, const void *g12, const void *g13, const void *g23)
static bool	dmma_tma_vector_aligned (const void *au, const void *av, const void *aw, const void *u, const void *v, const void *w, const void *h1, const void *g11, const void *g22, const void *g33, const void *g12, const void *g13, const void *g23)
static bool	dmma_tma_dudxyz_aligned (const void *du, const void *u, const void *dr, const void *ds, const void *dt, const void *jacinv)
static bool	dmma_tma_metrics_aligned (const void *drdx, const void *dsdx, const void *dtdx, const void *drdy, const void *dsdy, const void *dtdy, const void *drdz, const void *dsdz, const void *dtdz)
static bool	dmma_tma_opgrad_aligned (const void *u, const void *drdx, const void *dsdx, const void *dtdx, const void *drdy, const void *dsdy, const void *dtdy, const void *drdz, const void *dsdz, const void *dtdz)
static bool	dmma_tma_cdtp_aligned (const void *dtx, const void *x, const void *dr, const void *ds, const void *dt)
static bool	dmma_tma_conv1_aligned (const void *du, const void *u, const void *vx, const void *vy, const void *vz, const void *jacinv, const void *drdx, const void *dsdx, const void *dtdx, const void *drdy, const void *dsdy, const void *dtdy, const void *drdz, const void *dsdz, const void *dtdz)
static int	neko_dmma_tma_env ()

Macro Definition Documentation

DMMA_NG

#define DMMA_NG   7

Tensor Memory Accelerator (TMA) staging for the fp64 tensor core kernels, Hopper only.

The DMMA kernels in dmma_kernel.h stage an element into shared memory and hand the six tensor contractions to the fp64 tensor cores. What they do not do is overlap the staging with the contractions, and the staging is nearly all of the kernel: at lx = 8 an element is nine cubes of 4 kB – the field u, the seven geometric factors h1 and g11..g23, and the result w – of which the seven factors are 78% of the read traffic and are read in a single pointwise pass that sits between two __syncthreads() with nine multiply-adds per point to hide them behind. Nothing else in the block is running while those 28 kB arrive.

This variant issues that traffic up front, as bulk asynchronous copies on two mbarriers: the cube of u on one, waited on immediately, and the seven factor cubes on the other, waited on only at the pointwise step, so they are in flight underneath the first three contractions. The result cube goes back out the same way, one bulk store instead of DMMA_CUBE scalar ones.

The register-hoisting alternative – read the factors into registers before the contractions rather than into shared memory – is already implemented, in ax_helm_dmma_vector_elem(), and measured a loss on GH200 at lx = 8: the factors cost about sixty registers there and give back more occupancy than the overlap wins. Landing them in shared memory instead is the same hoist without the register bill, which is the argument for spending the shared memory here, and is what this variant is for.

It is not free either. Seven extra cubes take the block from 17.5 kB to 45.5 kB, and on an SM with 228 kB that is five resident blocks rather than the eight it takes to fill the 64 warp limit – 62.5% occupancy against 100%. The bet is that TMA replaces the memory level parallelism that the lost warps were providing: five blocks each with eight bulk copies in flight is 160 kB of outstanding traffic per SM, generated by one thread per block rather than by thousands of LDGs. Whether that bet pays is a measurement, so the variant is a tuner candidate like every other one and is never assumed to win.

Scope: lx == DMMA_P, double precision, sm_90.

The lx bound is what makes the cheap TMA path usable at all. A bulk copy moves a contiguous run of bytes, and at lx == DMMA_P a staged cube is the element's contiguous global chunk – dmma_pack< LX, 1 >::map() reduces to the identity, so the copy needs no addressing beyond a base pointer and a length. Every smaller lx either strides into the padded cube (PPA == 1) or scatters sub-cubes across it (PPA > 1), neither of which a bulk copy can express; those need cp.async.bulk.tensor and a CUtensorMap. That is not merely more work, it is the wrong tool here: a tensor map is built on the host by cuTensorMapEncodeTiled() and bakes in the base pointer, while u and w are Krylov vectors that change on every Ax call, so the descriptors would have to be re-encoded per call inside the solve. And the only lx values that would need it are the packed ones, where DMMA already measures 40% behind the kstep variant. So the tensor path is deliberately not taken.

Definition at line 101 of file dmma_tma_kernel.h.

DMMA_NG_CONV1

#define DMMA_NG_CONV1   13

Definition at line 111 of file dmma_tma_kernel.h.

DMMA_NG_OPGRAD

#define DMMA_NG_OPGRAD   9

Definition at line 107 of file dmma_tma_kernel.h.

NEKO_CONV1_TMA_SMEM

#define NEKO_CONV1_TMA_SMEM   ((int) sizeof(conv1_tma_smem))

Definition at line 401 of file dmma_tma_kernel.h.

NEKO_DMMA_TMA_BATCH_SMEM

#define NEKO_DMMA_TMA_BATCH_SMEM   ((int) sizeof(dmma_tma_batch_smem))

Definition at line 295 of file dmma_tma_kernel.h.

NEKO_OPGRAD_TMA_SMEM

#define NEKO_OPGRAD_TMA_SMEM   ((int) sizeof(opgrad_tma_smem))

Definition at line 372 of file dmma_tma_kernel.h.

NEKO_TMA_TOOLKIT

#define NEKO_TMA_TOOLKIT   0

Definition at line 122 of file dmma_tma_kernel.h.

NEKO_TUNE_LOG_DMMA_TMA

#define NEKO_TUNE_LOG_DMMA_TMA	(		LX,
			T4
	)

Value:

  do {                                                                        \
    for (int c = 0; c < NEKO_DMMA_CANDIDATES; c++) {                          \
      if ((T4)[c] >= NEKO_TUNE_INIT) { continue; }                            \
      char lbl_[16];                                                          \
      sprintf(lbl_, "TMA   %dw", NEKO_DMMA_NW(c));                            \
      sprintf(neko_log_buf, "%-13s: %9.2f us/call", lbl_,                     \
              NEKO_TUNE_US((T4)[c], iters));                                  \
      log_message(neko_log_buf);                                              \
    }                                                                         \
  } while (0)

Definition at line 580 of file dmma_tma_kernel.h.

NEKO_TUNE_LOG_DMMA_TMA_BATCH

#define NEKO_TUNE_LOG_DMMA_TMA_BATCH	(		LX,
			T5
	)

Value:

  do {                                                                        \
    for (int c = 0; c < NEKO_DMMA_CANDIDATES; c++) {                          \
      if ((T5)[c] >= NEKO_TUNE_INIT) { continue; }                            \
      char lbl_[16];                                                          \
      sprintf(lbl_, "TMAB  %dw", NEKO_DMMA_NW(c));                            \
      sprintf(neko_log_buf, "%-13s: %9.2f us/call", lbl_,                     \
              NEKO_TUNE_US((T5)[c], iters));                                  \
      log_message(neko_log_buf);                                              \
    }                                                                         \
  } while (0)

Definition at line 567 of file dmma_tma_kernel.h.

Function Documentation

cuda_have_tma()

static bool cuda_have_tma ( )

inlinestatic

Returns true if the current device has a TMA engine and this build can actually reach it. Result is cached after the first query.

Definition at line 165 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

cuda_have_tma_batch()

static bool cuda_have_tma_batch ( )

inlinestatic

Definition at line 337 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

cuda_have_tma_conv1()

static bool cuda_have_tma_conv1 ( )

inlinestatic

Definition at line 422 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

cuda_have_tma_opgrad()

static bool cuda_have_tma_opgrad ( )

inlinestatic

Whether the device will hand a block each of those allocations. opgrad's is byte for byte the batched axhelm one; conv1's is larger and is queried separately rather than assumed to follow.

Definition at line 417 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

cuda_tma_smem_optin()

static int cuda_tma_smem_optin ( )

inlinestatic

Returns true if the current device will hand a block the batched variant's dynamic allocation. Past 48 kB a block only gets shared memory it has explicitly opted into, and the opt-in ceiling is a device attribute rather than a property of the architecture, so it is queried. Cached.

Definition at line 311 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

dmma_tma_aligned()

static bool dmma_tma_aligned ( const void *  w, const void *  u, const void *  h1, const void *  g11, const void *  g22, const void *  g33, const void *  g12, const void *  g13, const void *  g23  )

inlinestatic

Definition at line 455 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

dmma_tma_batch_lx_supported()

template<const int LX>

static bool dmma_tma_batch_lx_supported ( )

inlinestatic

Compile-time predicate for the LX values the batched variant supports. Same bound as the other two and written out for the same reason, see dmma_tma_vector_lx_supported().

MUST match NEKO_AX_HELM_DMMA_TMA_BATCH_DISPATCH in ax_helm_kernel.h.

Definition at line 435 of file dmma_tma_kernel.h.

Here is the call graph for this function:

dmma_tma_cdtp_aligned()

static bool dmma_tma_cdtp_aligned ( const void *  dtx, const void *  x, const void *  dr, const void *  ds, const void *  dt  )

inlinestatic

Definition at line 527 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

dmma_tma_cdtp_lx_supported()

template<const int LX>

static bool dmma_tma_cdtp_lx_supported ( )

inlinestatic

And for cdtp. Same bound, written out per operator for the same reason.

MUST match NEKO_CDTP_DMMA_TMA_DISPATCH in cdtp_kernel.h.

Definition at line 259 of file dmma_tma_kernel.h.

Here is the call graph for this function:

dmma_tma_conv1_aligned()

static bool dmma_tma_conv1_aligned ( const void *  du, const void *  u, const void *  vx, const void *  vy, const void *  vz, const void *  jacinv, const void *  drdx, const void *  dsdx, const void *  dtdx, const void *  drdy, const void *  dsdy, const void *  dtdy, const void *  drdz, const void *  dsdz, const void *  dtdz  )

inlinestatic

Definition at line 537 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

dmma_tma_conv1_lx_supported()

template<const int LX>

static bool dmma_tma_conv1_lx_supported ( )

inlinestatic

Definition at line 248 of file dmma_tma_kernel.h.

Here is the call graph for this function:

dmma_tma_dudxyz_aligned()

static bool dmma_tma_dudxyz_aligned ( const void *  du, const void *  u, const void *  dr, const void *  ds, const void *  dt, const void *  jacinv  )

inlinestatic

Definition at line 483 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

dmma_tma_dudxyz_lx_supported()

template<const int LX>

static bool dmma_tma_dudxyz_lx_supported ( )

inlinestatic

The same predicate for the derivative operator, written out for the same reason: an lx its dispatch does not specialise resolves to a no-op that times as free and wins the tuner. Note that the DMMA bound is NOT written out per operator in the same way – dmma_lx_supported() covers every operator whose dispatch specialises 2..DMMA_P, which dudxyz does.

MUST match NEKO_DUDXYZ_DMMA_TMA_DISPATCH in dudxyz_kernel.h.

Definition at line 229 of file dmma_tma_kernel.h.

Here is the call graph for this function:

dmma_tma_lx_supported()

template<const int LX>

static bool dmma_tma_lx_supported ( )

inlinestatic

Compile-time predicate for the LX values the TMA staging supports, see the scope note above: double precision and lx == DMMA_P exactly.

MUST match NEKO_AX_HELM_DMMA_TMA_DISPATCH in ax_helm_kernel.h. An LX the dispatch does not specialise resolves to the no-op primary template, which times as free and wins the tuner comparison – the same trap documented at dmma_vector_lx_supported().

Definition at line 199 of file dmma_tma_kernel.h.

Here is the call graph for this function:

dmma_tma_metrics_aligned()

static bool dmma_tma_metrics_aligned ( const void *  drdx, const void *  dsdx, const void *  dtdx, const void *  drdy, const void *  dsdy, const void *  dtdy, const void *  drdz, const void *  dsdz, const void *  dtdz  )

inlinestatic

Definition at line 493 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

dmma_tma_opgrad_aligned()

static bool dmma_tma_opgrad_aligned ( const void *  u, const void *  drdx, const void *  dsdx, const void *  dtdx, const void *  drdy, const void *  dsdy, const void *  dtdy, const void *  drdz, const void *  dsdz, const void *  dtdz  )

inlinestatic

Definition at line 511 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

dmma_tma_opgrad_lx_supported()

template<const int LX>

static bool dmma_tma_opgrad_lx_supported ( )

inlinestatic

And for opgrad and conv1. Same bound, written out per operator for the same reason as above.

MUST match NEKO_OPGRAD_DMMA_TMA_DISPATCH in opgrad_kernel.h and NEKO_CONV1_DMMA_TMA_DISPATCH in conv1_kernel.h.

Definition at line 242 of file dmma_tma_kernel.h.

Here is the call graph for this function:

dmma_tma_ptr_aligned()

static bool dmma_tma_ptr_aligned ( const void *  p )

inlinestatic

A bulk copy requires both of its addresses to be 16 byte aligned and its length to be a multiple of 16, and gives no diagnostic when they are not. The length is DMMA_CUBE doubles and the per element offset is a multiple of DMMA_CUBE, so everything rests on the nine pointers the operator is handed: they are whole device_map()ed fields today and so are allocation aligned, but that is a property of every caller rather than of this kernel, and one caller passing an odd element offset would make the copies silently undefined. Cheap enough to check once per tune and fall back.

Definition at line 450 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

dmma_tma_vector_aligned()

static bool dmma_tma_vector_aligned ( const void *  au, const void *  av, const void *  aw, const void *  u, const void *  v, const void *  w, const void *  h1, const void *  g11, const void *  g22, const void *  g33, const void *  g12, const void *  g13, const void *  g23  )

inlinestatic

Definition at line 469 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

dmma_tma_vector_lx_supported()

template<const int LX>

static bool dmma_tma_vector_lx_supported ( )

inlinestatic

The same predicate for the vector operator. Identical to the scalar bound today, and written out anyway rather than aliased: the pair it mirrors, dmma_lx_supported() and dmma_vector_lx_supported(), are not the same, and an lx the vector dispatch does not specialise silently resolves to a no-op that times as free and wins the tuner.

MUST match NEKO_AX_HELM_DMMA_TMA_VECTOR_DISPATCH in ax_helm_kernel.h.

Definition at line 214 of file dmma_tma_kernel.h.

Here is the call graph for this function:

neko_dmma_tma_env()

static int neko_dmma_tma_env ( )

static

Definition at line 555 of file dmma_tma_kernel.h.

Here is the call graph for this function:

Here is the caller graph for this function:

tma_arch_compiled()

static bool tma_arch_compiled ( )

inlinestatic

Definition at line 149 of file dmma_tma_kernel.h.

Here is the caller graph for this function: