[TIR] Handle implicit buffers in PrimFunc specialization

[Lunderberg]

Prior to this commit, the tvm::tir::Specialize function would replace the specialized parameter when it occurs in the shape, strides, or element offset of a buffer, but only when the buffer occurs in specific locations. As a result, a tir::Buffer in the PrimFuncNode::buffer_map and BlockNode would be updated, but tir::Buffer objects in a BufferLoad or BufferStore node could be left unmodified and therefore contain undefined variables.

This commit updates PrimFunc::Specialize to update buffer objects in BufferLoad and BufferStore.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @junrushao, @quic-sanirudh, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[tqchen]

Shall we have wellform check to ensure all buffer are being declared somewhere?

[Lunderberg]

Since we don't currently have a requirement for a tir::Buffer to have a specific declaration prior to use, I don't think the use of Specialize should add that as an additional requirement. That said, as soon as the DeclBuffer nodes are required for buffer declarations, it would be better to have that as part of a well-formed check, and the handling of implicit buffers would no longer be needed.

[tqchen]

I think it is a good idea to require DeclBuffer, makes analysis and other things simpler. Given DeclBuffer is already in i think it is a good idea to make that req

[Lunderberg]

I agree, though I think it's a longer-term item. @vinx13 , I think last we checked, the DeclBuffer requirement described in RFC#70 was low priority and not actively developed. Is that still the case?

[tqchen]

I think we should start to enforce DeclBuffer to reduce the need to be compact to more complicated form in the passes

[Lunderberg]

Sounds good, and closing this PR.

[Lunderberg]

Cleaned up a dev branch into draft PR #14778, which would enforce DeclBuffer usage.