[Bug] TVMScript does not parse correctly with Python 3.8

[guberti]

While fixing #13330, I discovered that I could not get tests/python/unittest/test_tir_te_extern_primfunc.py to pass on main. When I ran it using Python 3.8.10, it would give the error:

==================================== ERRORS ====================================
____ ERROR collecting tests/python/unittest/test_tir_te_extern_primfunc.py _____
tests/python/unittest/test_tir_te_extern_primfunc.py:36: in <module>
    def func_1(A: T.Buffer[(16,), "float32"], C: T.Buffer[(1,), "float32"]):
python/tvm/script/tir/prim_func.py:40: in prim_func
    result = from_source(input_func)
python/tvm/script/parser.py:1368: in from_source
    result = to_ast(input_func, TVMDiagnosticCtx(), parser)
../.local/lib/python3.8/site-packages/synr/compiler.py:731: in to_ast
    prog = compiler.compile_module(program_ast)
../.local/lib/python3.8/site-packages/synr/compiler.py:100: in compile_module
    new_func = self.compile_def(stmt)
../.local/lib/python3.8/site-packages/synr/compiler.py:174: in compile_def
    params = self.compile_args_to_params(args)
../.local/lib/python3.8/site-packages/synr/compiler.py:164: in compile_args_to_params
    ty = self.compile_type(arg.annotation)
../.local/lib/python3.8/site-packages/synr/compiler.py:622: in compile_type
    return self._expr2type(self.compile_expr(ty))
../.local/lib/python3.8/site-packages/synr/compiler.py:605: in _expr2type
    assert isinstance(
E   AssertionError: Expected subscript call to have lhs of Var, but it is <class 'synr.ast.Attr'>
=========================== short test summary info ============================
ERROR tests/python/unittest/test_tir_te_extern_primfunc.py - AssertionError: ...
!!!!!!!!!!!!!!!!!!!! Interrupted: 1 error during collection !!!!!!!!!!!!!!!!!!!!
=============================== 1 error in 0.22s ===============================

As the error says, Python does not like the type declarations here:

@T.prim_func
def func_1(A: T.Buffer[(16,), "float32"], C: T.Buffer[(1,), "float32"]):
    # Do stuff

Solution

After some debugging, I realized the test ran correctly with Python 3.7.15. It seems like this is an issue with Python 3.8 (and possibly future versions too), but I haven't looked closely enough to be sure. In the meantime, I'll work around it by using Python 3.7.

• tir:parser

[junrushao]

Hey @guberti, this issue should be fixed in the latest HEAD. Testing script:

from tvm.script import tir as T

@T.prim_func
def func_1(
    A: T.Buffer[(16,), "float32"],
    C: T.Buffer[(1,), "float32"],
):
    ...

func_1.show()

Output in python 3.8:

# from tvm.script import tir as T
@T.prim_func
def func(A: T.Buffer[16, "float32"], C: T.Buffer[1, "float32"]):
    # body
    T.evaluate(0)

[guberti]

@junrushao just tested it, the issue is fixed on the latest HEAD. Your testing script and the unit test tests/python/unittest/test_tir_te_extern_primfunc.py both work. Thanks!