Can SparseTIR support operations other than multiplication and addition?

uwsampl / SparseTIR

SparseTIR: Sparse Tensor Compiler for Deep Learning

https://sampl.cs.washington.edu/SparseTIR/

Apache License 2.0

131 stars 14 forks source link

Closed gxiaotian closed 3 weeks ago

gxiaotian commented 4 months ago

Hello,

I am new to SparseTIR and wonder whether it can support operations other than multiplication and addition.

For example, the below expression is valid SparseTIR or not?

with T.sp_iter([I, J], "SR") as [i, j]:
       with T.init():
              Y[i, j] = 0.0
        Y[i] = Y[i, j] + (X[i] - X[j]) ** 2

Thanks!

yzh119 commented 3 months ago

Sorry for the very late reply, I just noticed this issue.

Yes of course, but I'm not sure if the parser can parse ** correctly. A safer way would be T.power(X[i] - X[j], 2)