A significant driver of progress in deep learning has been advances in computational resources. While those resources are often limited, the is a trend to replace dense computation in DNN with sparse computation for speeding up / saving memory to enable larger models. For example: [neural network pruning](https://github.com/he-y/Awesome-Pruning), [sparse transformer](https://openai.com/blog/sparse-transformer/). Also some new workloads like GNN relies on sparse support. It would be great if TVM can represent sparse computation workload.
There exists some sparse support in TVM already. Overall, it use the dense tensors to describe the sparse CSR/BSR tensors based on the existing Tensor DSL like here https://github.com/apache/incubator-tvm/blob/master/topi/python/topi/nn/sparse.py. However, this approach have some obvious drawbacks: - it is quite tedious to describe the sparse computation, while you need to deal with the indexing manually. - It does not provide proper abstractions for sparse kernel scheduling. This RFC would like to discuss how to add native sparse support in TVM. ## Sparse Workloads Here are some sparse workloads that we would like to keep in mind and take into consideration during design. - **Graph Neural Networks**: GNN is a type of Neural Network which directly operates on the graph structure, which has gained increasing popularity in various domains, including social network, knowledge graph, recommender system, and even life science. The graph data are often sparse, so that there exist urgent demand that optimizing sparse kernels for GNN workloads, like: sparse matrix-matrix multiplication (SPMM), Sampled dense-dense matrix product (SDDMM). segment_sum, segment_min, segment_max, segment_mm, etc. - **Block Sparse**: Even though sparse operations need less compute and memory relative to their dense counterparts, the speed-up observed by using sparse operations is less than expected on different hardware platforms. The block sparse representation (BSR) would be more friendly for hardwares and easier to be optimized. There also exist some works to induce block sparsity in RNNs/Transformer by pruning blocks of weights. >From the above workloads, we can summary some requirements that our sparse >support need to achieve: - It should be able to represent *common sparse formats*: CSR, RSR, RSR, etc. - Although most workloads are focused on 2D sparse matrics, but it would be better that if it can represent *multiple dimension tensor* so that fit with the original TVM Tensor abstraction. After some investigation, we found that the tree hierarchy representation used by TACO and ExTensor is a good candidate. ## The Tree Hierarchy Representation The tree hierachy representation can represent tensors of any order, by constructing formats from a bounded number of primitives, e.g., specifying whether each dimension is dense of sparse. (TACO also supports many other types like *range*, *hash*, etc. but we can expand it in the future depends on the demand.) With this approach, a CSR matrix can be represented as `SparseTensor([Dense, Sparse])`, RSR as `SparseTensor([Sparse, Dense])`, BSR as `SparseTensor([Dense, Dense, Sparse, Sparse])`. We can found that a general/sparse tensor is actually composed by several dense arrays with the tree hierarchy representation: - An array `A_val` is used to represent the non-zero elements of tensor A. - For every dense axis: an integer `Ai_size` is used to represent the size of tensor A's i-th dimension. - For every sparse axis: two index arrays, `Ai_pos` and `Ai_idx`, together form a segmented vector with one segment per entry in the previous dimension (parent node in the tree). The `Ai_idx` array stores all the non-zero indices in the dimension, while the `Ai_pos` array stores the location in the idx array where each segment begins. ### Understanding the Representation with Examples Here we will show with a 2D case to understand how the sparse tensor is represented under different formats: ``` example tensor: [ a, 0, b, c, 0, 0, 0, 0, d, 0, 0, e, ] ``` ``` Format: [Dense, Dense] Storage: axis 0 A0_size = 3 axis 1 A1_size = 4 values of A A_val = [a, 0, b, c, 0, 0, 0, 0, d, 0, 0, e] Access: produce B { for (i, 0, m) { for (j, 0, n) { B[((i*n) + j)] = A[((i*n) + j)] } } } ``` ``` Format: [Dense, Sparse] Storage: axis 0 A0_size = 3 axis 1 A1_pos = [0, 2, 2, 5] A1_idx = [0, 3, 0, 2, 3] values of A A_val = [a, b, c, d, e] Access: for (i, 0, A0_size) { for (j, A1_pos[i], A1_pos[i+1]) { idx = {i, A1_idx[j]} val = A_vals[j]; } } ``` ``` Format: [Sparse, Dense] Storage: axis 0 A0_pos = [0, 2] A0_idx = [0, 2] axis 1 A1_size = 4 A_val = [a, 0, b, c, 0, 0, 0, 0, d, 0, 0, e] Access: for (i, A0_pos[0], A0_pos[1]) { for (j, 0, A1_size) { idx = {A0_idx[i], j} val = A_vals[A0_idx[i] * A1_size + j]; } } ``` ``` Format: [Sparse, Sparse] Storage: axis 0 A0_pos = [0, 2] A0_idx = [0, 2] axis 1 A1_pos = [0, 2, 5] A1_idx = [0, 3, 0, 2, 3] values of A A_val = [a, b, c, d, e] Access: for (i, A0_pos[0], A0_pos[1]) { for (j, A1_pos[i], A1_pos[i+1]) { idx = {A0_idx[i], A1_idx[j]} val = A_vals[j]; } } ``` ## Implementation ### Format Declaration A tuple-like data structure can be introduced to declare the format with the sparsity on dimensions: `SparseFormat([Dense, Sparse])` ### Sparse Tensor As a counterpart of original dense `Tensor`, a `SparseTensor` class is a symbolic representation for sparse tensor, which is used during sparse code generation, composed by `pos_arrs`, `idx_arrs`, `val_arr`. ### DSL Enhancement To enhance existed DSL with ability to declare sparse computation, here are some approaches we can try. - Option 1: Adding New Sparse Operations ``` # demo code snippet import tvm.sparse as tvmsp # declare sparse format in_sformat = tvmsp.sparse_format([Dense, Sparse]) # declare dense format out_sformat = tvmsp.sparse_format([Dense, Dense]) # computation declaration n = tvm.var("n") m = tvm.var("m") A = tvmsp.placeholder((m, n), sformat=in_sformat, name='A') B = tvmsp.compute(A.shape, lambda i, j: A[i, j], sformat=out_sformat, name='B') s = tvm.create_schedule(B.op) ir = tvm.lower(s, [A, B], simple_mode=True) ``` This approach adds new operators like `SparsePlaceholder`, `SparseComputeOp`, with additional `sformat` field compared with origial operations. - Option 2: Enhancing `decl_buffer` ``` import tvm.sparse as tvmsp # declare sparse format sformat = tvmsp.sparse_format([Dense, Sparse]) # declare dense format sformat = tvmsp.decl_dense(3) # computation declaration n = tvm.var("n") m = tvm.var("m") A = tvm.placeholder((m, n), name='A') B = tvm.compute(A.shape, lambda i, j: A[i, j], name='B') s = tvm.create_schedule(B.op) Ab = tvm.decl_buffer(A.shape, A.dtype, sformat=in_sformat, name="Ab") Bb = tvm.decl_buffer(B.shape, B.dtype, sformat=out_sformat, name="Bb") ir = tvm.lower(s, [A, B], binds={A: Ab, B: Bb}, simple_mode=True) print(ir) ``` We can also enhances `decl_buffer` that let user can declare a sparse buffer and bind it with tensor while building. One thing is this approach separate sparse property with computation declaration. It sounds cool that we can represent dense and sparse computation with the same declaration, but it also means that we don't have such information until scheduling. ## Extension for Scheduling TODO -- You are receiving this because you are subscribed to this thread. Reply to this email directly or view it on GitHub: https://github.com/apache/incubator-tvm/issues/4332
