Is sparse kernel supported in the tensor expression language?

Questions
#1

I am wondering if there is a way for me to generate a sparse linear algebra code using TVM tensor expression language. It seems like relay.nn has sparse-related APIs like sparse_dense, but tvm.te does not. If I cannot generate sparse kernel using te, is there a way for me to use relay to manually generate scheduled sparse kernel (e.g., tiled)?

Thank you!

#2

CC: @tkonolige who is great at sparse kernels

1 Like
#3

@kwmaeng I’ve written the sparse_dense kernel for GPUs. It was a little bit of an arduous process, but here are my takeaways:

  • Using te only works for some sparse kernels. Sparse kernels are often written as functions over the input tensor. Unfortunately, te requires you to write your kernels as functions over indices in the output tensor. If you can restructure you kernel to be output driven, then you can get around this issue, but other wise your out of luck.
  • Even If you can write your kernel in te, you may not be able to schedule it. Currently, you cannot apply any scheduling to tensors that are used to control the bounds of a loop. Also, if you have a statement of the form x[y[i]] (x, y are tensors, i is an index), you cannot apply any transformations to pull y[i] out of x (for caching typically).
  • Given the above two points, you’ll probably have to write your kernel in TIR. Currently the best way to do this is IRBuilder (tvmscript is still a work in progress and I didn’t have as much luck using it). You can see the sparse_dense kernel I wrote here https://github.com/apache/incubator-tvm/blob/main/python/tvm/topi/cuda/sparse.py#L158 for an example of how to use IRBuilder.
3 Likes