#include <ax_helm_kernel.h>

Static Public Member Functions
static __device__ void	run (T __restrict__, T __restrict__, T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__, const T __restrict__)

Detailed Description

template<typename T, const int LX, const int NW>
struct ax_helm_dmma_tma_batch_dispatch< T, LX, NW >

Device kernel for the vector axhelm on the fp64 tensor cores, with the whole element staged by the TMA engine in one batch

The component-at-a-time variant above measured 4.9% behind the register hoisting DMMA kernel at nw = 4 on GH200 while running more blocks per SM than it (five against four, ptxas confirmed) – so residency is not what separates them. What separates them is how many bulk copies are outstanding. The scalar TMA kernel issues eight at once, covering every byte the element reads, and wins at 20 warps per SM against 48; the component-at-a-time vector kernel issues seven once and then, for components 1 and 2, a lone 4 kB copy that one thread issues and every other thread then waits on. A 4 kB copy is small, its fixed cost only disappears when several are in flight, and paying it alone three times an element is the deficit.

So this variant issues all ten of an element's input copies at entry – the three components and the seven geometric factors – and stores the three results without waiting on each other. Each component keeps its own cube, in and then out, which is what removes both serialisations at once: nothing has to be restaged, so no load waits on a store.

The cost is thirteen cubes rather than four, 54800 B, past the 48 kB a block gets for free – hence dynamic shared memory and the one-time opt-in in ax_helm_dmma_tma_batch_optin() below. That lands at four blocks per SM at nw = 4, which is exactly what the DMMA kernel it is competing with already runs at, so the batching is bought with no occupancy at all.

Only the first component's cube is waited on before work starts; the other nine copies are waited on at the first pointwise step, and components 1 and 2 never wait at all. See dmma_tma_kernel.h for the primitives, the struct that fixes the layout, the sm_90 and toolkit guards and the lx bound.

Definition at line 2135 of file ax_helm_kernel.h.

Member Function Documentation

◆ run()

template<typename T , const int LX, const int NW>

static __device__ void ax_helm_dmma_tma_batch_dispatch< T, LX, NW >::run	(	T *	__restrict__,
		T *	__restrict__,
		T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__,
		const T *	__restrict__
	)

inlinestatic

Definition at line 2136 of file ax_helm_kernel.h.

The documentation for this struct was generated from the following file:

/home/runner/work/neko/neko/src/math/bcknd/device/cuda/ax_helm_kernel.h

Static Public Member Functions

Detailed Description

Member Function Documentation

◆ run()