#include <dudxyz_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__, const int) |
Detailed Description
Device kernel for the derivative on the fp64 tensor cores
One element per block – or a packed cube of them at low order – NW warps, and the element staged into shared memory as four DMMA_P^3 cubes: the staged input and the three reference derivatives. This is the first half of ax_helm_dmma_elem() and nothing else: the same three D * U GEMMs handed to dmma_contract(), but the geometric factors turn the reference derivatives into the physical one pointwise and it goes straight back out to global memory, where axhelm contracts a second time with D^T. See dmma_kernel.h for the padded staging, the per axis matrix views and the arch and LX bounds.
Half the contractions for two thirds of the bytes is the reason to expect anything here. What won axhelm was not the tensor cores – it reached 82% of the achievable streaming bandwidth on GH200 at 7.8% of the fp64 tensor peak – but the access pattern the staging buys, and the cost of that staging is paid per contraction while the memory time it hides behind is paid per byte. At lx = 8 this operator has 8192 bytes per contraction against axhelm's 6144, so the trade is 33% better here, not worse.
Definition at line 357 of file dudxyz_kernel.h.
Member Function Documentation
◆ run()
The documentation for this struct was generated from the following file:
- /home/runner/work/neko/neko/src/math/bcknd/device/cuda/dudxyz_kernel.h
