Sparse Tensor Encoding

Compressed Sparse Row (CSR) is a canonical format for representing 2D sparse matrices. It uses three arrays:

• values: Stores only the non-zero values.
• col_indices: Stores the column index for each corresponding non-zero value.
• row_ptr: Stores the cumulative count of non-zero values up to the start of each row.

This format is highly efficient for row-wise operations, such as sparse matrix-vector multiplication (SpMV), because iterating through all non-zeros in a row is a contiguous memory access. It is the standard format in scientific computing libraries like SciPy.

Compressed Sparse Column (CSC) is the column-major counterpart to CSR. It also uses three arrays but transposes the access pattern:

• values: Stores the non-zero values.
• row_indices: Stores the row index for each non-zero.
• col_ptr: Stores the cumulative count of non-zero values up to the start of each column.

CSC is optimal for column-oriented operations. Converting between CSR and CSC is a common preprocessing step to align the data layout with the dominant access pattern of the computational kernel, minimizing cache misses.

Coordinate Format (COO) is the simplest sparse representation, storing a list of tuples for each non-zero entry. It typically uses three parallel arrays:

• values: The non-zero data.
• row_indices: The row coordinate for each value.
• col_indices: The column coordinate for each value.

COO is easy to construct and modify, making it ideal for incremental building of sparse matrices. However, it is generally less efficient for computation than CSR/CSC due to irregular memory access. It is often used as an intermediate format before conversion to a compressed format like CSR for final computation.

Block Compressed Sparse Row (BSR) extends CSR by grouping non-zero elements into dense, fixed-size blocks (e.g., 2x2, 4x4). Instead of storing individual scalar values, it stores small dense matrices.

• This format exploits sub-matrix structure common in problems from finite element analysis or certain convolutional layers.
• It reduces the storage overhead of indices (one index per block instead of per scalar) and enables the use of highly optimized dense linear algebra routines (BLAS) on the blocks, improving computational intensity and cache utilization.

The ELLPACK (ELL) format is designed for vector architectures and GPUs. It stores non-zeros in two dense 2D arrays:

• data: A num_rows x max_nnz_per_row matrix holding non-zero values, padded with zeros.
• indices: A corresponding matrix holding the column indices.

Performance is excellent when the number of non-zeros per row is relatively uniform, as it enables coalesced memory accesses and efficient SIMD/SIMT execution. For matrices with highly irregular row lengths, padding overhead becomes prohibitive, leading to hybrid formats like ELL-COO or ELLPACK-R.

Compressed Sparse Fiber (CSF) is a generalization of CSR/CSC to N-dimensional sparse tensors (order > 2). It recursively compresses along each mode (dimension).

• The structure is a tree where each level corresponds to a tensor mode, storing pointers to the next level and finally the non-zero values at the leaves.
• CSF is highly memory-efficient for high-order tensors but can have complex traversal logic. It is foundational for sparse tensor algebra (e.g., sparse tensor-times-matrix) in applications like data analytics and tensor decomposition for recommendation systems.

Block Sparse formats generalize CSR/CSC by storing small, dense blocks of non-zero values instead of individual scalars. This requires:

• data: A tensor of dense blocks.
• Block-specific indices and indptr arrays that point to blocks.

Use Case: Exploits structured sparsity patterns common in pruned neural networks (e.g., block-wise weight pruning). It dramatically improves computational efficiency by leveraging SIMD and tensor core units on modern GPUs and NPUs for block matrix multiplication. Implemented in libraries like cuSPARSE (bsr format) and specialized kernels in TensorRT.

Major ML frameworks abstract sparse encoding behind high-level APIs, each with distinct design choices:

• PyTorch: Uses COO as the primary internal format for flexibility, with CSR/CSC as auxiliary formats. Supports autograd for sparse tensors.
• TensorFlow: The tf.sparse.SparseTensor is always in COO-like format. For operations, it often converts temporarily to a densified representation or uses tf.raw_ops for direct sparse kernels.
• JAX: Offers jax.experimental.sparse with BCSR (Blocked CSR) as a first-class citizen, aligning with its strength in composable transformations and JIT compilation.

Consideration: The computational graph optimizer (like XLA or TorchScript) must recognize sparse tensor patterns to apply fusions and select optimal kernels.

Algorithms that systematically remove parameters from a neural network to induce sparsity, creating the tensors that sparse encoding formats store. Major methods include:

• Magnitude Pruning: Removes weights with the smallest absolute values.
• Structured Pruning: Removes entire channels, filters, or blocks to maintain hardware-friendly patterns.
• Iterative Pruning: Repeatedly prunes and retrains a model to recover accuracy. The resulting sparsity pattern (unstructured vs. structured) directly dictates which encoding format is most efficient.

Model optimization techniques co-designed with target silicon characteristics. For sparsity, this involves:

• Aligning pruning granularity (e.g., 2:4 or 1:2 structured sparsity) with the accelerator's native sparse compute units.
• Choosing an encoding format (e.g., Blocked-ELLPACK) that matches the hardware's memory access patterns.
• Compiler passes that restructure the graph to maximize the utilization of sparse execution units, a process tightly coupled with graph lowering.

The compiler process of transforming a high-level computational graph into low-level, hardware-specific instructions. For sparse tensors, lowering involves:

• Replacing dense operators (e.g., torch.mm) with their sparse counterparts (e.g., torch.sparse.mm).
• Injecting format conversion nodes when data flows between sparse and dense subgraphs.
• Selecting the optimal kernel from a registry based on the tensor's sparse encoding and the target hardware's capabilities.

Transforming the in-memory arrangement of tensor data to improve access efficiency. For sparse tensors, this goes beyond format choice (CSR vs. CSC) to include:

• Blocking/Clustering non-zero values to improve cache locality.
• Reordering rows/columns (matrix bandwidth reduction) to cluster non-zeros near the diagonal.
• Coordinating layout with loop tiling strategies in sparse kernels to minimize pointer-chasing overhead.

A compiler optimization that merges multiple sequential operations into a single kernel. With sparse tensors, fusion is critical to avoid expensive format conversion between operations. Examples include:

• Fusing a sparse linear layer with a subsequent ReLU activation, executing both in a single pass over the sparse data.
• Creating a fused SpMM + Bias Add kernel. Fusion minimizes intermediate memory reads/writes, which is especially beneficial for memory-bandwidth-bound sparse computations.