#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_vector_dispatch< T, LX, NW >

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

The same trade as ax_helm_dmma_tma_elem() makes for the scalar operator, applied to the thing that was holding the vector one back. The register variant above reads the seven geometric factors once and keeps them in registers across all three components, which is the right access pattern and the wrong place to put it: ~60 registers at nw = 4, measured on GH200 as a loss against the kstep variant even though it moves 13 arrays per element rather than 27. Here they are read once into shared memory instead, by the TMA engine, and the register file is left alone.

Everything else follows the scalar TMA kernel. The factor copies are issued before the component loop and waited on only at the first pointwise step, so they arrive underneath the staging and the first three contractions of component 0; after that they are simply resident for components 1 and 2. Each component is staged by one bulk copy and stored by another.

The block footprint is the scalar kernel's exactly – four working cubes plus seven factor cubes – because the components run one at a time through the same four. That leaves no room for a second component cube, so a component's staging is not overlapped with the previous component's contractions; doing that needs either the opt-in dynamic shared memory path or h1 moved back into registers to free a cube. Neither is done here.

See dmma_tma_kernel.h for the primitives, the sm_90 and toolkit guards and the lx == DMMA_P bound.

Definition at line 1862 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_vector_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 1863 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()