Toward Performance-Portable PETSc for GPU-based Exascale Systems

The first challenge for GPUs’ general use in scientific computing is the portability of application code across different hardware and software stacks. Different vendors support different programming models and provide different mathematical libraries (with different APIs) and even different synchronization models. AMD is revising its HIP and ROCm library support; how closely it will mimic the CUDA (cuBLAS, cuSPARSE) standards is unclear. On the other hand, NVIDIA continues to develop the cuSPARSE API, deprecating and refactoring many common usage patterns for sparse matrix computations. Kokkos is designed to guarantee application codes’ portability, but it does not provide distributed-memory MPI support and requires the user to work in its particular C++ programming style. Likewise, OpenCL has portable performance as an overarching goal, but implementation quality has been uneven, and performance has not met the standard needed for many application codes [10]. OpenMP offload poses unique difficulties for libraries. Since it is compiler based and the data structures and routines are opaque, it may require specific code for each OpenMP implementation.

In addition to application porting, high-throughput systems such as GPUs require developing and implementing new solver algorithms—in particular, algorithms that exploit high levels of data-parallel concurrency with low levels of memory bandwidth delivered in a data-parallel fashion. This approach generally involves algorithms with higher arithmetic intensity than most traditional simulations require. This challenge is beyond the scope of this paper; a robust research program developing advanced algorithms for GPU-based systems is being undertaken by the PETSc team and their collaborators.

The current focus of the PETSc library work for GPUs is high computational throughput for large simulations, keeping all the GPU compute units as busy as possible all of the time. In order to achieve this, each core must have an uninterrupted stream of instructions and a high-bandwidth stream of data within the constraints of the hardware and low-level software stack. The difficulty arises from the complex control flows and data flows, as indicated in Figure 3. With distributed-memory GPU computing, two levels of instruction flow exist:

• •

  high-level instructions flow from the CPU memory to the GPU stream queue and the GPU memory controller, and

• •

  kernel code flows from the GPU memory to the GPU compute cores.

The two levels of data flow are

• •

  between GPU memories through a combination of networks, switches, and RDMA transfers and

• •

  between the GPU memory and the GPU compute cores.

In this work, we are concerned only with the high-level instruction and data flow and assume that the low-level is suitably handled within the computational kernels and the hardware.

With a high-throughput computation on the GPU and the CPU cores, the same basic rule applies. Seamless high-level data control must exist between the CPU memory and the GPU memory, and the high-level instruction flow for the CPUs and GPUs must be coordinated and synchronized to ensure that neither is interrupted.¹¹1The OLCF Summit IBM/NVIDIA system includes an additional layer, whereby data control for transfers directly from the CPU caches to the GPU memory, but this will be ignored. Since the exascale nodes have many more CPU cores than GPUs, a mismatch exists between the two for both high-level instruction management and data flows; this is called the oversubscription problem.

For both the pure GPU throughput problem and the combined CPU and GPU case, one must understand in a little more detail how the CPU controls the GPU operation. The CPU sends two types of high-level instructions to the stream queue on the GPU:

• •

  Kernel launch instructions that control the next set of low-level instructions the GPU will run on its compute cores

• •

  Data transfer instructions that control the movement of data

High-level instructions are issued in the order they arrive in the queue. Should the queue become empty, the GPU computations and memory transfers will cease. One of the most challenging problems currently faced by large-scale distributed-memory GPU-based computing is preventing this queue from becoming empty. It is crucial that more entries be added to the queue without requiring the queue to be emptied first. The process of adding more entries while the queue is not empty is called pipelining.

Two other techniques besides pipelining can keep the GPU compute cores busy.

One approach is kernel fusion, which combines multiple kernels into a single kernel to increase the time spent in kernel computations while reducing the number of kernel launches, with the additional benefit of possibly increasing the arithmetic intensity of the computations. However, this optimization often results in code that is harder to read and more difficult to maintain. For thirty years, the general model for library development has been to write short single-concern functions that can be easily combined to promote code reuse. PETSc uses this style where each concern is its class method.

A second approach is using CUDA graphs, which allow the user to describe a series of operations (including CUDA kernels) and their dependencies as a graph. These graphs can be instantiated once and then executed many times without involving the CPU to amortize the high instantiation cost. This approach is similar to pipelining except that it offers more runtime flexibility in determining the next operation within the GPU based on the state of the computation, whereas with pipelining the order of operations is set in advance by the CPU.

GPU-aware MPI was introduced to hide latency and ease development of multi-node GPU code by allowing MPI calls to accept GPU memory address pointers, thereby allowing MPI implementations to overlap GPU-CPU copies with communication rendezvous, chunking, and completion, as well as on-node optimizations. For example, a CUDA-aware MPI can use NVIDIA GPUDirect point-to-point to copy data between two GPUs within a compute node or use NVIDIA GPUDirect RDMA to access the remote GPU memory across compute nodes without requiring CPU involvement.

Unfortunately, at the time of writing, even GPU-aware MPI implementations cannot pipeline MPI calls and kernel launches because the MPI API does not utilize the stream queue. Figure 3 shows that the GPU-aware MPI instructions pass directly to the memory controls (the black dashed lines) and do not enter the stream queue. Thus, the stream queue must be empty for every MPI call. Such expensive GPU synchronization for every MPI call involving GPU memory embodies the mismatch between MPI and GPU. This problem is analyzed in more detail in [11].

In contrast, NVIDIA provides both NCCL [12], a library of multi-GPU collective communication primitives, and NVSHMEM [13], an implementation of the OpenSHMEM [14] communications model for clusters of NVIDIA GPUs. Unlike MPI, these libraries can initiate communication from the CPU on streams to synchronize between computation and communication taking place on the GPU. NVSHMEM also supports GPU-initiated communication.

Exascale systems will have many more CPU cores than GPUs, and applications may use both the CPU and the GPU for their computations. Libraries and applications must manage the mismatch between the number of cores and GPUs. Two main alternatives can help manage this challenge.

Oversubscription occurs when multiple CPU processes use the same GPU simultaneously. On high-end NVIDIA systems, the Multi-Process Service (MPS), a runtime system enabling multiprocess CUDA applications to transparently use the Hyper-Q hardware-managed work queues on the GPU, allows funneling the work from multiple CUDA contexts into a single context. In some cases, this can improve resources utilization, but it can also reduce GPU efficiency. On Summit, for instance, we have observed reductions between 5% and 20% when sharing GPUs between ranks. Many types of overhead can contribute to this reduced efficiency, but we group them all as GPU oversubscription overhead. For example, in a standard strong scaling regime, increasing the number of processes sharing the same GPU while holding the computational load per physical GPU constant may result in more kernel launches associated with smaller chunks of data, ultimately degrading performance. It is unclear what this overhead will be on future systems.

A variant of this approach maintains two communicators, one for all the MPI processes and one with a single process for each physical GPU. While the GPUs can be used at all times, control of the GPUs may be dropped down to the smaller communicator during intense phases of GPU utilization. The individual ranks’ data in the GPU memory can be shared with the GPU’s single-rank process via interprocess communication (e.g., cudaIpc) calls. Alternatively, all the extra CPU cores associated with the GPU controlling rank can share a small communicator over which the needed communication is performed, either by using MPI shared-memory windows or via scatter/gather operations.

A different approach is to always use one MPI rank per GPU and use, for example, OpenMP shared-memory programming to employ additional cores controlled by the single MPI rank that also controls the GPU. When individual library or application components use pure MPI and others use hybrid MPI-OpenMP, a decrease in active ranks is also necessary.

Undersubscription occurs when a large computation has short phases that are best computed on a smaller set of the computational cores, in other words, fewer GPUs and CPUs; this is commonly encountered in coarse level solves of a multigrid algorithm. A complete library system must manage the reduction in computational cores used and the efficient movement of data required to utilize the fewer resources and the transition back to the full system.

CPU-GPU communications can limit the application’s overall performance. The best approach to deal with the communication costs is to reduce the amount of communication needed by moving computation to the GPU. Communication can also overlap with computation on the GPU by performing the computations in different streams.

On NVIDIA and AMD GPUs, the communication time can be reduced by using pinned memory. Pinned memory is CPU memory whose pages are locked to allow the RDMA system to employ the full bandwidth of the CPU-GPU interconnect; we observed bandwidth gains up to a factor of 4.

When running on mixed CPU and NVIDIA GPU systems, there are at least four types of memory: regular memory on the CPU, pinned memory, unified shared memory, and GPU memory. The first is usable only on the CPU; the second may be usable on some GPUs, but not directly via the host pointer; the third is usable on both the CPU and GPU; and the fourth only on the GPU. Pointers for all of them are allocated, managed, and freed by the CPU process. AMD’s HIP API and Intel’s discrete GPUs add to the diversity of memory models that library and application developers must support to achieve optimal performance across all systems.

When multiple types of memory are used, especially by different libraries, there must be a system to ensure that the appropriate version of a function is called for the given memory type. CUDA provides routines to determine whether the memory is GPU or CPU, but these calls are expensive (about 0.3 $\mu$s on OLCF Summit) and should be avoided when possible. OpenMPI and MPICH implementations that are GPU-aware call these routines repeatedly and thus pay this latency cost for each MPI call involving memory pointers. With unified memory, the operating system automatically manages transferring pages of unified memory between the CPU and GPU. However, this approach means that callers of libraries using unified memory must ensure they use the correct memory type in application code; otherwise, expensive data migration will be performed in the background.

Multiple streams are currently used by NVIDIA and AMD systems to synchronize different computation and computation phases directly on the GPU by providing multiple stream queues. This synchronization is much faster than synchronizations done on the CPU. Care must be taken not to launch too many overlapping streams, however, since they can slow the rate of the computation. Thus, it is best that a single entity be aware of all the streams; this requires extra effort when different libraries are all using streams. Moreover, if different libraries produce computations that require synchronizations between the library calls, those libraries must share common streams.

GPU systems often come with more single-precision than double-precision floating-point units. Even bandwidth-limited computations with single precision can be faster because they require half of the double-precision computations’ memory bandwidth. Sparse matrix operations do not usually see dramatic improvements, however, because the integer indices remain the same size.

With the compressed sparse row matrix format, the amount of data that needs to be transferred from the matrix is $64nz+32nz$ for double-precision computations with 32-bit column indices and $32nz+32nz$ for single-precision computations, a $2/3$ ratio. Further reductions in the factor require compression of the column indices or using a different storage format; see, for example, [15].

Recent GPUs also contain tensor computing units that use 16-bit floating-point operations. These are specialized and do not provide general-purpose 16-bit floating-point computations, but they have significant computing power; they can be used effectively for dense [16] or even sparse matrix computations [17]. See also a recent survey on multiprecision computations for GPUs [18].

Each PETSc back-end comprises two parts: the calls to numerical libraries that provide the per GPU algorithmic implementation and the glue code that connects these calls to the front-end code. The numerical libraries are cuBLAS and cuSPARSE for the CUDA back-end, ROCm for HIP, ViennaCL [21] for OpenCL (and potentially also for CUDA and HIP), MKL for SYCL, and Kokkos Kernels for Kokkos. Different PETSc back-end subclasses can be used in the same application.

Our highest priority at present is to prepare for the upcoming exascale systems. We already have a working Kokkos Kernels back-end. We are completing the HIP back-end to exploit the AMD GPUs planned for the upcoming Frontier system at Oak Ridge National Laboratory. The basic configuration for HIP is in place, as well as the initial vector operations. Development is in progress and being tested on early-access testbeds for Frontier. We have also begun implementation for a SYCL and MKL back-end for use with the Intel GPUs planned for the Aurora system at Argonne National Laboratory.

As our computation kernels’ development continues, we are also abstracting the fundamental types, their initialization, and the various libraries’ interfaces to reduce code duplication. The PETSc team will work with the Intel and AMD compiler groups to determine the information needed from the OpenMP offload compiler to share its data with the PETSc back-end.

Achieving high utilization requires a rethinking of outer-level algorithms to take advantage of operations with higher arithmetic intensity. Compact, dense quasi-Newton representations present an avenue for increasing arithmetic intensity and increasing GPU utilization for outer-level algorithms. Compared with conventional limited-memory matrix-free implementations, these formulations require fewer kernel launches, avoid dot products, and compute the approximate Jacobian action via a sequence of matrix-vector products constructed from accumulated quasi-Newton updates. Initial support for solving Krylov methods with multiple right-hand sides has been added to PETSc [22]. While having higher arithmetic intensity, these methods generate larger Krylov subspaces and typically converge in fewer iterations, and they may thus provide algorithmic speedup for calculating eigenvalues and for multiobjective adjoint computations. These block methods significantly decrease the number of sequentially launched kernels, providing outer loop fusion that reduces CPU-GPU latencies and calls; see challenge F2. Even with the multiple right-hand side optimizations, NVIDIA’s cuSPARSE on a V100 with extremely sparse matrix-vector multiplication achieves only 4% of peak, thus indicating the desirability of implementing algorithms that fundamentally change the performance for such simulations. Outer loop fusion will also be applied for other PETSc operations such as the multiple dot products operations needed in the Krylov methods (replacing a collection of BLAS 1 operations with a single BLAS 2 operation) to reduce the number of kernel launches and get more data reuse on the GPU and higher utilization.

A cross-utilization of the CPU and GPUs can occur even within what is conceptually a single computation. Consider the computation of matrix elements for a discretization scheme and their insertion into a matrix data structure. Within PETSc, value insertion is performed in a loop with the general routine MatSetValues (usage depicted in Listing 3), which efficiently adjusts the row sizes on the fly as values are inserted and takes care of communicating off-processor values, if needed.

However, the many repeated calls to MatSetValues run on the CPU and may cause a bottleneck when the assembling process runs on a GPU. To address this issue, we have developed a variant of the function callable from kernels that will guarantee fast GPU resident assembly. In a somewhat different scenario, common among complicated applications utilizing PETSc only for its robust solver functionality, the matrix values can be already available from the application itself; for such cases, we have recently developed a new API that accepts a general matrix description in a coordinate list (COO) format and directly assembles on the GPU utilizing efficient kernels from the Thrust library that is part of the CUDA Toolkit. For numerical results, see Section 5.5.

PETSc will use the simultaneous kernel launch capability of NVSHMEM to begin its parallel solvers, thus making them more tightly integrated and remaining in step during the solve process. This is a relatively straightforward generalization of the current solver launches, and it will be opaque to the users.

The PETSc back-end model is designed to support applications and libraries using a combination of CPUs and GPUs for computations. Each back-end class implementation provides two implementations, for host and device, and contains the code to manage data transfers. This strategy allows the users of PETSc to prescribe the next group of computations to take place by merely setting a flag on the object. With this approach, running, for example, the coarser levels of multigrid on the CPU and the finer levels on the GPU is easy. Unnecessary GPU-to-CPU communication is avoided with this approach, and any needed communication is handled internally by the object. The associated message passing required by the solver is managed automatically by PetscSF; see Response A2.

Since MPI does not use streams, inner products and ghost point updates require a CPU synchronization at these steps; one cannot, for example, pipeline an entire iteration of a solver together. We will provide support for NVSHMEM to allow the communication to utilize the stream queue; see Response A2. This is an important addition to PETSc to eliminate the empty streams queue problem.

PetscSF is the PETSc communication module, which has a simple but powerful interface to simplify communication code on both CPUs and GPUs [23]. It manages proper synchronizations needed by GPU-aware MPI, and has a series of optimizations to improve communication performance. PetscSF is incrementally replacing all direct use of MPI in PETSc. This will allow PETSc to utilize the features of NVSHMEM that are required to eliminate the problems of Challenge A2.

PetscSF uses the abstraction of star forests to represent communication patterns. A star is a simple tree consisting of one root vertex connected to zero or more leaves. A star forest is a disjoint union of stars. A PetscSF is created collectively by specifying, for each leaf on the current rank, the rank and offset of the corresponding root, shown in Figure 4. PetscSF then analyzes the graph and derives a communication pattern using persistent nonblocking MPI send and receive (default), other collectives, one-sided, or neighborhood collectives. PetscSF provides APIs such as PetscSF{Bcast,Reduce}Begin/End. The former broadcasts root values to leaves, and the latter reduces leaf values into roots with an MPI_Op. The Begin/End split-phase design allows users to insert computations in between, in order to overlap with communication.

Root and leaf indices are typically not contiguous per rank, so that the library has to pack and unpack the data for MPI sends and receives. If the data is in GPU memory, these pack/unpack routines must be implemented in GPU kernels. PetscSF APIs are raw pointer-based, for example,

where unit defines the data type of the vertices in the graph. PetscSF originally used cudaPointerGetAttributes to infer root/leaf memory spaces, but that turned out to slow down the operations, so we now keep track of the memory types that are passed to the PetscSF for communication on PETSc vectors. Currently, PetscSF supports CUDA and Kokkos; support for HIP and SYCL will be added soon. We will also provide a back-end that uses NVSHMEM or NCCL.

PetscSF prefers GPU-aware MPI, but it also supports non-GPU-aware MPI. Here we discuss some of the details of our GPU-aware support. There are multiple choices to build synchronization models, as shown in Figure 5, which uses CUDA as an example. In Figure 5(a), we assume that leafdata is produced by a kernel on a stream unknown to PetscSF. Therefore the sender has to call CUDADeviceSynchronize to synchronize with all the computations on the GPU before the pack kernel in order to make sure the leafdata is ready to pack. After the pack, the implementation calls CUDAStreamSynchronize(s1) to synchronize stream s1 and make sure the packed data in sbuf is ready for MPI_Isend. On the receiver side, when the MPI_Waitall returns, the received data is guaranteed to be in the receive buffer rbuf. But we might have to call the unpack kernel, for example, on stream s2, to unpack the data from rbuf to rootdata and call CUDAStreamSynchronize(s2) to synchronize unpack so that rootdata can be consumed by kernels on any stream.

Manipulating multiple streams is not an easy task. Like most codes, the PETSc default is to use only the default stream. A simplified synchronization model is shown in Figure 5(b), where we launch pack/unpack on the default stream and omit the CUDADeviceSynchronize before pack and CUDAStream
Synchronize(s2) after unpack. But obviously, we still need CUDAStreamSynchronize(NULL) before MPI_Isend. A solution for the MPI-GPU mismatch problem is to expose the CUDA streams to the MPI API; doing so will allow MPI function calls on GPU memory to behave more like kernels where the user selects the appropriate streams to allow kernel pipelining. The MPI Forum has begun an Accelerator working group, and the PETSc team is an eager customer willing to quickly work with the prototypes the working group develops.

MPI processes often need to communicate with themselves as well. Traditionally, programmers chose not to distinguish local (i.e., self-to-self) and remote (i.e., self-to-others) communications and let MPI handle the distinction, resulting in a simpler and more uniform user code. Distinguishing these types of communications with GPU computation is essential, however. At least two benefits accrue: (1) we can directly scatter source data to its destination without intermediate send or receive buffers for local communication, and (2) local communication is done by a GPU kernel (doing a memory copy with index indirection). The kernel can be perfectly executed asynchronously with remote MPI communication, overlapping local and remote communications.

PetscSF performs index analysis at the setup phase to provide optimizations to lower the cost on GPUs further. If leaf or root indices are contiguous, we do not need pack/unpack and intermediate buffers: data can be sent directly to save kernel launch and execution time. For noncontiguous indices, the indices might have specific patterns that can be taken advantage of to reduce the packing or unpacking time. Currently, PetscSF can detect whether a set of indices are for points in a rectangular subdomain. For example, in a rectangular 3D domain, the ghost regions (faces) and interior region are all such qualified subdomains. Packing entries in such subdomains only need parameters describing the subdomains, instead of all entry indices. Also, in PetscSFReduce, multiple leaves might be reduced into the same root. If so, the unpack kernel needs to take care of the data race between threads. Instead of blindly using atomic operations all the time, PetscSF leverages index analysis, and it uses nonatomic instructions when no duplicated indices are found.

We are focusing on the most critical case for PETSc exascale application codes: simulations that perform phases of the computations jointly on the CPUs and GPUs, with other phases exclusively on the GPUs all using a purely MPI model. This case will use the approach outlined above, using a large MPI communicator for all MPI cores with shared process usage of the GPUs combined with a smaller communicator just for the GPUs with the highest performance single process usage of the GPUs. Since the bulk of the computation will take place without sharing the GPU, this approach tolerates the overhead that arises from sharing mentioned in Challenge A3. Shared memory on the GPUs will allow joint ownership of the data used in the different phases of the computation; communication for the data will be managed by PetscSF, transparent to the application code. Other use cases can easily be added as needed by the application teams. PetscSF will also provide transparent support for the undersubscription problem faced, for example, by multigrid; see Section 5.3.

PETSc, by default, uses pinned memory for all data that may be copied to or from the GPU memory. As mentioned above, PETSc uses the lazy-mirror model to reduce memory copies. At the lowest level of the PETSc software stack, data access is controlled via get operations such as VecGetArrayRead, VecGetArrayWrite, and VecGetArray, and the corresponding restoring partners, for example VecRestoreArray. Careful use of these read and write accessors avoids unnecessary copies; the reader does not invalidate the data in the other memory, while the writer does not require copying, and it merely flags the other memory as invalid. SYCL also uses this read/write model. Once a computational phase, even with many stages, begins on the GPU, the data remains automatically in the GPU.

PETSc does not yet provide a mechanism to prefetch memory of an object, but that can be added by using the tracking of valid data and the get/restore model [8]. Prefetch would use streams to efficiently pipeline operations that depend on the required data.

We are testing multiple related approaches for managing the life cycle of allocated memory; its allocation, use, and release. Support for malloc with normal, pinned, and unified memory can be managed by marking the memory, either with flags carried in objects that hold the memory, within a header to the actual memory, or with a separate small allocation that contains the identifier and the actual pointer. This last approach also supports GPU memory since the identifier header is accessible on the CPU, while the GPU memory is not.

PETSc will incorporate streams into PETSc objects, both at the data level (vectors and matrices) and at the solvers level (linear, nonlinear, and time integrators). With this capability, entire iterations of solves can be launched on different steams so their computations are interlaced on the GPUs. This provides a straightforward way to effectively use GPUs to solve many different-size problems simultaneously.

Simple use of single-precision floating point on GPUs within a library is straightforward at the software level. In PETSc, each class will carry additional information: the precision of the data stored and the requested computations’ precision. Initially, PETSc will support multiprecision computation by utilizing single precision only on the GPU; this approach reduces the code complexity since it does not affect the CPU user APIs. Also provided will be the capability of converting the formats directly on the GPU when different phases of more extensive calculation can benefit from different precision. The BLIS [24] package has an excellent system for managing the control of the various mixed-precision computations, and PETSc may adopt a similar approach. When data is transferred from or to the GPUs, it will be converted between precisions on the fly. One will control the precision by a low-level object such as a vector or by a higher-level operation such as an entire linear solve by setting the flags appropriately on the controlling data object. Utilization of the tensor floating-point unit for general sparse matrix computations in PETSc is currently out of scope; but if other GPU numerical libraries use them, PETSc can benefit by providing an efficient interface to these packages. As application teams complete their ports to use GPUs, the PETSc team is ready to work with them on their exact needs to determine which parts of the computations may benefit from using lower-precision computations. At this point, the determination of the optimal combinations will be mostly experimental; because of the PETSc back-end separation from user code, the studies can all be done at runtime and would not require manually recompiling code for different precisions.

Additionally, the communication infrastructure must be precision-aware. We will add support for PetscSF to reduce double-precision input to single-precision output, for example, for the overlapping Schwarz method.