[TIR] Keep trivial LetStmt in tir.Simplify when used in buffer decl

[Lunderberg]

Prior to this commit, any trivial let binding of var1 = var2 is inlined. However, buffer definitions are not updated, so this can result in dangling tir::Var instances. This commit updates the tir.Simplify pass to keep trivial let bindings if they are used as part of a buffer definition.

Ideally, the trivial LetStmt variable would be inlined into the buffer definition as well as other expressions. However, because a buffer may be implicitly declared, the first usage may be within a constrained context. If that happens, the simplified shape/strides expression cannot be used to update the buffer definition, as that simplification is not valid at all possible usage points of the buffer.

for i in range(n):
    elem_offset = i
    view = T.Buffer(1, data=buf, elem_offset = elem_offset)
    if i == 0:
        # First occurrence in TIR is here, where elem_offset would
        # simplify to zero.
        view[0] = 1
    else:
        # But the same buffer is used here, where elem_offset doesn't
        # simplify to zero.
        view[0] = 2

This will be resolvable after #14778 lands, requiring all buffers to be declared with DeclBuffer prior to usage.

for i in range(n):
    elem_offset = i
    # All variables used by the DeclBuffer are valid across the entire
    # body of the DeclBuffer.
    view = T.decl_buffer(1, data=buf, elem_offset = elem_offset)
    if i == 0:
        view[0] = 1
    else:
        view[0] = 2

[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}

[Lunderberg]

@tvm-bot rerun

[masahi]

Is this test relevant? Looks like a dup from #14988

[Lunderberg]

Hmm, it was a dup from #14982, which I had based this PR on top of to make allow the CI to use that bugfix. I think the "squash and merge" should remove it from this PR as part of merging, but just to be on the safe side, rebasing onto main.

[Lunderberg]

And after rebase, no changes remain in test_tir_transform_split_host_device.py.