Kris Rowe, a computational scientist at the Argonne Leadership Computing Facility (ALCF), is leading efforts to optimize the computational dynamics solver NekRS for deployment on Aurora, the ALCF’s forthcoming 2-exaflop system. The ALCF is a U.S. Department of Energy (DOE) Office of Science User Facility located at Argonne National Laboratory.

Researchers use NekRS for mission-critical DOE problems like the simulation of coolant flow inside small modular reactors, as the Exascale Computing Project’s (ECP) ExaSMR effort aims to accomplish.

Like a number of computational science engineering applications, NekRS relies on OCCA, an open-source, vendor-neutral framework and library for parallel programming on diverse architectures. OCCA enables the transparent generation of raw backend code, as well as just-in-time compilation that permits SYCL/DPC++-compatible kernel-caching.

NekRS makes heavy use of the OCCA framework to achieve performance portability, leveraging the OCCA runtime to manage GPU memory, enqueue computations on GPUs, and effect just-in-time compilation of math kernels. The runtime translates human-readable source code into machine-executable binary while the application is still running. Furthermore, all NekRS math kernels are implemented using OCCA Kernel Language (OKL).

Development of the DPC++ backend for the OCCA portability framework included implementation of a source-to-source compiler for translating OKL to DPC++.

Version changes to the Nek application

NekRS, which began as an offshoot of libParanumal (an experimental set of finite element flow solvers developed at Virginia Tech), is composed, in contrast to its Fortran-based predecessor, predominantly of C++ code.

Beyond providing a user interface similar to that of computational dynamics solver Nek5000, the construction of NekRS represents a major departure from its predecessor. However, the NekRS developers have engineered their code to include significant backwards compatibility, enabling users to port their existing case files with minimal effort.

Nek5000 was originally designed for strong scaling—that is, to minimize the time-to-solution for a given problem. NekRS is likewise designed for strong scaling, but in terms of a much greater magnitude: exascale computational power will help carry out far larger multiscale and multiphysics simulations than were possible with previous generations of high-performance computing (HPC) systems.

The Aurora-optimized version of NekRS is expected to differ only slightly from other iterations of the NekRS source code. Because differences in the generation of compiler code and between GPU microarchitectures can significantly impact performance, code customization is concentrated among the most performance-critical compute kernels. The ALCF maintains a fork of NekRS (a copy of the source code) on its GitHub page to enable enhanced user accessibility and for further development and hardware-specific optimization.

Optimizing NekRS for Aurora

The NekRS optimization strategy, as with many HPC science codes, is to climb a ladder of increasingly complex problems. That is, the developers begin with small issues, often fit for a single compute node or server, before working to resolve large-scale difficulties.

A single science code can exhibit a wide range of performance characteristics depending on how it is being used or on the problem to which it is being applied. Due to such variations in behavior, it is important that performance be analyzed using sample inputs representative of the challenges that the code will be used to solve.

The ECP-supported ExaSMR effort provides an illustrative example. ExaSMR aims to model the entire core of a small module reactor, which as a simulation is far too large for a single node to generate. The NekRS developers, however, provided Rowe with a simulation of a single reactor rod as a representative test problem to execute on a single node.

This test case, in conjunction with the Intel performance tools VTune and Advisor, enabled Rowe to identify which stages of the NekRS solver process require the most time. Indeed, as is the case with many other science codes, a small number of math kernels are responsible for most of each GPU’s compute time. Identifying information of this kind—which is capable of reducing the amount of code being targeted from tens of thousands of lines to just a few hundred—is critical for accelerating optimization efforts.

Subsequent to identifying critical kernels, roofline analysis (as performed with Intel Advisor) calculates the difference between actual measured performance and the hardware’s theoretical peak. Moreover, the insights gleaned via roofline analysis can help determine which optimization techniques are most appropriate for a given situation, or even if an entirely new algorithm should be considered.

Such considerations of new algorithms have driven what Rowe considers to be some of his most interesting recent work. In one instance the researchers added an algorithm that exploits an algebraic structure known as the tensor product. The tensor product of the discretization methods used by NekRS permits a given kernel to be computed in multiple equivalent ways. For example, data for a single mesh element (subdomain) can be represented either as a large vector or as a cube of numbers. Both approaches have strengths and weaknesses and offer different advantages and disadvantages depending on the particular kernel, problem size, and hardware characteristics.

Once NekRS has been optimized for a single Aurora compute node, Rowe will use a sequence of larger test problems to optimize multi-node performance. Particular attention will be paid to MPI (message-passing interface) communication patterns exercised by the application. The final step will be to analyze the performance of NekRS as applied to the full ExaSMR reactor core problem at scale and running on the complete Aurora system.

Keeping OCCA up to date

The OCCA framework is compatible with kernels written in OKL in addition to those written in backend-specific code, such as SYCL. The primary goal in optimizing NekRS kernels is to improve performance to the maximum extent permitted by OKL kernels; OKL kernels were selected because they are immediately portable to other vendor architectures.

However, in certain cases additional performance improvements can be realized only by employing newer features that the SYCL programming model enables and for which no equivalents exist in OCCA. Subsequently, developers survey the programming models of alternative OCCA backends like CUDA, HIP, and OpenMP, for analogous features. In instances where such analogues exist, the feature can be introduced into the OCCA framework in a unified manner.

This is illustrative of one of OCCA’s strengths—its relative agility enables developers to it incorporate bleeding-edge developments as they occur.

To gauge the needs of OCCA users and better coordinate development, Rowe has led the effort to establish the OCCA Technical Advisory Forum (TAF). Holding open, biweekly meetings, the OCCA TAF brings together stakeholders from the DOE, U.S. universities, Intel, AMD, Shell, and CCG, among others.

=================

The Best Practices for GPU Code Development series highlights researchers’ efforts to optimize codes to run efficiently on the ALCF’s Aurora exascale supercomputer.