CSR Sparse Matrix Multiplication

Introduction

The compressed sparse row (CSR) format is used for encoding sparse matrix. Depending on the level of sparsity, the memory consumption and the computation cost of some of the matrix operations could be significantly reduced.

There are existing software which accelerates sparse matrix operations, such as cuSPARSE and SciPy. But high-level users usually don’t care how sparse matrix operations were implemented.

In this blog post, I would like to quickly discuss the CSR matrix and how CSR matrix multiplication is performed.

CSR Matrix

Suppose we have a (row-major) matrix $\mathbf{M} \in \mathbb{R}^{m \times n}$ in which there are $k$ non-zero values.

To represent the matrix using dense matrix format, we will need the following information.

• The value array $\mathbf{v} \in \mathbb{R}^{mn}$.
• The matrix shape $(m, n)$.

To represent the matrix using sparse CSR matrix format, we will need the following information.

• The value array $\mathbf{v} \in \mathbb{R}^{k}$.
• The column index array $\mathbf{c} \in \mathbb{R}^{k}$. This is also sometimes called the index array.
• The row index array $\mathbf{r} \in \mathbb{R}^{m + 1}$. This is also sometimes called the pointer array. The first element in $\mathbf{r}$, $\mathbf{r}[0]$, always equal to 0. The last element in $\mathbf{r}$, $\mathbf{r}[m]$, always equal to $k$. The number of non-zero values in row $i$ is $\mathbf{r}[i + 1] - \mathbf{r}[i]$.
• The matrix shape $(m, n)$.

For example, suppose we have the following matrix $\mathbf{M} \in \mathbb{R}^{5 \times 7}$,

$$
\begin{align}
\mathbf{M} &=
\begin{bmatrix}
10 & 20 & 0 & 0 & 0 & 0 & 0 \\
0 & 30 & 0 & 40 & 0 & 0 & 0 \\
0 & 0 & 50 & 60 & 70 & 0 & 0 \\
0 & 0 & 0 & 0 & 0 & 80 & 0 \\
0 & 0 & 0 & 0 & 0 & 0 & 0 \\
\end{bmatrix}
\end{align}
$$

The corresponding CSR matrix representation should be

$$
\begin{align}
\mathbf{v} &=
\begin{bmatrix}
10 & 20 & 30 & 40 & 50 & 60 & 70 & 80 \\
\end{bmatrix} \\
\mathbf{c} &=
\begin{bmatrix}
0 & 1 & 1 & 3 & 2 & 3 & 4 & 5 \\
\end{bmatrix} \\
\mathbf{r} &=
\begin{bmatrix}
0 & 2 & 4 & 7 & 8 & 8 \\
\end{bmatrix} \\
\end{align}
$$

To save memory, the number of non-zero values $k$ must be

$$
k < \frac{mn -m - 1}{2}
$$

Advantages of the CSR format

• Efficient arithmetic operations CSR + CSR, CSR $\times$ CSR, etc.
• Efficient row slicing.
• Fast matrix vector products.

Disadvantages of the CSR format

• Slow column slicing operations.
• Changes to the sparsity structure are expensive.

CSR Matrix Multiplication

Because it is inefficient to slice along the the column for CSR matrix, one of the matrix multipliers was transposed before multiplication. The transpose takes $O(N \log N)$ operations where $N$ is the number of non-zero elements in the matrix. The multiplication takes an additional $O(MN)$ operations where $M$ and $N$ are the numbers of non-zero elements in the two multiplier matrices. So if $M \gg \log N$, we could say the asymptotic complexity for CSR matrix multiplication is $O(MN)$.

Python Implementation

The Python implementation is used for demonstrating how to compute CSR matrix multiplication algorithmically. While the implementation is definitely not efficient because it is in Python and has never been optimized, the algorithm complexity should be asymptotically optimal.

csr_mm.py

from __future__ import annotations
from typing import Tuple
import numpy as np


class CSRMatrix:

    def __init__(self, indptr: np.ndarray, indices: np.ndarray,
                 data: np.ndarray, shape: Tuple[int, int]) -> None:

        # Row index
        self.indptr = indptr
        # Column index
        self.indices = indices
        self.data = data
        self.shape = shape

        self.dtype = self.data.dtype

    def toarray(self) -> np.ndarray:
        """Convert CSR matrix to Numpy array.

        Returns:
            np.ndarray: Dense matrix.
        """

        array = np.zeros(self.shape).astype(self.data.dtype)
        num_rows = self.shape[0]
        for i in range(num_rows):
            num_vals = self.indptr[i + 1] - self.indptr[i]
            for k in range(num_vals):
                val = self.data[self.indptr[i] + k]
                j = self.indices[self.indptr[i] + k]
                array[i][j] = val

        return array

    def transpose(self) -> CSRMatrix:
        """Transpose CSR matrix.

        O(NlogN) where N is the number of non-zero values.

        Returns:
            CSRMatrix: Transposed CSR matrix.
        """

        col_2d_idx = self.indices

        # Compute row 2d idx
        # O(N)
        row_2d_idx = np.zeros_like(col_2d_idx)
        k = 0
        num_rows = self.shape[0]
        for i in range(num_rows):
            num_vals = self.indptr[i + 1] - self.indptr[i]
            for j in range(num_vals):
                row_2d_idx[k + j] = i
            k += num_vals
        assert k == self.indptr[-1]

        # Col 2d index becomes row 2d index
        # Row 2d index becomes col 2d index
        new_row_2d_idx = col_2d_idx
        new_col_2d_idx = row_2d_idx

        # https://numpy.org/doc/1.22/reference/generated/numpy.lexsort.html
        # Sort by new_row_2d_idx, then by new_col_2d_idx
        # O(NlogN)
        ind = np.lexsort((new_col_2d_idx, new_row_2d_idx))
        new_row_2d_idx = new_row_2d_idx[ind]
        new_col_2d_idx = new_col_2d_idx[ind]

        # Create CSR matrix
        # O(N)
        indices = new_col_2d_idx
        data = self.data[ind]
        shape = (self.shape[1], self.shape[0])
        num_rows = shape[0]
        indptr = np.zeros(num_rows + 1).astype(np.int32)

        for i in new_row_2d_idx:
            indptr[i + 1] += 1

        for i in range(num_rows):
            indptr[i + 1] += indptr[i]

        indices = np.array(indices).astype(np.int32)
        data = np.array(data).astype(self.dtype)

        csr_matrix = CSRMatrix(indptr=indptr,
                               indices=indices,
                               data=data,
                               shape=shape)
        return csr_matrix

    def append_csr_rows(self, row_mat: CSRMatrix) -> CSRMatrix:
        """Append an another CSR matrix. 

        Args:
            row_mat (CSRMatrix): Another CSR matrix that has the same width.

        Returns:
            CSRMatrix: The resulted new CSR matrix.
        """

        assert len(row_mat.shape) == 2
        assert row_mat.shape[1] == self.shape[1]
        assert self.dtype == row_mat.dtype

        data = np.append(self.data, row_mat.data)
        indices = np.append(self.indices, row_mat.indices)
        shape = (self.shape[0] + row_mat.shape[0], self.shape[1])

        indptr = row_mat.indptr.copy() + self.indptr[-1]
        indptr = np.append(self.indptr[:-1], indptr)

        csr_matrix = CSRMatrix(indptr=indptr,
                               indices=indices,
                               data=data,
                               shape=shape)
        return csr_matrix

    def get_csr_row(self, row: int) -> CSRMatrix:
        """Get one row as a CSR matrix.

        Args:
            row (int): Row index.

        Returns:
            CSRMatrix: Row CSR matrix.
        """

        assert 0 <= row and row < self.shape[0]

        data = self.data[self.indptr[row]:self.indptr[row + 1]]
        indices = self.indices[self.indptr[row]:self.indptr[row + 1]]
        indptr = self.indptr[row:row + 2]
        indptr = indptr - indptr[0]
        shape = (1, self.shape[1])

        csr_matrix = CSRMatrix(indptr=indptr,
                               indices=indices,
                               data=data,
                               shape=shape)
        return csr_matrix

    def dot_vector_transposed(self, vec_transposed: CSRMatrix) -> CSRMatrix:
        """The dot product of the CSR matrix with another transposed vector CSR matrix.

        Suppose the CSR matrix is of shape [A, B], the transposed vector CSR matrix
        should be of shape [1, B]. The resulted dot product transposed CSR matrix is
        of shape [1, B].

        O(MN) where M is the non-zero values in the CSR matrix and N is the number of
        non-zero values in the vector CSR matrix.

        Args:
            vec_transposed (CSRMatrix): The vector CSR matrix which has been transposed.

        Returns:
            CSRMatrix: The resulted CSR matrix.
        """

        assert len(vec_transposed.shape) == 2
        assert vec_transposed.shape[1] == self.shape[1]
        assert vec_transposed.shape[0] == 1

        num_vals_vec = vec_transposed.indptr[-1] - vec_transposed.indptr[0]

        num_rows = self.shape[0]
        shape = (1, num_rows)
        indptr = np.zeros(2).astype(np.int32)
        indices = []
        data = []

        # O(MN)
        # where M is the non-zero elements of matrix and
        # N is the number of non-zero elements of vector
        for i in range(num_rows):
            num_vals_matrix = self.indptr[i + 1] - self.indptr[i]

            idx_1 = 0
            idx_2 = 0

            val = 0
            while idx_1 < num_vals_matrix and idx_2 < num_vals_vec:
                if self.indices[self.indptr[i] +
                                idx_1] == vec_transposed.indices[idx_2]:
                    val += self.data[self.indptr[i] +
                                     idx_1] * vec_transposed.data[idx_2]
                    idx_1 += 1
                    idx_2 += 1
                elif self.indices[self.indptr[i] +
                                  idx_1] < vec_transposed.indices[idx_2]:
                    idx_1 += 1
                else:
                    idx_2 += 1

            if val != 0:
                data.append(val)
                indices.append(i)

        indices = np.array(indices).astype(np.int32)
        data = np.array(data).astype(
            np.result_type(self.dtype,
                           vec_transposed.dtype))  # Type Inference TBD
        indptr[-1] = len(data)

        csr_vec_transposed = CSRMatrix(indptr=indptr,
                                       indices=indices,
                                       data=data,
                                       shape=shape)

        return csr_vec_transposed

    def dot_matrix_transposed(self, mat_transposed: CSRMatrix) -> CSRMatrix:
        """The dot product of the CSR matrix with another transposed CSR matrix.

        O(MN) where M is the non-zero values in the CSR matrix and N is the number of
        non-zero values in the another transposed CSR matrix.

        Args:
            mat_transposed (CSRMatrix): Another transposed CSR matrix.

        Returns:
            CSRMatrix: The resulted CSR matrix.
        """

        assert len(mat_transposed.shape) == 2
        assert mat_transposed.shape[1] == self.shape[1]

        num_rows = self.shape[0]
        shape = (mat_transposed.shape[0], num_rows)

        row_vec = mat_transposed.get_csr_row(row=0)
        csr_mat_transposed = self.dot_vector_transposed(vec_transposed=row_vec)

        for i in range(1, shape[0]):
            row_vec = mat_transposed.get_csr_row(row=i)
            product_row_vec = self.dot_vector_transposed(
                vec_transposed=row_vec)
            csr_mat_transposed = csr_mat_transposed.append_csr_rows(
                row_mat=product_row_vec)

        return csr_mat_transposed

    def dot_vector(self, vec: CSRMatrix) -> CSRMatrix:
        """The dot product of the CSR matrix with another vector CSR matrix.

        Args:
            vec (CSRMatrix): Another vector CSR matrix.

        Returns:
            CSRMatrix: The resulted CSR matrix.
        """

        assert len(vec.shape) == 2
        assert vec.shape[0] == self.shape[1]
        assert vec.shape[1] == 1

        # Transpose the CSR vector
        vec_transposed = vec.transpose()

        csr_vector_transposed = self.dot_vector_transposed(
            vec_transposed=vec_transposed)

        csr_vector = csr_vector_transposed.transpose()

        return csr_vector

    def dot_matrix(self, mat: CSRMatrix) -> CSRMatrix:
        """The dot product of the CSR matrix with another CSR matrix.

        O(MN) where M is the non-zero values in the CSR matrix and N is the number of
        non-zero values in the another SR matrix.

        Args:
            mat (CSRMatrix): Another CSR matrix.

        Returns:
            CSRMatrix: The resulted CSR matrix.
        """

        assert len(mat.shape) == 2
        assert mat.shape[0] == self.shape[1]

        mat_transposed = mat.transpose()

        csr_mat_transposed = self.dot_matrix_transposed(
            mat_transposed=mat_transposed)

        csr_mat = csr_mat_transposed.transpose()

        return csr_mat


def create_csr_matrix_from_dense_matrix(dense_matrix: np.ndarray) -> CSRMatrix:
    """Create a CSR matrix from a dense matrix.

    Args:
        dense_matrix (np.ndarray): Dense matrix.

    Returns:
        CSRMatrix: The resulted CSR matrix.
    """

    shape = dense_matrix.shape
    assert len(shape) == 2
    num_rows = shape[0]
    num_cols = shape[1]
    # The size of indptr array is num_rows + 1
    indptr = np.zeros(num_rows + 1).astype(np.int32)
    indices = []
    data = []
    for i in range(num_rows):
        for j in range(num_cols):
            if dense_matrix[i][j] != 0:
                data.append(dense_matrix[i][j])
                indices.append(j)
            indptr[i + 1] = len(indices)
    indices = np.array(indices).astype(np.int32)
    data = np.array(data).astype(dense_matrix.dtype)
    csr_matrix = CSRMatrix(indptr=indptr,
                           indices=indices,
                           data=data,
                           shape=shape)
    return csr_matrix


def create_random_matrix(shape: Tuple[int, int]) -> np.ndarray:
    """Create random matrix.

    Args:
        shape (Tuple[int, int]): Matrix shape.

    Returns:
        np.ndarray: Resulted random matrix.
    """

    matrix = np.random.randint(low=1, high=100, size=shape)

    return matrix


if __name__ == "__main__":

    num_tests = 100
    np.random.seed(0)

    for _ in range(num_tests):

        a = np.random.randint(low=1, high=20)
        b = np.random.randint(low=1, high=20)
        c = np.random.randint(low=1, high=20)
        mat_1 = create_random_matrix(shape=(a, b))
        mat_2 = create_random_matrix(shape=(b, c))
        mat_3 = mat_1.dot(mat_2)

        csr_mat_1 = create_csr_matrix_from_dense_matrix(dense_matrix=mat_1)
        csr_mat_2 = create_csr_matrix_from_dense_matrix(dense_matrix=mat_2)
        csr_mat_3 = csr_mat_1.dot_matrix(csr_mat_2)

        assert np.array_equal(mat_3, csr_mat_3.toarray())

References

• Sparse Matrix - Wikipedia
• CSR Matrix - SciPy