Researchers unlock accelerator performance beyond CUDA, enabling exa-scale computing for all platforms

Modern supercomputers depend on specialised processors like GPUs to deliver the immense performance needed for complex scientific simulations, but programming for these diverse systems presents a significant challenge. Researchers at Oak Ridge National Laboratory and Diditco, led by M. Cianciosa, D. Batchelor, and W. Elwasif, address this issue with a new computational framework that separates the underlying physics of a problem from the specific hardware it runs on. This innovative approach, built around a graph-based representation of equations, automatically generates optimised code for CPUs, Apple GPUs, and conventional GPUs, offering a versatile solution for a range of applications. The team demonstrates the framework’s potential by applying it to radio frequency ray tracing, a crucial component in the development of fusion energy, paving the way for more portable and efficient scientific software.

Different computer vendors currently limit the portability of scientific code, as a program written for one type of graphics processor will not operate on another. Frameworks capable of abstracting the underlying physics from the specific computer hardware are therefore needed to fully exploit current and future computing capabilities. This research presents the development of a graph computation framework which compiles physics equations into optimised kernel code for the target hardware.

Graph Compilation for Physics Equation Optimisation

This work introduces a system for physics simulations based on representing equations as a graph data structure. This graph is then used to compile the equations into optimised code for various hardware platforms, addressing the challenge of manually optimising code for different processors. The graph structure enables automatic calculation of gradients, essential for many simulations, and allows for symbolic simplification of equations, potentially removing redundant terms and exploiting symmetries. This approach offers hardware portability, reducing the need for platform-specific optimisation, and promotes modularity, making it easier to add new physics models or modify existing ones.

The system represents equations and operations as nodes within the graph, then traverses and transforms this graph into executable code for the target hardware. The work draws on principles from compiler design, such as abstract syntax trees and intermediate representations, and leverages the LLVM framework for code generation and optimisation. Modern programming languages like Fortran, C++, and Python are used in the implementation. This research represents a shift from manually optimised code to a more automated, graph-based compilation approach, increasing the productivity of physicists and computational scientists by reducing the effort required to develop and optimise simulations. The graph-based approach could enable new types of physics simulations currently too computationally expensive or difficult to implement. In essence, the authors propose a more flexible, automated, and scalable way to build and run physics simulations, potentially accelerating scientific discovery in a variety of fields.

Unified Graph Computation for Diverse Processors

Researchers have developed a new graph computation framework designed to overcome compatibility issues between different types of computer processors, particularly the increasing reliance on Graphics Processing Units (GPUs) for high-performance computing. This framework addresses the challenge that code written for one vendor’s GPU often cannot operate on another, hindering the full exploitation of modern supercomputer architectures. The framework operates by translating physics equations into an intermediate representation, then optimising this representation before generating machine-readable code tailored to the specific hardware.

A unique feature is the ability to render these expression trees in LaTeX format, significantly aiding in debugging and verification. The system incorporates a node caching mechanism, identifying and reusing identical sub-graphs to minimise computational overhead and improve performance, and automatically simplifies expressions during construction to further enhance speed. This approach differs from existing machine learning frameworks, such as TensorFlow and PyTorch, which often struggle with performance and memory usage when applied to scientific computing, and are challenging to integrate into existing C/C++/Fortran-based codes. By focusing on symbolic representation and optimisation, the framework allows for efficient computation of even high-order derivatives, crucial for many scientific applications. The developers demonstrate the utility of this framework by applying it to Radio Frequency ray tracing problems relevant to fusion energy research, paving the way for broader adoption in other computationally intensive fields.

Graph Compilation for Accelerated Physics Simulations

This work demonstrates a graph-based computational framework capable of compiling physics equations into optimised code for diverse hardware, including CPUs and GPUs from multiple vendors. Representing equations as graphs enables both automatic differentiation and symbolic manipulation, offering significant advantages over traditional coding approaches. Specifically, the framework allows for the easy implementation of new physics models and automatically generates necessary gradient terms, streamlining the development process and reducing the potential for errors. Benchmarking using a ray tracing problem relevant to fusion energy research shows that the generated kernels efficiently scale across multiple devices, demonstrating the framework’s ability to leverage parallel processing capabilities. While the examples focus on wave and particle simulations in plasmas, the authors emphasise the broad applicability of this approach to any physics problem involving a large number of independent calculations.