How to Design a Local Variable

Development
3

TVMScript also provide doc.NodeVisitor and doc.NodeTransformer for AST level mutations. In our private fork we collect all lhs variables with multiple Assign or AugAssign bindings. They are implicitly transformed to scalar buffer load & store during parsing to structure same as @spectrometerHBH’s answer.

from tvm.script import tir as T

@T.prim_func
def function(X: T.Buffer([16], "int32")):
    v = 0
    X[0] = v
    for i in range(10):
        v = i
        X[i + 1] = v
    
print(function)

# default tir parser output
# different assignment to `v` treat as new let bindings
@T.prim_func
def function(X: T.Buffer((16,), "int32")):
    v: T.int32 = 0
    X[0] = v
    for i in range(10):
        v_1: T.int32 = i
        X[i + 1] = v_1

# scalar transformed output
@T.prim_func
def function(X: T.Buffer((16,), "int32")):
    # with T.block("root"):
    v = T.alloc_buffer((1,), "int32")
    v[0] = 0
    X[0] = v[0]
    for i in range(10):
        v[0] = i
        X[i + 1] = v[0]

In llvm-based backends, we could also introduce new storage scope or someother annotations to ensure the scalar buffer allocation map to alloca and finally transform to SSA value via mem2reg.

One important note is though this behavior is more “pythonic”, it break the concise scoping behaviour for default IR parsing. And it may be hard to keep round trip property. So from my understanding it could be useful when we use TVMScript as a programming language, but may not be proper or default behaviour for TVMScript’s origin serving purpose.

1 Like