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

[oraluben]

Background

We want to keep python-3.8 compatibility for a while. One major blocker is the new type hinting for containers like list[str] since python 3.9. We try to use from __future__ import annotations to enable that feature in 3.8.

Detail background of from __future__ import annotations

from __future__ import annotations enables lazy evaluation of type annotation, and makes all annotations to be string.

from __future__ import annotations will be deprecated and removed in the future (not before 3.13 EOL).

Even after python 3.14 with annotationlib.get_annotations, string annotations from from __future__ import annotations could not be resolved as expected.

https://docs.python.org/3/howto/annotations.html#manually-un-stringizing-stringized-annotations

---

Issue

Ideally, when variables used in type annotations (e.g. M in X: T.Tensor((M, ))) is used in the inner function, it will be captured in closure of inner function, and can be resolved later.
But when it's not, the type annotation could not be solved and remain string.

---

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.

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.

Possible solution

1. Disallow all use of from __future__ import annotations
   Pros:

   1. This statement is not encouraged and is planned to be removed

   Cons

   1. Need some other way to maintain/test 3.8 support
   2. We don't have control of user code

2. Manually access variable to bound to closure
   Pros:

   1. Minimal changes to framework

   Cons:

   1. Need to manually add a useless statement, might be hard to be aware of

3. When creating prim_func, manually read stack to read locals of upper frame, and add them in annotation resolving
   Pros:

   1. Seems to works perfectly in simple cases

   Cons:

   1. Hard to maintain
   2. edge case?

[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()

[LeiWang1999]

cc @Hzfengsy @kurisu6912 , that's interesting..

[kurisu6912]

In frontend v2 #1120 , we experimentally suggest a prim func generator based expression.

In both tvm and tilelang, T.Tensor is treated as placeholders. What T.prim_func does is just extracting the argument annotations and passing them as the argument of the prim func implicitly. In the new frontend, I am trying to explicitly allow this manner, as a workaround for annotations.

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()

[Hzfengsy]

As 3.8 is deprecated, from __future__ import annotations is no longer needed

[kurisu6912]

@Hzfengsy are you sure? I see even python 3.14 has from __future__ import annotations
https://docs.python.org/3/library/__future__.html

[oraluben]

from __future__ import annotations will be deprecated and removed in the future (not before 3.13 EOL).

It's not deprecated at this point, and seems will be in the future.

[Hzfengsy]

@Hzfengsy are you sure? I see even python 3.14 has from future import annotations

You are right, from __future__ import annotations will exist. But my opinion is that we don't need to use it after 3.9 or 3.10, if my understanding is correct.

Please correct me if I missed anything. But if it's true, letting users do not use from __future__ import annotations might be a solution.

[kurisu6912]

Thanks, you are right.

Python 3.14 will introduce a native support for annotations, all annotations are lazily evaluated python objects, not strings. Since from __future__ import annotations will eventually being deprecated and removed, we should not use this.

That's very interesting, thanks very much.

[kurisu6912]

@oraluben if from __future__ import annotations is really required, please try the @T.prim_func(generator=True) method.

tilelang/tilelang/language/v2/builder.py

Lines 609 to 613 in 1768cbe

	>>> @prim_func(generator=True)
	... def parameterized_kernel(size: int = 128):
	... # function body using size parameter
	... pass
	>>> kernel = parameterized_kernel(size=256)

Or you can disable it, since it will be eventually deprecated and removed.