Pseudo-spectral direct numerical simulation (DNS) of the Taylor-Green vortex on AWS Neuron: NKI kernels on Inferentia2 and Trainium1, 3-D FFTs as matmuls, all-to-all collectives inside the kernel and a libnrt C driver, in fp32 up to 512^3, checked against an fp64 oracle and Trainium1 references. Sample code for HPC and CFD engineers. -
View it on GitHub