Neko 1.99.9
A portable framework for high-order spectral element flow simulations
Loading...
Searching...
No Matches
Static Public Member Functions | List of all members
cdtp_dmma_tma_dispatch< T, LX, NW > Struct Template Reference

#include <cdtp_kernel.h>

Static Public Member Functions

static __device__ void run (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 cdtp_dmma_tma_dispatch< T, LX, NW >

Device kernel for D^T x on the fp64 tensor cores, with the element staged by the TMA engine

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

The overlap has to be arranged differently here, and this is the one interesting thing about the variant. Every other TMA kernel in the tree contracts a field that arrives first and consumes its factors afterwards, so the factor copies fly underneath three contractions for free. cdtp needs x and a factor before it can form anything at all, and needs all three factors before the last contraction – a single wait on everything would leave the copies overlapping nothing.

So the four copies are all issued at entry but waited in two stages: x and dr on the first barrier, ds and dt on the second, which is waited only after the axis 0 contraction has been issued. Half the factor traffic flies underneath a contraction; the other half cannot, by the shape of the operator.

That costs one cube over the phase 1 kernels – the weighted fields overwrite their own staged factors, but x has to outlive all three of them, so the result needs its own – for 22032 B and 10 blocks per SM, the smallest TMA block and the best occupancy of any of these variants, and the only one that stays under the 48 kB a block gets without opting in.

Definition at line 584 of file cdtp_kernel.h.

Member Function Documentation

◆ run()

template<typename T , const int LX, const int NW>
static __device__ void cdtp_dmma_tma_dispatch< T, LX, NW >::run ( 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 585 of file cdtp_kernel.h.

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