[Bugfix][TVMScript] Handle LetStmt for var1 = var2 expressions

[Lunderberg]

Usually, when using TVMScript to represent a PrimFunc variable definition var_name = expr defines LetStmt with a variable named var_name bound to the expression expr. However, prior to this commit, if expr is a tir::Var, the TVMScript parser would instead silently omit the LetStmt, and rename all instances of that variable to var_name.

The root cause was in the VarTable.exist check, which erroneously returned False in all cases. This was due to a value is v check, which checked if the value was the same as the stack of maybe-shadowing values that share the same name. Replacing the 'value is vcheck with avalue in v` check resolves this issue.

This bug dates to the initial implementation of the new TVMScript parser in #12496.

[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, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[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, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

[junrushao]

@cyx-6 would you love to take a look? I very vaguely remember that the behavior was due to some backward compatibility issue that we have to support

[Lunderberg]

(This had an interesting interaction with #14318, because the accidental constant-folding could result in i == j // dynamic_param to become i == j when the dynamic parameter was 1. As a result, changing a numeric value could result in entirely different TIR being generated.)

[cyx-6]

Hi @Lunderberg, do you have a reproducible TVMScript code to illustrate the bug you are fixing? That will be helpful for us to get the point.

[junrushao]

(This had an interesting interaction with #14318, because the accidental constant-folding could result in i == j // dynamic_param to become i == j when the dynamic parameter was 1. As a result, changing a numeric value could result in entirely different TIR being generated.)

In this case, perhaps we should instead do i == T.FloorDiv(j, dynamic_param)

[Lunderberg]

@junrushao Thank you, and I can work around the issue by using an explicit T.FloorDiv, so this isn't by any means a blocking issue for me. It's more that it was an unexpected change, such as in the example below, and I wanted to avoid it becoming a similar issue for others.

@cyx-6 Certainly. The example below shows the weird interaction between the dynamic parameter and the variable name, which is what led to my initial investigation.

import tvm
from tvm.script import tir as T


@pytest.mark.parametrize("split_factor", [1, 2, 4])
def test_function(split_factor):
    @T.prim_func
    def explicit_tir_node(A: T.Buffer(16, "int32"), B: T.Buffer(16, "int32")):
        for i in T.serial(T.FloorDiv(16, split_factor)):
            chunk_start = T.Mul(i, split_factor)
            for j in T.vectorized(split_factor):
                B[chunk_start + j] = A[chunk_start + j]

    @T.prim_func
    def using_operator(A: T.Buffer(16, "int32"), B: T.Buffer(16, "int32")):
        for i in T.serial(16 // split_factor):
            chunk_start = i * split_factor
            for j in T.vectorized(split_factor):
                B[chunk_start + j] = A[chunk_start + j]

    using_operator.show()
    tvm.ir.assert_structural_equal(explicit_tir_node, using_operator)

This passes when the dynamic parameter split_factor is 2 or 4, but fails when split_factor == 1. In that case, it generates the PrimFunc shown below. For this case, the chunk_start = i * 1 was constant-folded to chunk_start = i, which was then used to rename i to chunk_start in the loop iterator.

@T.prim_func                                                        
def main(A: T.Buffer((16,), "int32"), B: T.Buffer((16,), "int32")):
    for chunk_start in range(16):
        for j in T.vectorized(1):      
            B[chunk_start + j] = A[chunk_start + j]

[Lunderberg]

Fixed a couple of CI errors, re-running to see what occurs next.

1. Explicitly call value.same_as(known_value), because value in value_stack would attempt to call PrimExpr.__bool__(value == value_stack[0]), which is not implemented.
2. Update cuda tensor intrinsic which introduced a variable binding between a T.block() and T.reads(). Using name = T.meta_var(expr) instead of name = expr follows the previous behavior of a variable rename, but done explicitly in the function definition.
3. Merged main into the PR branch to resolve conflict.

(I also opened #14399, to improve an error message encountered while debugging (2).)

[junrushao]

LGTM! Thanks for the patch!