search

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

[Bugfix] Use `reverse_cache_read` for spmm example #92

Closed yzh119 closed 1 year ago

yzh119 commented 1 year ago

The Bug

This PR fixes issue #90 , which is because our CompactBufferRegion pass cannot infer the extent of C_local when there are non-affine expressions:

allocate(C_local_8: Pointer(local float32), float32, [(min(1710902, ((max(I_1_3_indices_data_1[((i_1_3_1: int32*8) + threadIdx.y_3: int32)], I_1_3_indices_data_1[0]) + 1) - min(I_1_3_indices_data_1[((i_1_3_1*8) + threadIdx.y_3)], I_1_3_indices_data_1[0])))*2)]), storage_scope = local;

which should be:

allocate(C_local_8: Pointer(local float32), float32, [2]), storage_scope = local;

Solution

This PR uses reverse_cache_read (we are upstreaming its generalized form reindex_cache_read/write, see https://github.com/apache/tvm/pull/14161) whose generated buffer does not rely on CompactBufferRegion to determine its extent.