H2Opus: A distributed-memory multi-GPU software package for non-local operators

The representation of an $\mathcal{H}^{2}$ matrix exploits two kinds of hierarchies: a hierarchy of blocks at different levels of granularity and a hierarchy of bases that is used to represent the data in the blocks. At every level of the hierarchy the low rank blocks essentially form a block sparse matrix. In addition, there is a block sparse matrix that stores the dense blocks of the matrix that are not amenable to low rank representations.

Hierarchy of Blocks. Let $\mathcal{I}$ and $\mathcal{J}$ be the sets of indices of the rows and columns of a matrix, and $T_{\mathcal{I}}$ and $T_{\mathcal{J}}$ be hierarchical clusterings of these index sets, respectively. All blocks within the matrix can then be defined by cluster pairs $(t,s)\in T_{\mathcal{I}\times\mathcal{J}}$. Define an admissible block as a block that is either small enough to be stored in dense form of size $m\times m$, with $m$ denoting the so-called leaf size, or can be well approximated by a low rank matrix. A matrix $A$ in the $\mathcal{H}$-variant of hierarchical matrices can be partitioned into admissible blocks of various sizes that are organized hierarchically as the leaves of a matrix tree. A low rank matrix $A^{l}_{ts}$ on level $l$ of the tree and defined by the cluster pair $(t,s)$ is represented as the outer product of two tall matrices $B^{l}_{ts}$ and $C^{l}_{ts}$ of rank $k^{l}$: $A^{l}_{ts}=B^{l}_{ts}{C^{l}_{ts}}^{T}$.

The left panel of Figure 2(a) shows the leaves of this matrix tree which define a complete partitioning of the matrix into blocks, with dense leaves colored red and low rank leaves colored green. Figure 2(b) shows the various levels of the tree, where inner nodes are in cyan. The structure of this matrix tree depends on the application. For example, in problems involving a spatial discretization, we can leverage data from the particles or nodes generated by the discretization to generate it with the aid of a geometric admissibility condition [31]. Other applications might rely on the available algebraic data or heuristics to determine which blocks of the matrix are admissible as low rank blocks.

Hierarchy of Bases. The $\mathcal{H}^{2}$ representation uses block low rank factorizations of the form $A^{l}_{ts}=U^{l}_{t}S^{l}_{ts}{V_{s}^{l}}^{T}$ for every $ts$ block at every level $l$, where $U^{l}_{t}$ and $V^{l}_{s}$ are bases for the block rows $t$ and block column $s$ of level $l$, respectively. $S^{l}_{ts}$ is a small $k_{t}^{l}\times k_{s}^{l}$ coupling matrix representing the interaction between the row and column cluster pairs. These coupling matrices are organized in the matrix tree $\cal{S}$ defined by $T_{\mathcal{I}\times\mathcal{J}}$. In our implementation, we currently use a fixed rank per level that we refer to as $k^{l}$, or even simply $k$ when the context is clear, which allows us to use higher-performing fixed-size batch linear algebra kernels when executing basis operations.

An additional hierarchy is introduced in the row and column basis trees, where a basis node is expressed in terms of the bases of its children. Basis nodes are only stored explicitly at the leaves of the tree, and inner nodes can be computed from their children using small $k^{l}\times k^{l-1}$ interlevel transfer matrices. For example, one can compute an inner node $U^{l-1}_{i}$ from its two children, $U^{l}_{i_{1}}$ and $U^{l}_{i_{2}}$, and their corresponding transfer matrices $E^{l}_{i_{1}}$ and $E^{l}_{i_{2}}$ as

The explicit bases at the leaves and the interlevel transfer matrices are organized in a basis tree. Figure 3 shows an example of the nested basis $\cal{U}$, with the implicitly represented inner nodes shaded.

Putting it all together, the $\mathcal{H}^{2}$-matrix representation may be succinctly written as

where $A_{de}$ is a block sparse matrix with dense blocks of size $m\times m$, shown as the red leaves at the finest level of the matrix tree in Figure 2(a). $\cal{S}$ is a matrix tree of $k^{l}\times k^{l}$ coupling matrices, whose leaves, which appear at different levels, are shown as the cyan blocks in the right panel of Figure 2(a), and $\cal{U}$ and $\cal{V}$ are the block row and block column basis trees, each consisting of $m\times k$ explicitly stored bases $U,V$ at the leaf level, and $k^{l}\times k^{l-1}$ interlevel transfer matrices $E$ and $F$, as shown in Figure 3 for basis tree $\cal{U}$. The triple-tree product notation $<{\cal U},{\cal S},{\cal V}^{T}\!>$ is defined as the assembly of all blocks $ts$ of all levels $l$ where every block may be computed as $U^{l}_{t}S^{l}_{ts}{V_{s}^{lT}}$.

The description of the parallel distribution of the matrix and the algorithms operating on it uses the notation in Table 1. We use a left subscript to refer to local data owned by a GPU, and a triple bar subscript to refer to data that has been marshaled to allow batched GPU kernels to be invoked on it.

Treating each level of the matrix tree of the hierarchical matrix as a block sparse matrix, we decompose the levels into block rows, with each block row stored on a single GPU, as illustrated in Figure 4. The basis trees are similarly decomposed by level, assigning the nodes corresponding to the stored block rows/columns to the same GPU. This decomposition allows much of the upsweep and downsweep to be carried out independently as described in Section 3 for matrix-vector multiplication and in Section 5 for matrix compression. Above a critical C-level, the decomposition stops and a single root GPU owns all top levels. The interlevel transfer operators at the roots of each of the local branch basis trees ${{}_{p}}{\mathcal{U}}$ are duplicated at the leaf level of ${{}_{r}}{\mathcal{U}}$ (see Figure 4) to allow upsweep and downsweep operations to begin and end at the C-level.

We construct the distributed matrix by first splitting the row cluster tree $T_{\mathcal{I}}$ into $P$ independent branches at level $l=\log P$ and a local branch ${{}_{p}}{\mathcal{U}}$ of the basis tree on GPU $p$ is generated for each $p\in\left\{0\ldots P-1\right\}$. The top $l$ levels of the basis tree are kept on a master process and could potentially be split further if $P$ becomes large enough. The root branch of the matrix tree ${{}_{r}}{\mathcal{S}}$ can be generated on the master process in any number of ways, such as general admissibility dual tree traversal [31]. Once this matrix tree is constructed, a list of the basis nodes $L_{p}$ from ${{}_{r}}{\mathcal{U}}$ can be extracted for each GPU $p$ corresponding to the inadmissible nodes of the block row $p$ at level $l-1$ of the matrix tree. As an example, consider the distribution on $P=8$ GPUs of the fourth level of the matrix tree in Figure 2(b). The column indices corresponding to the cyan nodes are extracted for GPU $p=1$, setting $L_{1}=[0,1,4]$ from the second block row. Once the list for each process has been compiled, they are scattered from the master process to all other processes.

Since the nodes $q$ in $L_{p}$ were generated from the inadmissible nodes of the matrix tree $S^{l-1}_{pq}$ that have to be subdivided further, the structure of the matrix tree ${{}_{p}}{\mathcal{S}}$ on each GPU can then be generated independently using multiple dual tree traversals of the root node of ${{}_{p}}{\mathcal{U}}$ with the nodes $q$ in $L_{p}$, where each traversal generates a local branch ${{}_{p}}{\mathcal{S}}_{q}$. Once the structure has been determined, the entries of the transfer and leaf matrices of the basis trees, together with the entries of the matrix tree can be populated independently on each GPU using established techniques [16].

The upsweep phase is illustrated in Figure 6 and summarized in Algorithm 2. It proceeds from the leaf level where the explicitly stored block row bases, each of size $m\times k$, are simply multiplied by $x$. At every higher level, $\widehat{x}^{l-1}$ can be computed from the children at level $l$ using the transfer operators [16]. For a parent node $s$ with children $s_{1}$ and $s_{2}$ we have

In the distributed setting, the upsweep on each GPU $p$ can proceed in parallel on each of the local branches ${{}_{p}}{\mathcal{V}}$ independently, using the kernel shown in Algorithm 1. Once the upsweep reaches the roots of each branch, the data from all root nodes of the ${{}_{p}}{\widehat{x}}$ trees are gathered on the master process to populate the leaf level of the root branch ${{}_{r}}{\mathcal{V}}$, allowing the upsweep to complete. The regular upsweep Algorithm 1 can be used for the branch upsweep as well as the root upsweep by simply omitting the first batched operation on the leaves of the root branch, since that step is replaced by the gathering of the data.

From an efficiency perspective, the primary challenge in the upsweep comes from the fact that the individual operations in the basis tree nodes involve small matrix operators that do not possess a sufficient amount of data parallelism for effective GPU utilization. We overcome this problem by marshaling the appropriate level data to allow batched GPU kernels to be invoked, allowing $\widehat{x}$ to be computed with only $O(\log N)$ batched kernel executions. Except for the top few levels of the root process $r$, these executions happen concurrently on the $P$ GPUs. The marshaling GPU kernel is shown for a step of the upsweep operation in Algorithm 3. The marshaling kernel essentially plays the role of a fast static scheduler for all operations performed on a given level $l$ of the tree. It prepares for the execution of the operations by efficiently gathering data from the basis tree. Marshaling involves no data movement and therefore has very little overhead and in practice consumes only a tiny fraction of total execution time.

The second phase of the operation builds a vector tree $\widehat{y}$, where each node $t$ in a level $l$ is the product of the block row $t$ of level $l$ of the coupling matrix tree with the corresponding nodes in $\widehat{x}$. This operation can be expressed as

where $b_{t}$ is the set of column indices of the matrix blocks in the block row $t$. This is a block sparse matrix-vector multiplication at every level, and is illustrated in the middle portion of Figure 5, where the block-row, multi-GPU decomposition of each level is also highlighted. The scalar version of this problem is a well-studied computational kernel that has many possible high quality solutions [30, 41, 15, 23] which may be adapted to the block-sparse version. While we could rely on vendor kernels, their relatively low performance and lack of support for non-uniform blocks necessitates a different approach. Our solution relies on a key result regarding hierarchical matrices which puts a bound on the sparsity constant $C_{sp}$, the maximum number of blocks in any block row $t$ at any level $l$ of the matrix tree. Such constant is bounded by a dimension-independent $O(1)$ value that only depends on the specific structure of the matrix and the admissibility criterion[28, 16].

During the construction of the matrix tree, we generate conflict-free batches of matrix products that can then be executed by a series of non-uniform batched matrix-vector and matrix-matrix kernels. A conflict-free ordering of the batch indices can be obtained by assigning a batch index based on the position within the block row or column that increases from left to right, allowing us to marshal all the batches in a single marshaling kernel. While there will potentially be some kernels that perform little work, the bounded sparsity constant guarantees that there will be few of them, and they will thus represent a small portion of the total runtime. This could be optimized further by running those small batches on a separate low priority stream and then combining the results with the main execution stream.

The boundedness of the sparsity constants $C_{sp}$ also guarantees that the block row local to a GPU will require input from $\widehat{x}$ nodes belonging to a limited number of remote GPUs: we therefore consider a standard approach in distributed memory sparse matrix vector products, and split the block row into diagonal and off-diagonal parts. The contributions of these two parts can be overlapped with the needed communication as described in Section 4.1. Once the needed parts of $\widehat{x}$ are assembled, the multiplication phase can proceed on each GPU independently on the coupling matrices of ${{}_{p}}{\mathcal{S}}$ using the treeMultiply routine of Algorithm 4. Processing the multiplication of the root branch ${{}_{r}}{\mathcal{S}}$ on the master process to produce the root branch ${{}_{r}}{\widehat{y}}$ finalizes the phase as shown in Algorithm 5. The selectiveGather routine communicates and assembles the needed portion of the $\widehat{x}$ tree from the remote branches to use in the local treeMultiply. This is described in Section 4.1.

The vector tree $\widehat{y}$ now contains, at every level $l$, a partial contribution to the output vector $y$ expressed in terms of the bases $U^{l}$, and it could be computed as $U_{t}^{l}\widehat{y}_{t}$ if the bases at level $l$ were explicitly available. However, since we only store interlevel transfer operators, we express $\widehat{y}_{t}^{l}$ in terms of bases of the children of node $t$ and accumulate them in a downsweep phase through the ${\cal U}$ tree from root to leaves. This partial accumulation takes the form

between a parent node at level $l-1$ and children $t_{1}$ and $t_{2}$ at the finer level $l$.

The downsweep procedure, pictured in the left part of Figure 5, is summarized in Algorithm 7. A downsweep through the root subtree ${{}_{r}}{\widehat{y}}$ is first performed. Once the leaf level of ${{}_{r}}{\widehat{y}}$ has been updated, it can be scattered from the master process to all other GPUs and added to the roots of local trees ${{}_{p}}{\widehat{y}}$ containing data from the previous phase of the multiplication. When the roots of the local branches ${{}_{p}}{\widehat{y}}$ on the $P$ GPUs have their proper accumulation, obtained from the leaves of the tree ${{}_{r}}{\widehat{y}}$, we can sweep down ${{}_{p}}{\widehat{y}}$ independently on all GPUs setting each node to ${{}_{p}}{\widehat{y}}_{t_{i}}^{l}={{}_{p}}{\widehat{y}}_{t_{i}}^{l}+{{}_{p}}{E_{i}^{l}}\ {{}_{p}}{\widehat{y}}_{t}^{l-1}$; the level $l$ at each step now also contains the partial sum of ${{}_{p}}{y}$ for all levels above $l$ expressed in the basis of $U^{l}$. The leaf level will then contain the complete sum which is finally expanded into ${{}_{p}}{y}$ through a multiplication by the explicitly stored leaf level bases. We follow the same approach used for the upsweep, where each level is also processed in parallel by first marshaling the operations and then executing using a batch matrix vector product. This leads us to Algorithm 6 for computing ${{}_{p}}{y}$, i.e., the distributed vector $y$. The downsweep marshaling algorithm is structurally similar to the one described in Algorithm 3 and is omitted for brevity.

We first discuss a communication-optimized parallel block sparse matrix-vector multiplication for a given level matrix $S^{l}$ of $A_{lr}$, the low rank portion of the hierarchical matrix; the same logic applies to $A_{de}$ and it is omitted for brevity.

The matrix tree is first split into two distinct trees: a diagonal matrix tree ${{}_{p}}{S^{l}}_{p}$ and an off-diagonal matrix tree ${{}_{p}}{S^{l}}_{\overline{p}}$. While the off-diagonal portion could simply be represented as a set of trees, with one tree for every interaction ${{}_{p}}{S^{l}}_{q}$ with a GPU $q$, it is far more efficient to have them all in a single flattened one to allow marshaling operations on the off-diagonal blocks to be completed with a single kernel call.

Once the trees are split, the list of basis nodes that interact with ${{}_{p}}{S^{l}}_{\overline{p}}$ can be generated by iterating through its $(t,s)$ pairs and determining all unique $s$ values. Given the boundedness of the sparsity constant $C_{sp}$, on a given level $l$, a GPU $p$ will receive data only from a few other GPUs; we thus determine the list of GPUs that need to send data to $p$, as well as the list of nodes that should be received, and store such information in a compressed-storage format as representatively shown in Figure 7. The data needed from process pid$[i]$, corresponding to global indices listed in nodes from nodes_ptr$[i]$ to nodes_ptr$[i+1]$, is communicated among GPUs during the setup phase. A list of unique entries of nodes represents the compressed storage for the offiagonal block.

The diagonal multiplication phase can proceed without any communication while the off-diagonal phase needs input from other GPUs. The nodes at a level $l$ of the local ${{}_{p}}{\widehat{x}}$ branch which are needed by other GPUs are packed into a single buffer using a marshaling kernel and a single batched copy kernel populates a send buffer $B_{s}^{l}$, which is then used to issue non-blocking sends to the neighboring GPUs. Non-blocking receives are issued to populate a receive buffer $B_{r}^{l}$, which can then be directly used for the off-diagonal level of the tree, since the column basis indices have already been updated to use the compressed node format.

Algorithm 8 summarizes the overall communication-optimized multiplication phase. The exchangeData routine executes the non-blocking MPI sends and receives using the compressed off-diagonal data that is marshaled using the marshalOffdiag routine and copied using a batch kernel.

While the communication volume has been reduced and a portion of the communication can be hidden by overlapping it with the diagonal multiplication, there are still a few challenges left that hold back performance. Assuming the lack of hardware support for more advanced memory transfer features such as GPUDirect RDMA, transferring the data from the GPU to the host adds some overhead to the execution, as well as synchronization points that impact GPU usage. Moreover, not all MPI distributions guarantee that communication can progress in the background without explicit intervention from the process.

The effectiveness of overlapping the diagonal multiplication phase with the needed communication depends on the structure of the tree, the capabilities of the communication network, and the compute capabilities of the GPU. For a tailored, fine-grained control of the communications involved in the algorithm at hand, we explicitly create communication threads that queue up asynchronous copies on separate streams to overlap the transfers with the processing of each level of the local branch upsweep. The transfers can then be carried out on the same thread without interrupting or forcing synchronizations with the main execution stream, allowing communication and computation to overlap. When the diagonal block kernels have been queued up on the stream, the communication thread can then join the main thread. The multiplication of the root branch on the master process can also be hidden, to some extent, by overlapping it with the main stream of work, and scheduling it on a low priority stream. This will allow the work to be completed during phases of low GPU utilization, such as the smaller top levels of the basis and matrix tree. Finally, by adding the dense block multiplication phase to a low priority stream, GPU utilization on all GPUs can be increased as well.

Figure 8 shows the effect of overlapping communication with computation on the timeline of the overall distributed matrix-vector multiplication. As expected, the gaps in the timeline due to communication are significantly smaller when it is overlapped with the computation.

Consider how a new basis for a block row $A_{i}^{q}$ at the finest level $q$ would be generated. $A$ here denotes only the low rank portions of the hierarchical matrix, since the dense blocks are not compressed. $A_{i}^{q}$ consists of $b$ low rank blocks expressed at level $q$ as $U_{i}^{q}S_{ij}^{q}V_{j}^{qT}$ with $j=j_{1}\cdots j_{b}$, and additional pieces representing the restriction of blocks from coarser levels to their “i” rows.

A new more efficient basis can be generated by computing the SVD of $U_{i}^{q}B_{i}^{qT}$, and using the left singular vectors as the new basis. This would however require an expensive SVD of the $O(N)$-sized $B_{i}^{q}$. In order to avoid it, we first compute the QR decomposition of $B_{i}^{q}$ and then perform the SVD on the small $R$ factor.

The new basis for level $q$, $\overline{U}^{q}_{i}$ is effectively a reweighing of the columns of the previous basis $U_{i}^{q}$.

The task of computing $R_{i}^{q}$ of the QR decomposition of $B_{i}^{q}$ can be done efficiently by exploiting the nestedness of the bases. Let us assume that the QR decomposition of $B_{i^{+}}^{q-1}$, the parent block $i^{+}$ at level $q-1$, is available as $Q_{i^{+}}^{q-1}R_{i^{+}}^{q-1}$. Then,

with $B_{i}^{q}$ conveniently expressible as:

Assuming the $V^{q}$ bases are orthogonal, the block diagonal matrix in Eq. 4 is orthogonal and the QR of $B_{i}^{q}$ simply reduces to the QR of the small stack at the end of Eq. 4 which involves only $b+1$ blocks, each being a small $k\times k$ coupling/transfer matrix. Since this QR uses the $R^{q-1}$ matrix from level $q-1$, the overall computation starts from the root and goes down the tree computing all the $R^{l}_{i}$ matrices for all levels in a downsweep traversal.

As with previous operations, all blocks at a given level can be processed in parallel, and importantly for the distributed setting, all the processes owning the subtrees below the C-level can proceed independently. Therefore, the computational pattern is identical to the distributed downsweep of the matrix-vector multiplication. Above the C-level, a single GPU is responsible for processing the top levels of the basis tree. The leaves of the top subtree, which hold the $R^{l}$ factors at level $l$, are then scattered to all GPUs and seed the roots of the individual subtrees, which continue the downseep independently all the way to the leaf level of the basis. The computational work at every level being, for every block row, a QR decomposition of the small stack at the end of Eq. 4. Batched QR operations [21] are used within every GPU to achieve high performance.

Once the new reweighed basis is generated for all levels, it can then be truncated to the desired accuracy to generate the optimal basis $U^{\prime q}_{i}$. The coupling blocks are then compressed by projecting them on the new truncated basis.

The truncation operation has to preserve the nestedness property. This can be accomplished by starting the truncation at the leaf level and sweeping up the basis tree. Given the reweighed basis $\overline{U}$ ($\overline{U}_{t}$ at the leaves and $\overline{E}_{t}$ at nodes in the tree), we seek a basis $U^{\prime}$ ($U^{\prime}_{t}$ at leaves and $E^{\prime}_{t}$ at nodes in the tree) that spans the same subspace as $U$ to the desired accuracy.

At the leaf level, this may be done through an SVD of the explicitly available basis

$U^{\prime}_{t}$ is a subset of columns of the left singular vectors $U$. $T_{t}$ is a transformation matrix (which may be computed as $U^{\prime T}_{t}\overline{U}_{t}$) used to compute the new transfer matrix from a node to its parent as described next.

For non-leaves, the truncated basis is expressed in terms of new compressed transfer operators:

$E^{\prime}_{t_{1}}$ and $E^{\prime}_{t_{2}}$ are the new transfer matrices for the children of a non-leaf node $t$, and $T_{t}$ is a transformation that will be used to compute the new transfer matrix from $t$ to its parent, in a similar fashion. $\left[\begin{smallmatrix}E^{\prime l}_{t_{1}}\\ E^{\prime l}_{t_{2}}\end{smallmatrix}\right]$ are computed by first computing an SVD of the matrix $\left[\begin{smallmatrix}T^{l}_{t_{1}}\overline{E}^{l}_{t_{1}}\\ T^{l}_{t_{2}}\overline{E}^{l}_{t_{2}}\end{smallmatrix}\right]$. The left singular vectors corresponding to singular values below the target compression threshold are truncated, and the remaining subset is partitioned to generate the new transfer matrices $E^{\prime}_{t_{1}}$ and $E^{\prime}_{t_{2}}$. $T_{t}$ is computed as $\sum_{t_{i}}E^{\prime lT}_{t_{i}}T^{l}_{t_{i}}E^{l}_{t_{i}}$.

The structure of the truncation algorithm is identical to that of the upsweep phase in the matrix-vector multiplication. All GPUs start the truncation operations concurrently, each on its subtree, with no interprocess communication required. The computational work at every level being, for every block row, an SVD involving the leaf basis at the leaf level, or the stacked transfer operators for non-leaf levels. Within a GPU, batched SVDs are used [19]. Once all GPUs reach the C-level, a gather operation communicates the new transfer operators from the roots of the subtrees to the leaves of the root tree that is stored on a single GPU. This bootstraps the last phase of the upweep which proceeds on the root GPU.

Two details remain to be discussed. First, in the downsweep phase, the algorithm relied on the basis $V$ being orthogonal. When this is not the case, a pre-processing step is needed to orthogonalize the $\mathcal{V}$ basis tree. A basis is orthogonal if $V_{j}^{lT}V_{j}^{l}$ is the identity matrix for all levels $l$. Orthogonalizing a basis involves performing QR on the finest level basis and then going up the tree to compute new transfer matrices that allow higher level nodes to satisfy the orthogonality condition. This is also done in an upsweep pass that is very similar to the one described above for truncation, but replacing the SVD operations by QR operations. The distributed implementation therefore also proceeds independently on all GPUs up until the C-level. A gather operation is then performed to bootstrap the orthogonalization of the top levels of the basis tree which reside on a single GPU.

Finally, once a new compressed basis $U^{\prime}$ is computed, it remains to approximate the coupling blocks in it. Since $U^{\prime}$ is an orthogonal basis, the best approximation of a block is obtained by an orthogonal projection. Therefore, we can obtain the approximation $A^{\prime}$ by projecting every low rank block $A_{ts}=U_{t}S_{ts}V_{s}^{T}$ on the new basis to obtain:

The new $S^{\prime}_{ts}$ coupling blocks are computed via batched matrix multiplication operations. In this step, there is parallelism across all GPUs and across all levels. In addition, The GPUs are particularly efficient in GEMM operations, and in practice this projection step consumes much less time that the other operations, particularly those involving batched SVDs.

The H2Opus library as well as the code to generate and run all examples below is open source and is available at https://github.com/ecrc/h2opus. Batched kernels for GEMM operations are executed using MAGMA [12, 4], while SVD and QR batches are performed using the algorithms shipped with the open-source library KBLAS, available at https://github.com/ecrc/kblas-gpu. The marshaling GPU kernel and various other utility routines use Thrust [6].

All tests are conducted on Summit, the IBM Power System AC922 supercomputer installed at Oak Ridge National Laboratory. Individual nodes on Summit have 6 NVIDIA V100-SXM2 GPUs with 16GB of HBM2 memory each, but we only used 4 GPUs per node (2 per socket) in our runs. Summit has a fast host-to-device bandwidth which can deliver up to 50 GB/s over PCIe 4.0; in our application, we were able to use a significant portion of it, 40 GB/s, as measured by nvprof. It also has a fat-tree topology network for internode communication that delivers 200Gb/s of bandwidth. To assess the efficiency attained by our algorithms, we measure the performance of the single GPU batched GEMM implementation from MAGMA, with batch elements of size $64\times 64$. All computations are done in double-precision and every point in every plot has been generated as the average of 10 runs after discarding the fastest and slowest timings.

To test the performance and scalability of the matrix-vector multiplication and matrix compression implementations, we performed numerical experiments on two sample matrix sets with different structural characteristics. The first matrix set comes from a spatial statistics application using a point set placed on a 2D grid of side length $a$, and an exponential kernel with a correlation length of $0.1a$. The hierarchical matrix representation of this covariance matrix uses $m=64$ as the finest block size and size of the dense leaves. A geometric admissibility condition $\eta||C_{t}-C_{s}||\geq(D_{t}+D_{s})/2$ is used, where $C$ and $D$ refer to the center and diagonal size of a bounding box of the corresponding point set. We use a value of $\eta=0.9$ and set a rank $k=64$ in the low rank blocks, resulting in an approximation with relative accuracy better than $10^{-7}$ for all problem sizes. This accuracy is computed by sampling $10\%$ of the rows and computing $||Ax-A_{\mathcal{H}^{2}}x||/||Ax||$ with randomly generated vectors with entries from a uniform distribution. The resulting sparsity constant of the matrix, which is a proxy for how finely refined the matrix is in its off-diagonal portions, is $C_{sp}=17$. At the largest size, the matrix has 23 levels, with the top 10 levels on a single master GPU, and the bottom 13 levels on separate 1024 GPUs for a matrix size of 536M.

The second matrix set comes from a 3D Gaussian process and is intended to show the effect of memory pressure—due to a finer refinement in the off-diagonal blocks—on scalability. The matrices are constructed using a set of points on a 3D grid of side length $a$ and uses an exponential kernel with correlation length $0.2a$ [10], a similar admissibility condition to the previous case, and a rank $k=64$ for the low rank blocks. The resulting relative accuracy is now $10^{-3}$ for all problem sizes considered in this section, and the resulting matrix tree has many more leaves than in the 2D case with a larger sparsity constant $C_{sp}=30$.

To test the performance and scalability of the algebraic compression routines, we used the same sets of matrices described earlier. For the 2D tests, we start with the matrix defined as above, with point clusters of size $m=64$, and admissibility parameter $\eta=0.9$. Its low rank blocks are initially constructed using a Chebyshev polynomial approximation of the kernel in the bounding boxes of the point clusters. A $6\times 6$ grid is used, resulting in all low rank blocks having a uniform rank $k=36$. Compression seeks to reduce these ranks to maintain an accuracy of $\tau=10^{-3}$. The 2D test set was scaled using a local matrix size ${{}_{p}}{N}=2^{20}$, reaching a matrix size of 67M on 64 GPUs.

For the 3D tests, a point cluster size $m=64$ and an admissibility parameter $\eta=0.95$ are used in the matrix construction. The low rank blocks are initially computed using a tri-cubic Chebyshev polynomial approximation of the kernel in the bounding boxes of the point clusters, resulting in all low rank blocks having a uniform rank $k=64$. Compression seeks to reduce these ranks to maintain an accuracy of $\tau=10^{-3}$. The test test was scaled using a local matrix size ${{}_{p}}{N}=2^{18}$, reaching a matrix size of 16M on 64 GPUs.

We consider the performance of a Krylov solver for the solution of the integral equation $\mathcal{L}[u(\mathbf{x})]=b(x)$ where $\mathcal{L}$ is the fractional diffusion operator defined as:

The spatially varying diffusivity $a(\mathbf{x},\mathbf{y})$ is defined as the geometric mean of the usual diffusion coefficient, $a(\mathbf{x},\mathbf{y})=\kappa(\mathbf{x})^{1/2}\kappa(\mathbf{y})^{1/2}$, $\beta$ is the fractional order ($0.5<\beta<1$) and we assume it to be constant in this experiment, $\Omega\in R^{n}$ is the region in which the solution $u$ is sought, and $\Omega_{0}\in R^{n}$ is a surrounding region in which volume constraints, somewhat analogous to the Dirichlet boundary conditions in the classical diffusion equations, are imposed on the solution, $u(\mathbf{x})=0$ for $\mathbf{x}\in\Omega_{0}$. We solve the problem with $\Omega=[-1,1]^{2}$, $\Omega_{0}=[-3,3]^{2}\setminus[-1,1]^{2}$, and constant fractional order $\beta=0.75$. We use a diffusivity field of the form

with the bump function $f$ defined as:

and a right hand side $b(x)=1$ in $\Omega$.

The singularity of the kernel in (5) requires that the discretization of the integral be done carefully to allow standard quadrature rules to attain their theoretical convergence rate. To this end, we rewrite the integral as:

where $p_{\mathbf{x}}(\mathbf{y})$ is a function with local support around $\mathbf{x}$ chosen to remove the singularity of the integrand and allow a trapezoidal rule to achieve its quadratic convergence. The first integral in (8) can be then discretized on a regular grid while the second integral involves terms that can be integrated analytically. The details are derived in [8], which generalizes the treatment of [43] to the variable coefficient case. This correction has the property that it allows the hierarchical matrix treatment of the kernel to remain as is, similarly to the FMM-compatible treatment of integral equations on curves [34, 48] and other locally-corrected quadrature schemes [29].

In 2D, using a regular grid of $N$ points with spacing $h$ in $\Omega$, the final discretization results in a system of the form:

where $D$ is an $N\times N$ diagonal matrix with entries

$K$ is an $N\times N$ formally-dense matrix with zero diagonals and entries