Sparse Matrices

The idea in one sentence

When almost every entry of a matrix is zero, the honest move is to stop storing the zeros, and every sparse format is one answer to the question that follows: then how do you find what is left?

Where the zeros come from

Sparsity is not a trick somebody invented to save memory. It is what the data looks like.

A batch of token ids arriving at an embedding layer is a matrix of one-hot rows: one entry per row is 1 and the other fifty thousand are 0. A graph is an adjacency matrix, and a user who follows two hundred accounts out of two hundred million contributes a row that is one part in a million away from empty. An attention mask that only lets each position see its neighbours is a band around the diagonal. A pruned layer is a weight matrix somebody deliberately emptied.

The four patterns look different, and that difference is the whole subject. One entry per row, scattered, banded, and random are four different problems, and they get four different answers below.

The bargain

Lecture 23 said a dense matrix is a promise about addresses: give me a row and a column and I will compute where that entry lives, with one multiply and one add. Lecture 27 turned that promise into an object with an algebra. All of it depends on the same thing, that the address is computed rather than looked up.

A sparse matrix breaks that promise. Once you skip the zeros, the entry in row $i$ and column $j$ is no longer at a place you can compute, because you no longer know how many entries came before it. So the position has to be written down. Every sparse format is a different way of writing it down, and every one of them pays the same tax.

That tax is why sparsity has a threshold. Storing a float32 costs 4 bytes. Storing the column it sat in costs another 4. So a sparse format that keeps one index per value spends 8 bytes on what dense spends 4, and it wins only when it can skip more than half the matrix. Below that, the honest answer is to store the zeros. This is the first of two thresholds in this chapter, and it turns out to be the generous one.

Three arrays, in any order

The first idea anybody has is also the right one to start with. Keep a list of the entries that exist, each one carrying its own row and column. That is coordinate format, universally called COO.

Here is the matrix this chapter uses. Six by six, nine entries that exist, and one row that is entirely empty, which will matter shortly.

COO is easy to build. You can append to it, you can build it out of order, and duplicate entries can be summed at the end. What it is bad at is being read. Ask for row 3 and there is nothing to do but scan all nine entries and collect the ones whose row is 3. On a matrix with a hundred million entries that is a hundred million comparisons to answer a question about one row.

Two cousins solve the same building problem in other ways. Dictionary of keys stores a hash map from an index pair to a value, so a single entry can be written in constant time. List of lists keeps one list of column indices and one list of values per row, so a row can grow without touching the others. Both are for construction. Neither is for arithmetic. The habit that follows is worth keeping: build in one format, convert, compute in another.

Compress the rows

Look at the row array in Figure 2 again: 0, 0, 1, 1, 2, 3, 3, 5, 5. Sorted, and full of repetition. Anything sorted with repeats can be replaced by the places where it changes.

That is compressed sparse row, CSR, and it is three arrays again, but the third one is different in kind. values and colIndex still have one slot per entry. The third array, rowPtr, has one slot per row plus one, and it holds where each row starts.

import numpy as np
import scipy.sparse as sp

A = np.array([[10,  0,  0, 12,  0,  0],
              [ 0,  0, 11,  0,  0, 13],
              [ 0, 16,  0,  0,  0,  0],
              [ 0,  0, 11,  0, 18,  0],
              [ 0,  0,  0,  0,  0,  0],
              [14,  0,  0,  0,  0, 15]], dtype=float)

S = sp.csr_array(A)
print(S.data)                # [10. 12. 11. 13. 16. 11. 18. 14. 15.]
print(S.indices)             # [0 3 2 5 1 2 4 0 5]
print(S.indptr)              # [0 2 4 5 7 7 9]
print(S.nnz, A.size)         # 9 36

Read indptr against the picture. Row 0 runs from slot 0 to slot 2, so it owns values 10 and 12. Row 2 runs from 4 to 5, one entry, the 16. Row 3 runs from 5 to 7, the 11 and the 18. Row 4 runs from 7 to 7, which is nothing at all, and that is the empty row showing up as two equal numbers in a row.

The pointer pair is the whole trick

rowPtr holds boundaries, not counts. Its entries are useful only in adjacent pairs, and the pair is what a row is.

Write the pair down and the loop writes itself. Everything a row needs is a contiguous range of slots, and the columns for that row sit in the same range of colIndex.

def csr_matvec(values, col_index, row_ptr, x):
    out = np.zeros(len(row_ptr) - 1)
    for i in range(len(out)):
        for k in range(row_ptr[i], row_ptr[i + 1]):    # this row's slice
            out[i] += values[k] * x[col_index[k]]
    return out

x = np.array([1., 2., 3., 4., 5., 6.])
print(csr_matvec(S.data, S.indices, S.indptr, x))     # [ 58. 111.  32. 123.   0. 104.]
print(A @ x)                                          # [ 58. 111.  32. 123.   0. 104.]

Nine multiplies, and the same answer as the dense product, which does thirty six. Row 4 runs zero times around the inner loop and produces its zero for free.

Two properties fall out of that loop. The rows are independent, so a machine with many cores can have each one take a range of rows and never talk to its neighbours. And values is read straight through, in order, once, which is the one thing about sparse code that memory likes. What memory does not like is x[col_index[k]], and that indirect read is the reason for everything in the last third of this chapter.

The twin, and two more

Swap the roles of rows and columns and you get compressed sparse column, CSC: a colPtr with one slot per column plus one, a rowIndex, and the same values in a different order. CSC of a matrix is CSR of its transpose, and nothing else about it is new.

Which one you want is decided by the loop you are about to run, exactly as row-major against column-major was in Lecture 23. This is that choice again, one level up, with the strides replaced by stored indices.

Two more formats are worth knowing by name. Block sparse row keeps small dense blocks instead of single entries, so the index is paid once per block rather than once per number, and the blocks themselves can go to a dense kernel. It is the right format when the nonzeros come in clumps, which pruning schemes can be designed to produce. Diagonal format stores each occupied diagonal as a full vector and one offset, which is nearly free for the banded mask in Figure 1 and disastrous for the scattered graph. Every format is a bet about where the nonzeros are, and a format that bets wrong stores more than dense would.

The product that never touches a zero

Here is the loop you just ran, drawn.

Multiplying two sparse matrices is the same instinct one level up. Walk the nonzeros of $A$ one row at a time. Each nonzero $A_{ik}$ activates row $k$ of $B$, and the whole of that row gets scaled and accumulated into row $i$ of the result. Zeros are never visited, because you never had them.

The count of that work is worth having in closed form, because it explains everything below. Every nonzero in column $k$ of $A$ meets every nonzero in row $k$ of $B$, so the number of multiplications is the sum over $k$ of those two counts multiplied together. Nothing about the matrix sizes appears in it. Only the structure does.

That is a devastating saving on paper. At one percent density on a $2048 \times 2048$ matrix it is eight hundred and sixty thousand multiplies against the dense product’s eight and a half billion, a factor of ten thousand.

n = 2048
rng = np.random.default_rng(0)
M = sp.random_array((n, n), density=0.01, format='csr', dtype=np.float32, rng=rng)

per_row = np.diff(M.indptr)                      # nonzeros in each row of M
per_col = np.bincount(M.indices, minlength=n)    # nonzeros in each column of M
print(int((per_col * per_row).sum()))            # 860444
print(n ** 3)                                    # 8589934592

Now measure it.

When sparse actually wins

Two costs, two thresholds, and they are nowhere near each other.

Memory first, because it is the easy one. A CSR matrix of float32 values with 32-bit indices spends 8 bytes per stored entry against dense’s 4 per position, plus one pointer per row, which disappears at any real size. So the stored size is about twice the density, and it crosses dense at a density of one half.

for d in (0.01, 0.25, 0.5):
    M = sp.random_array((n, n), density=d, format='csr', dtype=np.float32, rng=rng)
    stored = M.data.nbytes + M.indices.nbytes + M.indptr.nbytes
    print(d, M.indices.dtype, round(stored / (n * n * 4), 3))

# 0.01 int32 0.02
# 0.25 int32 0.5
# 0.5 int32 1.0

Exactly as predicted, and exactly where predicted. Now time, on the same machine, with the same matrices.

from timeit import repeat
bench = lambda f, k: min(repeat(f, number=k, repeat=7)) / k

D = rng.standard_normal((n, n)).astype(np.float32)
dense = bench(lambda: D @ D, 3)
print(round(dense * 1e3, 2), "ms for the dense product")     # 6.19 ms

for d in (0.001, 0.002, 0.005, 0.01, 0.02, 0.05):
    M = sp.random_array((n, n), density=d, format='csr', dtype=np.float32, rng=rng)
    print(d, round(bench(lambda: M @ M, 1) / dense, 3))

# 0.001 0.009
# 0.002 0.024
# 0.005 0.154
# 0.01  0.738
# 0.02  4.404
# 0.05  13.141

At one percent the sparse product does one ten-thousandth of the arithmetic and takes three quarters of the time. At two percent it is four and a half times slower than multiplying every zero by hand.

That comparison deserves a caveat in each direction. It is harsh on sparse, because randomly placed nonzeros are the worst case for locality and a matrix with structure would do better. It is also harsh on sparse because the dense product goes to Apple’s Accelerate library, which uses the chip’s matrix hardware and more than one core, while the sparse product is a single thread of ordinary C. Grant both and the gap in arithmetic is still ten thousand to one, and nothing on either list accounts for a gap that size. The arithmetic was never the problem.

The problem is x[col_index[k]]. A dense kernel knows every address before it runs, which is what let Lecture 24 tile it and Lecture 26 hand it to thousands of threads. A sparse kernel learns each address from memory it just read, one at a time, and can predict nothing. It fetches a whole cache line to use one number from it, exactly the failure Lecture 23 measured, and it cannot be reordered out of existence because the order is data.

The hardware’s compromise

Put all of this next to a GPU and it gets worse. The whole model in Lecture 25 rests on thirty two neighbouring threads wanting thirty two neighbouring addresses, and a sparse row hands them thirty two addresses it looked up. Tensor cores are worse still: they eat a fixed rectangular block and have no opinion about which of its entries are zero.

So the hardware made an offer. Instead of allowing sparsity anywhere, demand it everywhere, in a fixed pattern that the silicon can decode. NVIDIA’s article Accelerating Inference with Sparsity Using the NVIDIA Ampere Architecture and NVIDIA TensorRT (July 20, 2021) describes the rule as: in each contiguous block of four values, two must be zero. Half the weights go, the survivors are stored with metadata giving their positions, and the sparse tensor core reads that metadata and skips the rest. The article’s table gives 624 sparse TOPS against 312 dense for FP16 input, and reports ResNet-50 at 76.2 percent top-1 accuracy sparse against 76.1 dense after a three-step retraining recipe.

That is a strange bargain, and it is worth seeing why anyone took it. Fifty percent sparsity is far less than the ninety plus percent an unconstrained pruner can reach on the same network. What the hardware buys with the constraint is that the address becomes computable again, from two bits sitting at a known place which the silicon decodes as it reads. The pattern stays regular enough to keep the tensor core fed, and that is the only currency here.

Where this is going

Part VII asked one question in six chapters: where does a number live, and what does it cost to go and get it. Memory is one long line. A layout is the rule for walking it. A CPU that walks it well beats one that walks it badly by a factor nothing else in the chain can match. A GPU is thousands of walkers who must agree on a direction. Layout algebra is that agreement written down. And this chapter is what happens when you give up the rule and store the answers instead: you save the memory you expected and lose the speed you did not expect to lose.

If there is one habit to keep from all of it, it is that the arithmetic is rarely what costs. Counting multiplies told you the sparse product was ten thousand times cheaper, and it was three quarters as fast. The number that predicted the truth was the number of unpredictable addresses.

This part ends here as a text and starts somewhere else as a practice. Go and write a kernel, badly, and measure it. The Library has the three places worth starting: Triton Puzzles for thinking in blocks, LeetGPU for CUDA practice judged on real hardware, and the GPU MODE kernel leaderboard for problems nobody has finished optimising. The linear algebra in the earlier parts tells you what to compute. This part tells you what it will cost. The rest is measurement.

And the machine you now understand has one more part to serve. Everything Part VII built runs a model as it is; Lecture 29 opens Part VIII by asking what it costs to change one, and answers with this course’s own algebra: rank, subspaces, and rotations, hired to fine-tune. The matrices stop being examples there and start being the workload.