Sparse Matrix Matrix Multiplication Support

[anchitDave]

Hi, just a quick question if FreeTensor supports sparse matrix multiplication? Also if you have some documentation for the same?
Thank You in advance!

[roastduck]

Yes and no.

Sparse matrix multiplication is not yet provided by the libop API. But since FreeTensor supports general loops and indirect indexing, it is possible to implement your own sparse matrix multiplication in FreeTensor.

Implementing a sparse matrix multiplication in single thread is completed OK. But parallelizing and optimizing for sparse operations is non-trivial. Sparse matrix multiplication can refer to sparse-matrix-mutiply-dense-matrix (SpMM) or sparse-matrix-multiplying-sparse-matrix (SpGEMM). Most of the optimizations for SpMM is supported by existing schedules provided by FreeTensor, except preprocessing or reordering of the sparse matrix. Optimizations for SpGEMM may additionally require advanced data structures like trees or hash tables in the runtime code, which is not supported by FreeTensor. Complex scheduling on sparse code may additionally require a no_deps hint to FreeTensor, to help analyzing the dependence (currently undocumented, described here).

You can find https://github.com/roastduck/FreeTensor_experiments/blob/main/gat/ours/main.py#L34 for a SpMM-like program. This is one of the programs used for Artifact Evaluation of FreeTensor's paper. (Since the written date, some of the usage of FreeTensor in this program is removed or not encouraged anymore, see the current doc for the latest usage).