GPU-capable RF Ray tracing using a domain-specific compiler

Designing an optimized fusion pilot plant requires rapid prediction of whole-facility performance. High-fidelity models such as full wave radio-frequency (RF) for determining power deposition profiles require too much time and resources to explore a vast array of design parameters. Geometric optics, commonly called ray tracing offer a faster approach to determine such models. However, legacy codes have not been built to exploit the hardware architectures of modern computers and their graphical processing units (GPUs). To compound the problem, GPU coding is not compatible between different manufacturers. Without proper abstraction, codes written for the CUDA architecture of Summit need to be rewritten to support the AMD architecture of Frontier. We present a new ray tracing code built using a graph computational framework. Physics equations are compiled into a graph of mathematical operations. Transformations of that graph apply auto differentiation to build ray update expressions. Reductions of the graph to reduce the problem to its simplest form. This graph of operations can be just in time (JIT) compiled to an optimized GPU or CPU kernels for a specific architecture. Using this framework, we demonstrate correct fully 3D ray behavior for realistic tokamak geometries for different plasma dispersion functions.