[Bug] TVMScript fails with LetStmt

[liangW-intellif]

Hi there,
I encountered an error that when generating TVMScript with LetStmt, the generated script cannot be parsed if the var in LetStmt appears the first time in the IR.

import tvm
from tvm.script import tir as T
from tvm import te


def test_let_roundtrip():
    x = te.var("x")
    y = te.var("y")
    let0 = tvm.tir.LetStmt(x, 0, tvm.tir.Evaluate(0))
    let1 = tvm.tir.LetStmt(y, 1, tvm.tir.Evaluate(0))
    body = tvm.tir.SeqStmt([let0, let1])
    func = tvm.tir.PrimFunc([], body)
    print(func.script())
    # Parse fail
    roundtrip = tvm.script.from_source(func.script())
    tvm.ir.assert_structural_equal(func, roundtrip, True)

The above case generates TVMScript as below

@T.prim_func
def main():
    with T.let(x, 0):
        T.evaluate(0)
    y: T.int32 = 1
    T.evaluate(0)

But when parsing this TVMScript, it reports an error that 'x' is an undefined variable so that the parse failed. The TVMScript with LetStmt is not round-trippable in this case.

test_let_roundtrip.py::test_let_roundtrip # from tvm.script import tir as T

@T.prim_func
def main():
    with T.let(x, 0):
        T.evaluate(0)
    y: T.int32 = 1
    T.evaluate(0)
error: Undefined variable: x
 --> <str>:5:16
   |  
 5 |      with T.let(x, 0):
   |                 ^     
FAILED

For the semantic of LetStmt, I think'x' should be identified as a definition, the script should be parsed successly.

Please let me know if there is any confusion, thanks very much!

Environment

latest TVM mainline

@junrushao @wrongtest-intellif

[junrushao]

CC @cyx-6 would you like to take a look? Thanks!

[cyx-6]

Thanks for reporting this bug. I believe that this is due to the bug of printing LetStmt. Since we forget to print the x = T.var("int32") before the with T.let(x, 0):. This bug will be fixed in the #13900. And more details are in the description of that PR too. Hope this will solve your issue.

[junrushao]

Closed via #13900