Sparse Tensor Support in MatX

[aartbik]

MatX is in the process of adding experimental support for sparse tensors. The implementation is based on the Universal Sparse Tensor (UST) type that uses a tensor format DSL (Domain Specific Language) to describe a vast space of storage formats. Although the UST type can easily define many common storage formats (such as dense vectors and matrices, sparse vectors, sparse matrices in COO, CSR, CSC, DCSR, DCSC, or BSR format, with generalizations for sparse tensors), it can also define many less common storage formats.

Despite the forward-looking design of using the UST type inside MatX, the current experimental supported API is limited to sparse matrices in COO, CSR, or CSC format. Such sparse matrices can be constructed using "make" methods on the constituent buffers (similar constructors can be found in e.g. SciPy or CuPy or torch sparse):

auto Acoo = make_tensor_coo(vals, row, col, {m, n});
auto Acsr = make_tensor_csr(vals, rowp, col, {m, n});
auto Acsc = make_tensor_csc(vals, row, colp, {m, n});

Alternatively, sparse matrices can be constructed from dense matrices as follows.

Acoo = dense2sparse(A);

Other supported operations include conversions back to dense matrices, Sparse-Matrix times Vector (SpMV), Sparse-Matrix times Matrix (SpMM), and solving a sparse system of linear equations.

  A = sparse2dense(Acoo);
  V = matvec(Acoo, W);        // SpMV
  C = matmul(Acoo, B);        // SpMM
  X = solve(Acsr, Y);

More details can be found in the Sparse Tensor Type Section of the MatX documentation.

Give it a try and let us know what you think! Given sufficient interest, the assortment of supported formats and operations is expected to grow!

[aartbik]

MatX now also support the DIA format through the UST type. For an m x n matrix, this format uses a linearized storage where each nonzero diagonal has n entries, padded with zeros on the right for the lower triangular part and padded with zeros on the left for the upper triagonal part.

An example follows:

//
 // Creates a 6x6 DIA matrix with 3 nonzero diagonals.
 //
 // |  4  1  0  0  0  0 |
 // | -1  4  1  0  0  0 |
 // |  0 -1  4  1  0  0 |
 // |  0  0 -1  4  1  0 |
 // |  0  0  0 -1  4  1 |
 // |  0  0  0  0 -1  4 |
 //
 auto dvals = make_tensor<float>({3 * 6});
 auto doffsets = make_tensor<int>({3});
 dvals.SetVals({-1, -1, -1, -1, -1, 0, 4, 4, 4, 4, 4, 4, 0, 1, 1, 1, 1, 1});
 doffsets.SetVals({-1, 0, 1});
 auto Adia = experimental::make_tensor_dia(dvals, doffsets, {6, 6});

Currently, the only supported operation is SpMV.

V = matvec(Adia, W); // SpMV

Since the format is not supported by cuSPARSE, this operation is implemented with a hand-written kernel.

For matrices with only a few nonzero diagonals that are close to the main diagonal, this format is much more compact than either COO or CSR. Furthermore, for such cases, the hand-written kernel performs much better than using cuSPARSE SpMV with either the COO or CSR format.

Give it a try and let us know what you think!

[aartbik]

More information on the Universal Sparse Tensor can be found in this blog posting:

Establishing a Scalable Sparse Ecosystem with the Universal Sparse Tensor

[aartbik]

A second blog posting on the Universal Sparse Tensor is up as well:

Simplify Sparse Deep Learning with Universal Sparse Tensor in nvmath-python

Even though this focuses on nvmath-python, it sketches some of the design difference of the DSL with MatX.