[BUG][Runtime] `from __future__ import annotations` breaks tvm's type annotation

[oraluben]

This issue is blocking #963 . Note if M is accessed in the prim_func, e.g. T.Kernel(M, threads=128), the problem disappear. I'm still investigating, but based on the description, it seems that when the function body do not access variable from outer scope, they'll not be included in function closure and therefore the typing annotation string cannot be resolved at runtime.

https://docs.python.org/3/reference/compound_stmts.html#annotations

If the future statement from future import annotations is present, all annotations are instead stored as strings:

One possible workaround is to create T.Tensor before defining T.prim_func.

Required prerequisites

• I have read the documentation https://tilelang.com.
• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

What version of TileLang are you using?

main

System information

Problem description

As title

Reproducible example code

from __future__ import annotations

import tilelang
import tilelang.language as T


@tilelang.jit(out_idx=[-1])
def kernek_main(M):

    @T.prim_func
    def k(
        X: T.Tensor((M, )),
    ):
        with T.Kernel(128 * 128, threads=128) as (bx, ):
            pass

    return k


def main(M=8192):

    kernel = kernek_main(M)
    print(kernel.get_kernel_source())


if __name__ == "__main__":
    main()

Traceback

$ python test.py 
2025-10-20 17:56:08  [TileLang:tilelang.env:WARNING]: Loading tilelang libs from dev root: /home/yyc/repo/tilelang/build
error: Mismatched type on argument #1 when calling: `script.ir_builder.tir.Arg(0: ffi.String, 1: ffi.Object) -> ffi.Object`. Expected `ffi.Object` but got `const char*`
 --> /home/yyc/repo/tilelang/test.py:11:5
    |  
 11 |      def k(
    |      ^^^^^^
note: run with `TVM_BACKTRACE=1` environment variable to display a backtrace.

Expected behavior

No response

Additional context

No response

[oraluben]

patch in tvm to have a better error msg:

diff --git a/python/tvm/script/parser/tir/parser.py b/python/tvm/script/parser/tir/parser.py
index 467b9cc64..b3b460ba8 100644
--- a/python/tvm/script/parser/tir/parser.py
+++ b/python/tvm/script/parser/tir/parser.py
@@ -421,6 +421,7 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
                         ann = func_annotation.get(arg.arg, None)
                         if ann is None:
                             raise
+                    assert not isinstance(ann, str), f'Unresolved type hint {repr(ann)}'
                     param = T.arg(arg.arg, ann)
                     self.var_table.add(arg.arg, param)
                 self.visit_body(node.body)

[oraluben]

For upstream tvm we have similar issue:

from __future__ import annotations

from tvm.script import tir as T


def f(M=1):

    @T.prim_func
    def f(A: T.Buffer((M,), "float32")):
        pass

    return f


f()