This is a subproject of the ELLIIT GPAI C05 General Purpose AI Computing project. The idea is to apply ML methods and AI hardware to scientific computing for domains such as Computational Fluid Dynamics. The specific focus of this project is solving partial differential equations using tensor operations and, if necessary, physics informed neural networks.
Preliminary results and a prototype framework already allow us to describe and generate proper backend code for CPU, GPU and TPU individually. Experimental results show promising performance with CUDA code on GPU, on par with hand-written state of the art solvers. However, testing and optimizing our solver and framework for TPU or NPUs or any similar ML hardware accelerators is beyond our local capabilities. We plan to involve the NAISS Alvis infrastructure to evaluate and improve the performance of our PDE framework. Depending on the stability and solution convergence speed we might also involve physics informed neural networks (PINN) to fine tune our solver. We suspect the number representation available in the ML hardware will affect solver accuracy and convergence in a different manner than in the CUDA/GPU code. PINN training may offer a way to overcome accuracy issues in a different way.
If successful, we will be able to speed up solving e.g. Navier-Stokes equations by employing heterogeneous systems of both GPUs and tensor cores in parallel. This will also allow us to identify bottlenecks in current TPU architectures and suggest specific software patterns and hardware architectures to overcome such bottlenecks.
We have already published some results here:
https://doi.org/10.1007/s10766-025-00791-6
Since then we have improved the framework and we can get performance close to hand written CUDA on GPUs. We do not yet know how this translates to ML hardware such as TPUs or NPUs. This is what we want to investigate.