[Bug][TIR][StorageRewrite] bufferload unvisited when used in builtin::address_of()

[supersat]

We encountered a bug where using T.address_of results in memory being erroneously reused and tracked it down to a bug in src/tir/transforms/storage_rewrite.cc where only the buffer index expressions are visited, and not the buffer itself. @csullivan has a fix (tlc-pack/relax@c029f74) but it needs tests and to be upstreamed.

• tir:transform

cc: @kparzysz-quic

cc @Hzfengsy @junrushao @quic-sanirudh @shingjan

[wrongtest-intellif]

Thanks for the findings!

Generally, BufferLoad in T.address_of would not happen and is just a sugar to represent buffer address. It does not represent an actual live usage to related memory. Thus in my understanding the previous implementation is by intention. Could you kindly provide some detail on how the memory being erroneously reused? Thanks!

[supersat]

I am working on a minimal test case to demonstrate the issue.

Essentially, we were writing TIR to offload some compute to an NPU. In this primfunc, we allocate some NPU RAM, copy partial tensors in to it, invoke the compute, and copy partial results out of the allocated NPU RAM. We had a rolling buffer similar to what's generated with the rolling_buffer schedule primitive, but hand-written. We invoke the NPU through a call_extern and pass in the rolling buffer offsets computed by T.address_of. The actual memory access happens inside the external function, but since TVM can't see this, it ended up reusing memory where the output buffer was reusing part of the input rolling buffer, generating nonsense results.

[tqchen]

Reading on this particular case. We can indeed go for a more conservative view or optimistic pov. Given that there is a possibility that the particular address might be shared by the underlying microkernel.

I agree that in this particular case of storage rewrite we can be conservative here. Ideally we have more explicit access pointers that marks read write with a region.

So a fix could be just remove the special handling of address of and add a comment that we decided to go conservative here for micro-kernel use-cases.