UST: Diagonal Formats

[aartbik]

Diagonal formats provide efficient storage for sparse matrices where the nonzeros are primarily clustered along diagonals. As was shown in our first blog, the UST describes the diagonal (or anti-diagonal) relation between row and column indices using a dimension variable subtraction (or addition) in a compressed level. A subsequent range level defines one of these dimension variables as an actual index into the dense storage of each diagonal, as shown below.

(i, j) -> (j-i : compressed, i : range)	# DIA-I
(i, j) -> (j-i : compressed, j : range)	# DIA-J

Consider the following example diagonal matrix.

The DIA-I storage is shown below. The metadata for this scheme is minimal, and consists of a level 0 position array to store the number of nonzero diagonals and a coordinate array to store the offsets for these diagonals (j-i is -4, 0, or 5). To accommodate access through the row-index, each diagonal is stored using a fixed size of 8 elements, where padding is added to the left (prefix) for the lower triangular diagonals and to the right (suffix) for the upper triangular diagonals. This external padding represents indices outside the original matrix space and must not be accessed. Conversely, internal padding is required between elements, such as between -5 and -6, to preserve the structural density of the diagonals. These zeros are routinely accessed, sacrificing operations on zero values to gain storage efficiency.

The DIA-J storage is analogous, but indexes through the column-index instead. Similar to DIA-I, each diagonal uses a fixed size of 8 elements, but the roles of the external padding are reversed: padding is added to the right (suffix) for the lower triangular diagonals and to the left (prefix) for the upper triangular diagonals.

Because the required padding is proportional to the row or column index within DIA-I and DIA-J schemes, respectively, DIA-I offers better storage efficiency for wide matrices, while DIA-J is more suited for tall matrices.

The UST in nvmath-python provides zero-cost interoperability with the DIA formats supported by SciPy and CuPy.

 dia = scipy.sparse.dia_array(a)
 ust = Tensor.from_package(dia)
 print(ust.tensor_format.name)  # prints DIAJ

Alternatively, compact diagonal storage can be defined by adding an unpadded property on the range level format.

(i, j) -> (j-i : compressed, i : range(unpadded)) 	# DIA-I-Compact
(i, j) -> (j-i : compressed, j : range(unpadded))	# DIA-J-Compact

The resulting memory layout is depicted below. Within this scheme, the extent of every diagonal is defined by the index boundary. This approach reduces the storage requirements compared to the aforementioned padded variants, but at the expense of additional overhead during data retrieval to account for variable diagonal lengths.

Because SciPy and CuPy lack native support for unpadded schemes, explicit conversion is required. While this introduces some data movement overhead, it is straightforward to express in nvmath-python (the current reference implementation for such conversions can be costly for large matrices, but is subject to continuous improvement).

ust_out = ust_in.convert(tensor_format=NamedFormats.DIAICompact)

The UST integrates seamlessly with the stateful generic Matmul operation in nvmath-python, enabling the amortization of planning overhead across successive runs. This is particularly advantageous for workloads in High-Performance Computing (HPC) and Deep Learning (DL), such as iterative solvers or high-throughput inference servers. Leveraging these formats within nvmath-python is straightforward. The following snippet illustrates the basic structure for managing operands a, b, and c—whether they represent vectors, matrices, or tensors—using the stateful Matmul operation:

with Matmul(a, b, c) as mm:
    mm.plan()       # initial setup overhead, potentially expensive
    mm.execute()    # repeated execution amortizes the setup overhead

For standard configurations, such as Sparse Matrix-Vector Multiplication (SpMV) using CSR, the planning phase primarily involves configuring the cuSPARSE library. In contrast, specialized structures such as the aforementioned DIA formats may lack native library support, causing the planning phase to invoke a more intensive, one-time JIT/LTO sequence for dynamic CUDA kernel generation. Code caching ensures this overhead is not revisited, even across distinct planning cycles for identical operation-format pairings. Upon completion of the compilation phase during planning, specialized kernel execution for these custom formats frequently outperforms generic library implementations.

To illustrate these concepts, we measured the float32 SpMV execution times (after planning) of CSR, DIA-I, and DIA-I-Compact on a NVIDIA B200 for matrices from the CUTh Benchmark Suite (symmetric positive definite band matrices with 7 diagonals arising from large-scale finite difference/volume discretization of the heat equation on uniform grids; in this test, symmetry remains unexploited by processing both triangular sections explicitly; more on that later). The results are visualized in the following plot, with execution time on the primary y-axis alongside the count of stored entries on the secondary y-axis.

Both DIA formats significantly outperform the standard CSR format in execution speed. While performance is comparable between the padded and unpadded DIA variants, the unpadded format matches the memory efficiency of CSR, thereby delivering an optimal balance between execution performance and memory footprint. However, note that the planning phase for DIA is more costly due to the initial JIT/LTO step; consequently, the performance benefits are only realized when the SpMV operation is executed repeatedly.

A few final observations in closing. First, the UST seamlessly generalizes these principles to multi-dimensional diagonal structures within higher-order tensors. Finally, note that the UST functionality is currently experimental; expect the syntax and internal implementation to evolve over time. Lastly, while our mapping (associating (i, j) with (j-i, i)) uses two levels, Stephen Chou’s thesis explores a three-level DIA representation (associating (i, j) with (j-i, i, j)), with each approach presenting unique trade-offs.