search

apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.42k stars 3.4k forks source link

[Transform] Modify FuseTIR pass to propagate buffer attributes #17075

Closed quic-sanirudh closed 3 weeks ago

quic-sanirudh commented 1 month ago

Arguments of a fused TIR PrimFunc generated from a fused relax function do not retain all the buffer attributes from their original PrimFuncs as the buffers are created from the StructInfo of the Relax vars. This patch collects a mapping of relax vars to its corresponding TIR buffers in a fused relax function and uses that info to propagate its buffer attributes such as axis_separators and storage_scope

quic-sanirudh commented 1 month ago

cc @Hzfengsy @Lunderberg @slyubomirsky

quic-sanirudh commented 3 weeks ago

@Lunderberg I think I've addressed your comments. When you get a chance, could you please take a look?

quic-sanirudh commented 3 weeks ago

Thanks for taking the time to review and provide feedback.