#include <opgrad_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__, const T *__restrict__)

Detailed Description

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

Device kernel for the weak gradient on the fp64 tensor cores, with the element staged by the TMA engine

Same three contractions and the same cube layout as opgrad_dmma_elem(); at lx == DMMA_P the padding is empty and the cube offset is the flat point index, so the arithmetic is identical and only the staging differs.

All ten input copies – u and the nine metric cubes, 77% of the traffic – are issued together at kernel entry, u on its own mbarrier and waited immediately, the metrics on a second and waited only at the pointwise step that consumes them. Batching every read into one issue is what separated the scalar axhelm TMA kernel that won from the component-at-a-time vector one that did not: a lone 4 kB bulk copy never amortises the engine's fixed latency.

Thirteen cubes is 54800 B, past the 48 kB a block gets for free, so the allocation is dynamic and its layout is opgrad_tma_smem in dmma_tma_kernel.h. That is byte for byte the batched axhelm block, so it runs at the same four blocks per SM and reuses the same device gate.

The three outputs leave as ordinary coalesced stores, see the note on the struct: only u's cube is free by then, and the stores are 19% of the traffic against the loads' 77%.

Definition at line 617 of file opgrad_kernel.h.

Member Function Documentation

◆ run()

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

static __device__ void opgrad_dmma_tma_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__,
		const T *	__restrict__
	)

inlinestatic

Definition at line 618 of file opgrad_kernel.h.

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

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

Static Public Member Functions

Detailed Description

Member Function Documentation

◆ run()