uwsampl / SparseTIR

SparseTIR: Sparse Tensor Compiler for Deep Learning
https://sampl.cs.washington.edu/SparseTIR/
Apache License 2.0
131 stars 13 forks source link

[Tracking Issue] Specify the dimension order of allocated buffers for reverse cache read/write #54

Closed yzh119 closed 1 year ago

yzh119 commented 1 year ago

Problem

Currently, the reverse cache read/write will allocate a temporary buffer whose dimension order was determined by the order of iter vars in the block.

However, user might customize the storage of allocated buffers (e.g. the default dimension order might be [N, M] but the user needs an [M, N] storage instead).

Solution

Add an argument dim_order=[0, 2, 1] which accepts a list of integers, which specifies the order of dimensions in allocated buffers. The numbers in provided dimension order list indicate the block itervar index.