This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new bb9bc20a82 [SCRIPT] Support PEP 695 symbolic variables in Relax and
TIR (#20107)
bb9bc20a82 is described below
commit bb9bc20a8294fb19a2a40f029fe57baa546a8206
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Aug 9 22:53:09 2026 +0800
[SCRIPT] Support PEP 695 symbolic variables in Relax and TIR (#20107)
Relax and TIRx dependent type fields need symbolic integer variables to
be declared in a function signature without persisting parser-only
metadata in IR.
- Accept standard module-level `TypeVar` declarations on Python 3.10+
and bare or int-bound PEP 695 parameters on Python 3.12+.
- Treat declarations as syntax sugar for fresh int64 TIR variables used
by dependent signature fields; unused declarations disappear and no
function attrs, type-parameter fields, or runtime parameters are added.
- Parse both versioned forms through shared signature plumbing while
retaining dialect-specific annotation and body handling.
Representative TIRx Python 3.12+ output uses PEP 695 and postponed
annotations, so dependent expressions are unquoted:
```python
from __future__ import annotations
@T.prim_func(private=True)
def main[M](A: T.Buffer((M, M * T.int64(2)), "float32")):
A[0, 0] = T.float32(1.0)
```
Before Python 3.12, the portable TIRx output uses a module-level
`TypeVar` and quotes compound dependent expressions:
```python
M = TypeVar("M")
@T.prim_func(private=True)
def main(A: T.Buffer((M, "M * T.int64(2)"), "float32")):
A[0, 0] = T.float32(1.0)
```
The same declaration and printing rules apply to Relax tensor
signatures.
Validation:
- Latest-main native build completed successfully.
- Consolidated Relax and TIRx TypeVar round-trip tests passed.
- Related Relax/TIRx parser, printer, and well-formedness suites passed
661 tests with 4 expected failures.
- Touched-file pre-commit hooks and diff checks passed.
- The implementation does not change Relax well-formedness analysis.
---
include/tvm/script/printer/doc.h | 15 +++-
python/tvm/relax/script/parser/parser.py | 6 +-
python/tvm/runtime/script_printer.py | 12 ++-
python/tvm/script/parser/_core.py | 2 +-
python/tvm/script/parser/core/diagnostics.py | 2 +-
python/tvm/script/parser/core/doc.py | 7 +-
python/tvm/script/parser/core/doc_core.py | 68 +++++++++++++++++
python/tvm/script/parser/core/entry.py | 3 +-
python/tvm/script/parser/core/parser.py | 85 +++++++++++++++++++++-
python/tvm/script/printer/doc.py | 3 +
python/tvm/tirx/script/parser/parser.py | 18 ++++-
src/relax/script/printer/dependent_type.cc | 10 ++-
src/relax/script/printer/function.cc | 52 +++++++------
src/relax/script/printer/tir.cc | 7 +-
src/relax/script/printer/utils.h | 4 +
src/script/printer/doc.cc | 10 ++-
.../printer/doc_printer/python_doc_printer.cc | 5 ++
src/script/printer/ir/ir.cc | 34 ++++++++-
src/script/printer/utils.h | 62 +++++++++++++++-
src/tirx/script/printer/buffer.cc | 20 +++--
src/tirx/script/printer/function.cc | 61 ++++++++++++++--
src/tirx/script/printer/utils.h | 5 +-
tests/python/relax/test_tvmscript_parser.py | 2 +
tests/python/relax/test_tvmscript_type_vars.py | 68 +++++++++++++++++
tests/python/tirx/test_tvmscript_type_vars.py | 67 +++++++++++++++++
.../python/tvmscript/test_tvmscript_printer_doc.py | 1 +
.../test_tvmscript_printer_python_doc_printer.py | 20 +++++
27 files changed, 588 insertions(+), 61 deletions(-)
diff --git a/include/tvm/script/printer/doc.h b/include/tvm/script/printer/doc.h
index 0442852692..5e877db958 100644
--- a/include/tvm/script/printer/doc.h
+++ b/include/tvm/script/printer/doc.h
@@ -1212,6 +1212,14 @@ class FunctionDocNode : public StmtDocNode {
ffi::Optional<ExprDoc> return_type{std::nullopt};
/*! \brief The body of function. */
ffi::Array<StmtDoc> body;
+ /*!
+ * \brief The PEP 695 type parameters of the function.
+ *
+ * Possible actual types:
+ * - ExprDoc (a bare parameter like ``T``)
+ * - AssignDoc (an annotated parameter like ``T: int``)
+ */
+ ffi::Array<Doc> type_params;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -1220,7 +1228,8 @@ class FunctionDocNode : public StmtDocNode {
.def_ro("args", &FunctionDocNode::args)
.def_ro("decorators", &FunctionDocNode::decorators)
.def_ro("return_type", &FunctionDocNode::return_type)
- .def_ro("body", &FunctionDocNode::body);
+ .def_ro("body", &FunctionDocNode::body)
+ .def_ro("type_params", &FunctionDocNode::type_params);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.FunctionDoc",
FunctionDocNode, StmtDocNode);
};
@@ -1239,9 +1248,11 @@ class FunctionDoc : public StmtDoc {
* \param decorators The decorator of function.
* \param return_type The return type of function.
* \param body The body of function.
+ * \param type_params The PEP 695 type parameters of the function.
*/
explicit FunctionDoc(IdDoc name, ffi::Array<AssignDoc> args,
ffi::Array<ExprDoc> decorators,
- ffi::Optional<ExprDoc> return_type, ffi::Array<StmtDoc>
body);
+ ffi::Optional<ExprDoc> return_type, ffi::Array<StmtDoc>
body,
+ ffi::Array<Doc> type_params = {});
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionDoc, StmtDoc,
FunctionDocNode);
};
diff --git a/python/tvm/relax/script/parser/parser.py
b/python/tvm/relax/script/parser/parser.py
index bef126a884..083aa3ea93 100644
--- a/python/tvm/relax/script/parser/parser.py
+++ b/python/tvm/relax/script/parser/parser.py
@@ -32,7 +32,7 @@ from tvm.relax.script.builder.frame import BindingBlockFrame
from tvm.relax.utils import convert_to_expr
from tvm.script.ir_builder import ir as I
from tvm.script.ir_builder.base import IRBuilder
-from tvm.script.parser._core import Parser, dispatch, doc
+from tvm.script.parser._core import Parser, collect_signature_type_vars,
dispatch, doc
from tvm.tirx.script import builder as T
from .entry import (
@@ -231,9 +231,11 @@ def collect_symbolic_var_from_prelude(
def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) ->
None:
- symbolic_vars = {}
+ symbolic_vars = collect_signature_type_vars(self, node)
prim_params = set()
with self.var_table.with_frame():
+ for var_name, var in symbolic_vars.items():
+ self.var_table.add(var_name, var)
for arg in node.args.args:
if arg.annotation is None:
self.report_error(arg, "Type annotation is required for
function parameters.")
diff --git a/python/tvm/runtime/script_printer.py
b/python/tvm/runtime/script_printer.py
index 3f67e285de..973069021a 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -17,6 +17,7 @@
"""Configuration of TVMScript printer"""
import os
+import sys
from collections.abc import Sequence
from tvm_ffi import get_global_func, register_object
@@ -199,6 +200,10 @@ class Scriptable:
merged_extra: dict = {}
if extra_config is not None:
merged_extra.update(extra_config)
+ if "script.use_pep695" not in merged_extra:
+ merged_extra["script.use_pep695"] = merged_extra.get(
+ "relax.use_pep695", sys.version_info >= (3, 12)
+ )
# Only auto-switch if the caller has not already set a tirx.prefix
override.
if "tirx.prefix" not in merged_extra:
@@ -271,6 +276,11 @@ class Scriptable:
obj_to_underline: list[Object] | None = None,
obj_to_annotate: dict[Object, str] | None = None,
) -> str:
+ merged_extra = dict(extra_config or {})
+ if "script.use_pep695" not in merged_extra:
+ merged_extra["script.use_pep695"] = merged_extra.get(
+ "relax.use_pep695", sys.version_info >= (3, 12)
+ )
return _relax_script(
self,
PrinterConfig(
@@ -286,7 +296,7 @@ class Scriptable:
num_context_lines=num_context_lines,
syntax_sugar=syntax_sugar,
show_object_address=show_object_address,
- extra_config=extra_config,
+ extra_config=merged_extra,
path_to_underline=path_to_underline,
path_to_annotate=path_to_annotate,
obj_to_underline=obj_to_underline,
diff --git a/python/tvm/script/parser/_core.py
b/python/tvm/script/parser/_core.py
index b1d29d0623..d87214f22f 100644
--- a/python/tvm/script/parser/_core.py
+++ b/python/tvm/script/parser/_core.py
@@ -21,4 +21,4 @@
from .core import dispatch, doc, utils
from .core.dispatch import OpMethod, register_op
from .core.entry import parse, scan_macro
-from .core.parser import Parser
+from .core.parser import Parser, collect_signature_type_vars
diff --git a/python/tvm/script/parser/core/diagnostics.py
b/python/tvm/script/parser/core/diagnostics.py
index 4cc7c3a48b..974c5c593f 100644
--- a/python/tvm/script/parser/core/diagnostics.py
+++ b/python/tvm/script/parser/core/diagnostics.py
@@ -186,7 +186,7 @@ def findsource(obj):
if len(tokens) > 1:
name = None
if tokens[0] == "def":
- name = tokens[1].split(":")[0].split("(")[0] + "<locals>"
+ name = tokens[1].split(":")[0].split("(")[0].split("[")[0] +
"<locals>"
elif tokens[0] == "class":
name = tokens[1].split(":")[0].split("(")[0]
# pop scope if we are less indented
diff --git a/python/tvm/script/parser/core/doc.py
b/python/tvm/script/parser/core/doc.py
index 8c94d445c2..7d44dd50d6 100644
--- a/python/tvm/script/parser/core/doc.py
+++ b/python/tvm/script/parser/core/doc.py
@@ -307,7 +307,12 @@ def _register_default():
DefaultTranslator(
getattr(ast, cls_name),
from_doc,
- doc_cls._FIELDS, # pylint: disable=protected-access
+ [
+ field
+ for field in doc_cls._FIELDS # pylint:
disable=protected-access
+ if field in getattr(getattr(ast, cls_name), "_fields",
())
+ or field in {"lineno", "col_offset", "end_lineno",
"end_col_offset"}
+ ],
)
)
diff --git a/python/tvm/script/parser/core/doc_core.py
b/python/tvm/script/parser/core/doc_core.py
index aa05629d5c..74b7e57e3d 100644
--- a/python/tvm/script/parser/core/doc_core.py
+++ b/python/tvm/script/parser/core/doc_core.py
@@ -77,6 +77,7 @@ class FunctionDef(stmt):
"body",
"decorator_list",
"returns",
+ "type_params",
"lineno",
"col_offset",
"end_lineno",
@@ -90,6 +91,7 @@ class FunctionDef(stmt):
body,
decorator_list,
returns,
+ type_params,
lineno,
col_offset,
end_lineno,
@@ -101,6 +103,7 @@ class FunctionDef(stmt):
self.body = body
self.decorator_list = decorator_list
self.returns = returns
+ self.type_params = type_params
class ClassDef(stmt):
@@ -356,6 +359,67 @@ class expr(AST):
self.end_col_offset = end_col_offset
+class type_param(AST):
+ _FIELDS = ["lineno", "col_offset", "end_lineno", "end_col_offset"]
+
+ def __init__(self, lineno, col_offset, end_lineno, end_col_offset):
+ super().__init__()
+ self.lineno = lineno
+ self.col_offset = col_offset
+ self.end_lineno = end_lineno
+ self.end_col_offset = end_col_offset
+
+
+class TypeVar(type_param):
+ _FIELDS = [
+ "name",
+ "bound",
+ "default_value",
+ "lineno",
+ "col_offset",
+ "end_lineno",
+ "end_col_offset",
+ ]
+
+ def __init__(self, name, bound, default_value, lineno, col_offset,
end_lineno, end_col_offset):
+ super().__init__(lineno, col_offset, end_lineno, end_col_offset)
+ self.name = name
+ self.bound = bound
+ self.default_value = default_value
+
+
+class ParamSpec(type_param):
+ _FIELDS = [
+ "name",
+ "default_value",
+ "lineno",
+ "col_offset",
+ "end_lineno",
+ "end_col_offset",
+ ]
+
+ def __init__(self, name, default_value, lineno, col_offset, end_lineno,
end_col_offset):
+ super().__init__(lineno, col_offset, end_lineno, end_col_offset)
+ self.name = name
+ self.default_value = default_value
+
+
+class TypeVarTuple(type_param):
+ _FIELDS = [
+ "name",
+ "default_value",
+ "lineno",
+ "col_offset",
+ "end_lineno",
+ "end_col_offset",
+ ]
+
+ def __init__(self, name, default_value, lineno, col_offset, end_lineno,
end_col_offset):
+ super().__init__(lineno, col_offset, end_lineno, end_col_offset)
+ self.name = name
+ self.default_value = default_value
+
+
class BoolOp(expr):
_FIELDS = ["op", "values", "lineno", "col_offset", "end_lineno",
"end_col_offset"]
@@ -1016,6 +1080,7 @@ __all__ = [
"NotEq",
"NotIn",
"Or",
+ "ParamSpec",
"Pass",
"Pow",
"RShift",
@@ -1030,6 +1095,8 @@ __all__ = [
"Subscript",
"Try",
"Tuple",
+ "TypeVar",
+ "TypeVarTuple",
"UAdd",
"USub",
"UnaryOp",
@@ -1050,6 +1117,7 @@ __all__ = [
"mod",
"operator",
"stmt",
+ "type_param",
"unaryop",
"withitem",
]
diff --git a/python/tvm/script/parser/core/entry.py
b/python/tvm/script/parser/core/entry.py
index 418639db29..ec5b5d7ab1 100644
--- a/python/tvm/script/parser/core/entry.py
+++ b/python/tvm/script/parser/core/entry.py
@@ -17,7 +17,7 @@
"""The entry point of TVM parser."""
import inspect
-from typing import Any
+from typing import Any, TypeVar
import tvm
@@ -62,6 +62,7 @@ def _default_globals() -> dict[str, Any]:
"Tx": _tirx_tile,
"tirx": _tirx_dsl,
"Axis": _tirx_layout.Axis,
+ "TypeVar": TypeVar,
}
diff --git a/python/tvm/script/parser/core/parser.py
b/python/tvm/script/parser/core/parser.py
index 36349b26a4..771be8f3d1 100644
--- a/python/tvm/script/parser/core/parser.py
+++ b/python/tvm/script/parser/core/parser.py
@@ -17,16 +17,17 @@
"""The core parser"""
import abc
+import ast
import inspect
from collections import defaultdict
from collections.abc import Callable
from contextlib import contextmanager
-from typing import Any
+from typing import Any, TypeVar
import numpy as np
from tvm.error import DiagnosticError
-from tvm.ir import GlobalVar
+from tvm.ir import GlobalVar, Var
from tvm.runtime import Object
from ...ir_builder import IRBuilder
@@ -42,6 +43,74 @@ DEFAULT_VISIT = {
}
+def collect_signature_type_vars(parser: "Parser", node: doc.FunctionDef) ->
dict[str, Var]:
+ """Collect symbolic variables declared by function type-parameter
syntax."""
+
+ annotation_names = set()
+ parameter_names = {arg.arg for arg in node.args.args}
+
+ class AnnotationNameCollector(doc.NodeVisitor):
+ def visit_Name(self, name): # pylint: disable=invalid-name
+ annotation_names.add(name.id)
+
+ def visit_Constant(self, constant): # pylint: disable=invalid-name
+ if not isinstance(constant.value, str):
+ return
+ try:
+ expression = ast.parse(constant.value, mode="eval")
+ except SyntaxError:
+ return
+ annotation_names.update(
+ child.id for child in ast.walk(expression) if
isinstance(child, ast.Name)
+ )
+
+ name_collector = AnnotationNameCollector()
+ for arg in node.args.args:
+ name_collector.visit(arg.annotation)
+ name_collector.visit(node.returns)
+
+ symbolic_vars = {}
+ for binding_name, value in parser.var_table.get().items():
+ if (
+ binding_name not in annotation_names
+ or binding_name in parameter_names
+ or not isinstance(value, TypeVar)
+ ):
+ continue
+ if value.__name__ != binding_name:
+ parser.report_error(
+ node,
+ f"TypeVar binding {binding_name!r} must match its declared
name {value.__name__!r}",
+ )
+ if value.__constraints__ or value.__bound__ is not None:
+ parser.report_error(
+ node,
+ f"Symbolic TypeVar {binding_name!r} must not have constraints
or a bound",
+ )
+ symbolic_vars[binding_name] = Var(binding_name, "int64")
+
+ for type_param in node.type_params or []:
+ if not isinstance(type_param, doc.TypeVar):
+ parser.report_error(type_param, "Only PEP 695 TypeVar parameters
are supported")
+ # Both a bare parameter and the optional ``int`` bound declare the
+ # int64 PrimVar represented by this type parameter. The printer uses
+ # the bare spelling, while accepting the bound as explicit input.
+ if type_param.bound is not None and not (
+ isinstance(type_param.bound, doc.Name) and type_param.bound.id ==
"int"
+ ):
+ parser.report_error(
+ type_param,
+ f"Symbolic TypeVar {type_param.name!r} must be unannotated or
annotated as int",
+ )
+ if type_param.default_value is not None:
+ parser.report_error(
+ type_param,
+ f"Symbolic TypeVar {type_param.name!r} must not have a
default",
+ )
+ symbolic_vars[type_param.name] = Var(type_param.name, "int64")
+ return symbolic_vars
+
+
def _deferred(exit_f: Callable[[], None]):
"""Created context with certain exit function.
@@ -777,6 +846,18 @@ class Parser(doc.NodeVisitor):
self.report_error(node, "The parser does not understand the
decorator")
_dispatch_wrapper(func)(self, node)
+ def visit_ImportFrom(self, node: doc.ImportFrom) -> None: # pylint:
disable=invalid-name
+ """Accept postponed annotations emitted by the PEP 695 printer."""
+ is_future_annotations = (
+ node.module == "__future__"
+ and node.level == 0
+ and len(node.names) == 1
+ and node.names[0].name == "annotations"
+ and node.names[0].asname is None
+ )
+ if not is_future_annotations:
+ self.report_error(node, "Only 'from __future__ import annotations'
is supported")
+
def visit_arguments(self, node: doc.arguments) -> Any:
"""The general arguments visiting method.
diff --git a/python/tvm/script/printer/doc.py b/python/tvm/script/printer/doc.py
index 819a24a431..dd4d635e8f 100644
--- a/python/tvm/script/printer/doc.py
+++ b/python/tvm/script/printer/doc.py
@@ -462,6 +462,7 @@ class FunctionDoc(StmtDoc):
decorators: Sequence[ExprDoc]
return_type: ExprDoc | None
body: Sequence[StmtDoc]
+ type_params: Sequence[Doc]
def __init__(
self,
@@ -470,6 +471,7 @@ class FunctionDoc(StmtDoc):
decorators: list[ExprDoc],
return_type: ExprDoc | None,
body: list[StmtDoc],
+ type_params: list[Doc] | None = None,
):
self.__init_handle_by_constructor__(
_ffi_api.FunctionDoc, # type: ignore # pylint: disable=no-member
@@ -478,6 +480,7 @@ class FunctionDoc(StmtDoc):
decorators,
return_type,
body,
+ type_params or [],
)
diff --git a/python/tvm/tirx/script/parser/parser.py
b/python/tvm/tirx/script/parser/parser.py
index 324f59cdb5..ca881ac1c7 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -20,14 +20,14 @@ import ast
import contextlib
from copy import deepcopy
from functools import partial
-from typing import Any
+from typing import Any, TypeVar
import tvm
from tvm.ir import Expr, GlobalVar, PointerType, PrimType
from tvm.script.ir_builder import ir as I
from tvm.script.ir_builder.base import IRBuilder
from tvm.script.ir_builder.base import IRBuilderFrame as Frame
-from tvm.script.parser._core import Parser, dispatch, doc
+from tvm.script.parser._core import Parser, collect_signature_type_vars,
dispatch, doc
from tvm.script.parser.core.doc import from_doc
from tvm.tirx import Buffer, IterVar, Layout, buffer_data, is_buffer_var
from tvm.tirx.script import builder as T
@@ -429,9 +429,15 @@ def _eval_signature_annotation(
for child in ast.walk(expression):
if not isinstance(child, ast.Name) or not
isinstance(child.ctx, ast.Load):
continue
- if define_missing and child.id not in
self_parser.var_table.get():
+ current_value = self_parser.var_table.get().get(child.id)
+ is_shadowed_type_var = child.id in signature_dtypes and
isinstance(
+ current_value, TypeVar
+ )
+ if define_missing and (current_value is None or
is_shadowed_type_var):
# TIR match-scope indices default to int32. A later scalar
- # parameter keeps its explicitly declared dtype.
+ # parameter keeps its explicitly declared dtype. That
+ # runtime parameter also shadows a module TypeVar with the
+ # same name, such as one emitted for another function.
var = tvm.tirx.Var(child.id,
signature_dtypes.get(child.id, "int32"))
self_parser.var_table.add(child.id, var,
allow_shadowing=False)
self_parser._signature_match_vars[child.id] = var
@@ -835,6 +841,8 @@ def visit_function_def(self: Parser, node: doc.FunctionDef)
-> None:
persistent = find_decorator_annotation(node, "persistent", default=False)
self.function_annotations = None
with self.var_table.with_frame(), _signature_match_var_scope(self):
+ for name, var in collect_signature_type_vars(self, node).items():
+ self.var_table.add(name, var, allow_shadowing=False)
prim_func_ctx = T.prim_func(is_private=privacy, s_tir=s_tir,
persistent=persistent)
with prim_func_ctx:
T.func_name(node.name)
@@ -1139,6 +1147,8 @@ def visit_tvm_declare_function(self: Parser, node:
doc.FunctionDef) -> GlobalVar
ret_type = None
with self.var_table.with_frame(), _signature_match_var_scope(self):
signature_dtypes = _signature_prim_var_dtypes(node)
+ for name, var in collect_signature_type_vars(self, node).items():
+ self.var_table.add(name, var, allow_shadowing=False)
arg_annotations = []
for arg in node.args.args:
diff --git a/src/relax/script/printer/dependent_type.cc
b/src/relax/script/printer/dependent_type.cc
index 87f158d94b..af1f126fce 100644
--- a/src/relax/script/printer/dependent_type.cc
+++ b/src/relax/script/printer/dependent_type.cc
@@ -53,7 +53,15 @@ ExprDoc PrintShapeVar(const PrimExpr& e, const AccessPath&
e_p, const IRDocsifie
});
}
// Step 3. Stringify the PrimExpr if func var exists
- if (func_var_mode) {
+ bool is_bare_type_var = false;
+ if (f != nullptr && f->type_vars != nullptr) {
+ if (auto var = e.as<tirx::PrimVar>()) {
+ is_bare_type_var = f->type_vars->count(var.value().get());
+ }
+ }
+ bool use_postponed_annotations =
+ UsePEP695TypeVars(d) && f != nullptr && f->type_vars != nullptr &&
!f->type_vars->empty();
+ if (func_var_mode && !is_bare_type_var && !use_postponed_annotations) {
return ExprStringDoc(expr_doc, e_p);
}
return expr_doc;
diff --git a/src/relax/script/printer/function.cc
b/src/relax/script/printer/function.cc
index fafdad135b..869b842a10 100644
--- a/src/relax/script/printer/function.cc
+++ b/src/relax/script/printer/function.cc
@@ -16,6 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <algorithm>
+
#include "./utils.h"
namespace tvm {
@@ -51,6 +53,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
RelaxFrameNode::RegisterReflection(); }
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<relax::Function>("", [](relax::Function n, AccessPath n_p,
IRDocsifier d) -> Doc {
std::unordered_set<const VarNode*> func_vars;
+ std::unordered_set<const VarNode*> type_vars;
+ std::unordered_set<const VarNode*> prim_params;
With<RelaxFrame> f(d);
IdDoc func_name("");
@@ -64,6 +68,13 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
(*f)->AddDispatchToken(d, "relax");
(*f)->is_func = true;
(*f)->func_vars = &func_vars;
+ (*f)->type_vars = &type_vars;
+ (*f)->prim_params = &prim_params;
+ for (const Var& param : n->params) {
+ if (param->ty.as<PrimTypeNode>()) {
+ prim_params.insert(param.get());
+ }
+ }
// Step 1. Print params
ffi::Array<AssignDoc> params;
{
@@ -79,31 +90,22 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
ffi::Optional<ExprDoc> ret_type = d->AsDoc<ExprDoc>(n->ret_ty,
n_p->Attr("ret_ty"));
// Step 3. Clean up func variables
(*f)->func_vars = nullptr;
+ (*f)->type_vars = nullptr;
+ (*f)->prim_params = nullptr;
// Step 4. Print attributes
- if (!n->attrs->dict.empty()) {
- // If the function is a global function and has a global symbol,
- // then don't print the global symbol (it will be implicit from not
being private).
- // For a function without an IR module whose global symbol
- // doesn't match the function name, we should still print the global
symbol attribute.
- if (AtTopLevelFunction(d) &&
n->attrs->dict.count(tvm::attr::kGlobalSymbol) &&
-
n->attrs->dict.at(tvm::attr::kGlobalSymbol).as_or_throw<ffi::String>() ==
- func_name->name) {
- ffi::Map<ffi::String, Any> new_attrs;
- for (auto kv : n->attrs->dict) {
- if (kv.first != tvm::attr::kGlobalSymbol) {
- new_attrs.Set(kv.first, kv.second);
- }
- }
- if (!new_attrs.empty()) {
- (*f)->stmts.push_back(ExprStmtDoc(
- Relax(d, "func_attr") //
- ->Call({d->AsDoc<ExprDoc>(DictAttrs(new_attrs),
n_p->Attr("attrs"))})));
- }
- } else {
- (*f)->stmts.push_back(
- ExprStmtDoc(Relax(d, "func_attr") //
- ->Call({d->AsDoc<ExprDoc>(n->attrs,
n_p->Attr("attrs"))})));
+ ffi::Map<ffi::String, Any> printable_attrs;
+ for (const auto& [key, value] : n->attrs->dict) {
+ // A matching global symbol is implicit for a top-level function.
+ if (key == tvm::attr::kGlobalSymbol && AtTopLevelFunction(d) &&
+ value.as_or_throw<ffi::String>() == func_name->name) {
+ continue;
}
+ printable_attrs.Set(key, value);
+ }
+ if (!printable_attrs.empty()) {
+ (*f)->stmts.push_back(ExprStmtDoc(
+ Relax(d, "func_attr") //
+ ->Call({d->AsDoc<ExprDoc>(DictAttrs(printable_attrs),
n_p->Attr("attrs"))})));
}
// Step 5. Prepare the decorator (include purity if it's impure)
ExprDoc decorator = Relax(d, "function");
@@ -127,7 +129,9 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
// Step 6. Print body
ffi::Array<StmtDoc> body = PrintSeqExpr(n->body, n_p->Attr("body"), d,
/*use_ret=*/true);
(*f)->stmts.insert((*f)->stmts.end(), body.begin(), body.end());
- return HeaderWrapper(d, FunctionDoc(func_name, params, {decorator},
ret_type, (*f)->stmts));
+ auto type_var_docs = DefineTypeVarDocs(type_vars,
ffi::GetRef<Frame>((*f).get()), d);
+ return WrapFunctionDocWithTypeVars(
+ d, FunctionDoc(func_name, params, {decorator}, ret_type,
(*f)->stmts), type_var_docs);
});
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
diff --git a/src/relax/script/printer/tir.cc b/src/relax/script/printer/tir.cc
index dfd742a0b1..6ac8f11986 100644
--- a/src/relax/script/printer/tir.cc
+++ b/src/relax/script/printer/tir.cc
@@ -61,10 +61,15 @@ Doc PrintCanonicalVar(Var n, AccessPath n_p, IRDocsifier d)
{
if (f->func_vars) {
TVM_FFI_ICHECK(f->is_func);
f->func_vars->insert(n.get());
+ if (!f->prim_params->count(n.get())) {
+ f->type_vars->insert(n.get());
+ }
}
IdDoc var = d->Define(n, ffi::GetRef<Frame>(f), n->name.empty() ? "v" :
n->name);
var->source_paths.push_back(n_p);
- f->stmts.push_back(AssignDoc(var, PrintVarCreation(prim_var, n_p, d),
std::nullopt));
+ if (!f->func_vars || f->prim_params->count(n.get()) ||
!f->type_vars->count(n.get())) {
+ f->stmts.push_back(AssignDoc(var, PrintVarCreation(prim_var, n_p, d),
std::nullopt));
+ }
}
if (ffi::Optional<ExprDoc> doc = d->GetVarDoc(n)) {
return doc.value();
diff --git a/src/relax/script/printer/utils.h b/src/relax/script/printer/utils.h
index 6fe09a672b..cb7c8c4b46 100644
--- a/src/relax/script/printer/utils.h
+++ b/src/relax/script/printer/utils.h
@@ -42,6 +42,8 @@ class RelaxFrameNode : public FrameNode {
bool is_func = false;
bool module_alias_printed = false;
std::unordered_set<const VarNode*>* func_vars = nullptr;
+ std::unordered_set<const VarNode*>* type_vars = nullptr;
+ std::unordered_set<const VarNode*>* prim_params = nullptr;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -60,6 +62,8 @@ class RelaxFrame : public Frame {
n->d = d.get();
n->is_func = false;
n->func_vars = nullptr;
+ n->type_vars = nullptr;
+ n->prim_params = nullptr;
data_ = std::move(n);
}
diff --git a/src/script/printer/doc.cc b/src/script/printer/doc.cc
index a5a4e92c9a..53738361b9 100644
--- a/src/script/printer/doc.cc
+++ b/src/script/printer/doc.cc
@@ -261,13 +261,15 @@ ReturnDoc::ReturnDoc(ExprDoc value) {
}
FunctionDoc::FunctionDoc(IdDoc name, ffi::Array<AssignDoc> args,
ffi::Array<ExprDoc> decorators,
- ffi::Optional<ExprDoc> return_type,
ffi::Array<StmtDoc> body) {
+ ffi::Optional<ExprDoc> return_type,
ffi::Array<StmtDoc> body,
+ ffi::Array<Doc> type_params) {
ffi::ObjectPtr<FunctionDocNode> n = ffi::make_object<FunctionDocNode>();
n->name = name;
n->args = args;
n->decorators = decorators;
n->return_type = return_type;
n->body = body;
+ n->type_params = type_params;
this->data_ = std::move(n);
}
@@ -483,8 +485,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def("script.printer.FunctionDoc",
[](IdDoc name, ffi::Array<AssignDoc> args,
ffi::Array<ExprDoc> decorators,
- ffi::Optional<ExprDoc> return_type,
ffi::Array<StmtDoc> body) {
- return FunctionDoc(name, args, decorators,
return_type, body);
+ ffi::Optional<ExprDoc> return_type,
ffi::Array<StmtDoc> body,
+ ffi::Array<Doc> type_params = {}) {
+ return FunctionDoc(name, args, decorators,
return_type, body,
+ type_params);
});
}
diff --git a/src/script/printer/doc_printer/python_doc_printer.cc
b/src/script/printer/doc_printer/python_doc_printer.cc
index 13ccbde1ea..72fbbc1ec4 100644
--- a/src/script/printer/doc_printer/python_doc_printer.cc
+++ b/src/script/printer/doc_printer/python_doc_printer.cc
@@ -805,6 +805,11 @@ void PythonDocPrinter::PrintTypedDoc(const FunctionDoc&
doc) {
output_ << "def ";
PrintDoc(doc->name);
+ if (!doc->type_params.empty()) {
+ output_ << "[";
+ PrintJoinedDocs(doc->type_params, ", ");
+ output_ << "]";
+ }
output_ << "(";
PrintJoinedDocs(doc->args, ", ");
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index 640bc6c57e..67c42f0f31 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -64,9 +64,28 @@ struct SortableFunction {
}
};
+ffi::Optional<ffi::String> GetTypeVarDeclarationName(const StmtDoc& stmt) {
+ const auto* assign = stmt.as<AssignDocNode>();
+ if (assign == nullptr || !assign->rhs.has_value()) {
+ return std::nullopt;
+ }
+ const auto* lhs = assign->lhs.as<IdDocNode>();
+ const auto* call = assign->rhs.value().as<CallDocNode>();
+ if (lhs == nullptr || call == nullptr) {
+ return std::nullopt;
+ }
+ const auto* callee = call->callee.as<IdDocNode>();
+ if (callee == nullptr || callee->name != "TypeVar") {
+ return std::nullopt;
+ }
+ return lhs->name;
+}
+
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<IRModule>("", [](IRModule mod, AccessPath p, IRDocsifier d)
-> Doc {
std::vector<SortableFunction> functions;
+ ffi::Array<StmtDoc> type_var_decls;
+ std::unordered_set<ffi::String> declared_type_vars;
for (const auto& kv : mod->functions) {
functions.push_back(SortableFunction(kv));
}
@@ -101,6 +120,14 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
Doc doc = d->AsDoc(base_func, p->Attr("functions")->MapItem(gv));
d->cfg->binding_names.pop_back();
if (const auto* stmt_block = doc.as<StmtBlockDocNode>()) {
+ for (const StmtDoc& stmt : stmt_block->stmts) {
+ if (ffi::Optional<ffi::String> name =
GetTypeVarDeclarationName(stmt)) {
+ if (!declared_type_vars.count(name.value())) {
+ declared_type_vars.insert(name.value());
+ type_var_decls.push_back(stmt);
+ }
+ }
+ }
(*f)->stmts.push_back(stmt_block->stmts.back());
(*f)->stmts.back()->source_paths = std::move(doc->source_paths);
} else if (auto stmt = doc.as<StmtDoc>()) {
@@ -118,7 +145,12 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
<< " produced Doc type of " << doc->GetTypeKey();
}
}
- return HeaderWrapper(d, ClassDoc(module_doc, {IR(d, "ir_module")},
(*f)->stmts));
+ ClassDoc class_doc(module_doc, {IR(d, "ir_module")}, (*f)->stmts);
+ if (type_var_decls.empty()) {
+ return HeaderWrapper(d, class_doc);
+ }
+ type_var_decls.push_back(class_doc);
+ return HeaderWrapper(d, StmtBlockDoc(type_var_decls));
});
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
diff --git a/src/script/printer/utils.h b/src/script/printer/utils.h
index 99e98556b6..1ec51e3a43 100644
--- a/src/script/printer/utils.h
+++ b/src/script/printer/utils.h
@@ -29,6 +29,7 @@
#include <tvm/script/printer/ir_docsifier.h>
#include <tvm/script/printer/printer.h>
+#include <algorithm>
#include <sstream>
#include <string>
#include <unordered_set>
@@ -67,6 +68,9 @@ inline std::string Docsify(const ffi::ObjectRef& obj, const
IRDocsifier& d, cons
TVM_FFI_THROW(TypeError) << "Unexpected doc type: " << doc->GetTypeKey();
}
std::ostringstream os;
+ if (d->ir_usage.count("future_annotations")) {
+ os << "from __future__ import annotations\n\n";
+ }
if (!d->metadata.empty()) {
if (d->cfg->show_meta) {
os << "metadata = tvm.ir.load_json(\"\"\""
@@ -124,6 +128,9 @@ inline std::string DType2Str(DLDataType dtype) {
inline Doc HeaderWrapper(const IRDocsifier& d, const Doc& doc) {
if (d->ir_usage.size()) {
ffi::Array<StmtDoc> stmts;
+ if (d->ir_usage.count("type_var")) {
+ stmts.push_back(CommentDoc("from typing import TypeVar"));
+ }
if (d->ir_usage.count("ir")) {
stmts.push_back(CommentDoc("from tvm.script import ir as " +
d->cfg->ir_prefix));
}
@@ -140,12 +147,65 @@ inline Doc HeaderWrapper(const IRDocsifier& d, const Doc&
doc) {
d->cfg->GetExtraConfig<ffi::String>("relax.prefix", "R")));
}
stmts.push_back(CommentDoc(""));
- stmts.push_back(doc.as_or_throw<StmtDoc>());
+ if (const auto* stmt_block = doc.as<StmtBlockDocNode>()) {
+ stmts.insert(stmts.end(), stmt_block->stmts.begin(),
stmt_block->stmts.end());
+ } else {
+ stmts.push_back(doc.as_or_throw<StmtDoc>());
+ }
return StmtBlockDoc(stmts);
}
return doc;
}
+inline std::vector<std::pair<std::string, ExprDoc>> DefineTypeVarDocs(
+ const std::unordered_set<const VarNode*>& type_vars, const Frame& frame,
const IRDocsifier& d) {
+ std::vector<std::pair<std::string, ExprDoc>> type_var_docs;
+ type_var_docs.reserve(type_vars.size());
+ for (const VarNode* var_node : type_vars) {
+ Var var = ffi::GetRef<Var>(var_node);
+ ffi::Optional<ExprDoc> existing_doc = d->GetVarDoc(var);
+ ExprDoc var_doc = existing_doc.has_value()
+ ? existing_doc.value()
+ : d->Define(var, frame, var->name.empty() ? "v" :
var->name);
+ const auto* id_doc = var_doc.as<IdDocNode>();
+ TVM_FFI_ICHECK(id_doc != nullptr);
+ type_var_docs.emplace_back(id_doc->name, var_doc);
+ }
+ std::sort(type_var_docs.begin(), type_var_docs.end(),
+ [](const auto& lhs, const auto& rhs) { return lhs.first <
rhs.first; });
+ return type_var_docs;
+}
+
+inline bool UsePEP695TypeVars(const IRDocsifier& d) {
+ return d->cfg->GetExtraConfig<bool>("script.use_pep695",
+
d->cfg->GetExtraConfig<bool>("relax.use_pep695", false));
+}
+
+inline Doc WrapFunctionDocWithTypeVars(
+ const IRDocsifier& d, FunctionDoc function_doc,
+ const std::vector<std::pair<std::string, ExprDoc>>& type_var_docs) {
+ if (type_var_docs.empty()) {
+ return HeaderWrapper(d, function_doc);
+ }
+ if (UsePEP695TypeVars(d)) {
+ d->ir_usage.insert("future_annotations");
+ for (const auto& type_var_doc : type_var_docs) {
+ function_doc->type_params.push_back(type_var_doc.second);
+ }
+ return HeaderWrapper(d, function_doc);
+ }
+
+ d->ir_usage.insert("type_var");
+ ffi::Array<StmtDoc> stmts;
+ for (const auto& [name, var_doc] : type_var_docs) {
+ stmts.push_back(AssignDoc(
+ var_doc, IdDoc("TypeVar")->Call({LiteralDoc::Str(name,
ffi::Optional<AccessPath>())}),
+ std::nullopt));
+ }
+ stmts.push_back(function_doc);
+ return HeaderWrapper(d, StmtBlockDoc(stmts));
+}
+
/*! \brief Check if a string has multiple lines. */
inline bool HasMultipleLines(const std::string& str) {
return str.find_first_of('\n') != std::string::npos;
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index e161a03868..950b5acc34 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -30,8 +30,8 @@ namespace printer {
ffi::Map<ffi::String, ExprDoc> BufferAttrs(
tirx::BufferVar buffer, const AccessPath& buffer_p, const Frame& frame,
const IRDocsifier& d,
BufferVarDefinition var_definitions, ffi::Optional<Expr> data =
std::nullopt,
- bool stringify_undefined_shape = false,
- std::unordered_set<tirx::Var> stringify_shape_vars = {}) {
+ bool stringify_undefined_shape = false, std::unordered_set<tirx::Var>
stringify_shape_vars = {},
+ std::unordered_set<tirx::Var> stringify_compound_shape_vars = {}) {
using tvm::tirx::Var;
using tvm::tirx::VarNode;
ffi::Map<ffi::String, ExprDoc> kwargs;
@@ -88,6 +88,7 @@ ffi::Map<ffi::String, ExprDoc> BufferAttrs(
PrimExpr e = shape[i];
AccessPath e_p = shape_p->ArrayItem(i);
bool contains_new_var = false;
+ bool contains_compound_shape_var = false;
std::unordered_set<Var> vars_in_shape;
tirx::PostOrderVisit(e, [&](const ffi::ObjectRef& obj) {
if (const auto* var_node = obj.as<VarNode>()) {
@@ -95,14 +96,20 @@ ffi::Map<ffi::String, ExprDoc> BufferAttrs(
vars_in_shape.insert(var);
contains_new_var =
contains_new_var || !d->IsVarDefined(var) ||
stringify_shape_vars.count(var);
+ contains_compound_shape_var =
+ contains_compound_shape_var ||
stringify_compound_shape_vars.count(var);
}
});
if (is_new_var(e)) {
add_out_of_line_var_def(e.as_or_throw<Var>(), e_p);
}
ExprDoc result = d->AsDoc<ExprDoc>(e, e_p);
- results.push_back(stringify_undefined_shape && contains_new_var ?
ExprStringDoc(result, e_p)
- :
result);
+ bool is_bare_compound_shape_var =
+ e.as<VarNode>() &&
stringify_compound_shape_vars.count(e.as_or_throw<Var>());
+ bool stringify_compound_expr = contains_compound_shape_var &&
!is_bare_compound_shape_var;
+ results.push_back((stringify_undefined_shape && contains_new_var) ||
stringify_compound_expr
+ ? ExprStringDoc(result, e_p)
+ : result);
// A quoted shape expression defines every Var it contains. Do not quote
// later dimensions merely because they reuse a Var introduced here.
for (const Var& var : vars_in_shape) {
@@ -335,10 +342,11 @@ ExprDoc BufferDecl(const tirx::BufferVar& buffer, const
ffi::String& method,
}
ExprDoc BufferAttn(const tirx::BufferVar& buffer, const AccessPath& p, const
Frame& frame,
- const IRDocsifier& d, std::unordered_set<tirx::Var>
stringify_shape_vars) {
+ const IRDocsifier& d, std::unordered_set<tirx::Var>
stringify_shape_vars,
+ std::unordered_set<tirx::Var>
stringify_compound_shape_vars) {
ffi::Map<ffi::String, ExprDoc> attrs =
BufferAttrs(buffer, p, frame, d, BufferVarDefinition::MatchBuffer,
std::nullopt, true,
- std::move(stringify_shape_vars));
+ std::move(stringify_shape_vars),
std::move(stringify_compound_shape_vars));
if (!attrs.count("dtype")) {
attrs.Set("dtype", LiteralDoc::DataType(buffer->dtype->dtype,
p->Attr("dtype")));
}
diff --git a/src/tirx/script/printer/function.cc
b/src/tirx/script/printer/function.cc
index f6ba953b86..554dba11c5 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -34,6 +34,42 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
d->SetCommonPrefix(func, [](const ffi::ObjectRef& obj) {
return obj->IsInstance<tirx::VarNode>() ||
obj->IsInstance<tirx::BufferTypeNode>();
});
+ std::unordered_set<const VarNode*> runtime_params;
+ for (const tirx::Var& param : func->params) {
+ runtime_params.insert(param.get());
+ }
+ std::unordered_set<const VarNode*> type_vars;
+ auto collect_type_vars = [&](const PrimExpr& expr) {
+ for (const tirx::Var& var : tirx::UndefinedVars(expr)) {
+ const auto* var_ty_node = var->ty.as<PrimTypeNode>();
+ if (var_ty_node == nullptr) {
+ continue;
+ }
+ PrimType var_ty(var_ty_node->dtype);
+ if (!runtime_params.count(var.get()) && var_ty.IsScalar() &&
+ var_ty.MatchesElementType(DLDataTypeCode::kDLInt, 64)) {
+ type_vars.insert(var.get());
+ }
+ }
+ };
+ for (const tirx::Var& param : func->params) {
+ if (!param->ty.as<tirx::BufferTypeNode>()) {
+ continue;
+ }
+ tirx::BufferVar buffer(param);
+ for (const PrimExpr& extent : buffer->shape) {
+ collect_type_vars(extent);
+ }
+ for (const PrimExpr& stride : buffer->strides) {
+ collect_type_vars(stride);
+ }
+ collect_type_vars(buffer->elem_offset);
+ for (const PrimExpr& address : buffer->allocated_addr) {
+ collect_type_vars(address);
+ }
+ }
+ auto type_var_docs = DefineTypeVarDocs(type_vars,
ffi::GetRef<Frame>((*f).get()), d);
+ bool use_postponed_annotations = UsePEP695TypeVars(d) &&
!type_vars.empty();
int n_args = func->params.size();
// Step 1. Handle `func->params`
ffi::Array<AssignDoc> args;
@@ -55,21 +91,28 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
if (var->ty.as<tirx::BufferTypeNode>()) {
tirx::BufferVar buffer(var);
std::unordered_set<tirx::Var> stringify_shape_vars;
+ std::unordered_set<tirx::Var> stringify_compound_shape_vars;
std::unordered_set<tirx::Var> shape_vars;
for (const PrimExpr& shape : buffer->shape) {
tirx::PostOrderVisit(shape, [&](const ffi::ObjectRef& obj) {
if (const auto* shape_var_node = obj.as<tirx::VarNode>()) {
tirx::Var shape_var = ffi::GetRef<tirx::Var>(shape_var_node);
shape_vars.insert(shape_var);
- if (!bound_signature_vars.count(shape_var)) {
+ bool is_type_var = type_vars.count(shape_var.get());
+ if (!use_postponed_annotations &&
!bound_signature_vars.count(shape_var) &&
+ !is_type_var) {
stringify_shape_vars.insert(shape_var);
}
+ if (!use_postponed_annotations && is_type_var) {
+ stringify_compound_shape_vars.insert(shape_var);
+ }
}
});
}
IdDoc lhs = DefineBuffer(buffer, *f, d);
ExprDoc annotation =
- BufferAttn(buffer, var_p->Attr("ty"), *f, d,
std::move(stringify_shape_vars));
+ BufferAttn(buffer, var_p->Attr("ty"), *f, d,
std::move(stringify_shape_vars),
+ std::move(stringify_compound_shape_vars));
args.push_back(AssignDoc(lhs, std::nullopt, annotation));
for (const tirx::Var& shape_var : shape_vars) {
bound_signature_vars.insert(shape_var);
@@ -177,12 +220,14 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
ffi::Array<ExprDoc> pos_args;
decorator = std::move(decorator->Call(pos_args, kwargs_keys,
kwargs_values));
}
- return HeaderWrapper(d, FunctionDoc(
- /*name=*/func_name,
- /*args=*/args,
- /*decorators=*/{decorator},
- /*return_type=*/ret_type,
- /*body=*/(*f)->stmts));
+ return WrapFunctionDocWithTypeVars(d,
+ FunctionDoc(
+ /*name=*/func_name,
+ /*args=*/args,
+ /*decorators=*/{decorator},
+ /*return_type=*/ret_type,
+ /*body=*/(*f)->stmts),
+ type_var_docs);
});
TVM_REGISTER_SCRIPT_AS_REPR(tirx::PrimFuncNode, ReprPrintTIR);
diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h
index 2c38ec59e2..b2246a667f 100644
--- a/src/tirx/script/printer/utils.h
+++ b/src/tirx/script/printer/utils.h
@@ -316,10 +316,13 @@ ExprDoc BufferDecl(const tirx::BufferVar& buffer, const
ffi::String& method,
* \param d The IRDocsifier
* \param stringify_shape_vars Variables whose first shape use must be
stringified. The set is
* passed by value so entries can be consumed as dimensions are emitted.
+ * \param stringify_compound_shape_vars Variables whose compound shape
expressions must be
+ * stringified while their bare-name uses remain direct.
* \return The ExprDoc corresponding to the buffer declaration
*/
ExprDoc BufferAttn(const tirx::BufferVar& buffer, const AccessPath& p, const
Frame& frame,
- const IRDocsifier& d, std::unordered_set<tirx::Var>
stringify_shape_vars = {});
+ const IRDocsifier& d, std::unordered_set<tirx::Var>
stringify_shape_vars = {},
+ std::unordered_set<tirx::Var> stringify_compound_shape_vars
= {});
/*!
* \brief Print the creation of a Var
diff --git a/tests/python/relax/test_tvmscript_parser.py
b/tests/python/relax/test_tvmscript_parser.py
index b91a1e4544..1e7cddb116 100644
--- a/tests/python/relax/test_tvmscript_parser.py
+++ b/tests/python/relax/test_tvmscript_parser.py
@@ -995,6 +995,8 @@ def test_call_tir_with_tir_var():
Y[vi] = X[vi]
_check(Module)
+ portable = Module.script(show_meta=True,
extra_config={"script.use_pep695": False})
+ tvm.ir.assert_structural_equal(Module, tvm.script.from_source(portable))
def test_call_tir_with_grad():
diff --git a/tests/python/relax/test_tvmscript_type_vars.py
b/tests/python/relax/test_tvmscript_type_vars.py
new file mode 100644
index 0000000000..990a9582f7
--- /dev/null
+++ b/tests/python/relax/test_tvmscript_type_vars.py
@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import sys
+from typing import TypeVar
+
+import tvm
+import tvm.testing
+from tvm.script import relax as R
+
+M = TypeVar("M")
+UNUSED_GENERIC = TypeVar("UNUSED_GENERIC", bound=int)
+
+
+def test_type_vars_roundtrip():
+ @R.function(private=True)
+ def func(
+ x: R.Tensor((M, "M * 2"), "float32"),
+ ) -> R.Tensor((M, "M * 2"), "float32"):
+ return x
+
+ script = func.script()
+ if sys.version_info >= (3, 12):
+ assert script.startswith("from __future__ import annotations\n\n")
+ assert "def main[M](" in script
+ assert 'R.Tensor((M, M * 2), dtype="float32")' in script
+ typed = tvm.script.from_source(
+ """
[email protected](private=True)
+def func[M: int](x: R.Tensor((M, M * 2), "float32")):
+ return x
+"""
+ )
+ tvm.ir.assert_structural_equal(func, typed)
+ else:
+ assert "from __future__ import annotations" not in script
+ assert 'M = TypeVar("M")' in script
+ assert 'R.Tensor((M, "M * 2"), dtype="float32")' in script
+
+ portable = func.script(extra_config={"relax.use_pep695": False})
+ assert "from __future__ import annotations" not in portable
+ assert 'M = TypeVar("M")' in portable
+ assert 'R.Tensor((M, "M * 2"), dtype="float32")' in portable
+ assert "M = T.int64()" not in script
+ assert "UNUSED_GENERIC" not in script
+ assert [param.name for param in func.params] == ["x"]
+ assert not hasattr(func, "type_params")
+ assert func.attrs.get("relax.type_vars") is None
+ tvm.ir.assert_structural_equal(func, tvm.script.from_source(script))
+ tvm.ir.assert_structural_equal(func, tvm.script.from_source(portable))
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/tirx/test_tvmscript_type_vars.py
b/tests/python/tirx/test_tvmscript_type_vars.py
new file mode 100644
index 0000000000..c3d13f122b
--- /dev/null
+++ b/tests/python/tirx/test_tvmscript_type_vars.py
@@ -0,0 +1,67 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import sys
+
+import tvm
+import tvm.testing
+
+
+def test_type_vars_roundtrip():
+ func = tvm.script.from_source(
+ """
+M = TypeVar("M")
+UNUSED = TypeVar("UNUSED")
+
[email protected]_func(private=True)
+def func(A: T.Buffer((M, M * 2), "float32")):
+ A[0, 0] = T.float32(1)
+"""
+ )
+
+ script = func.script()
+ if sys.version_info >= (3, 12):
+ assert script.startswith("from __future__ import annotations\n\n")
+ assert "def main[M](" in script
+ assert 'T.Buffer((M, M * T.int64(2)), "float32")' in script
+ typed = tvm.script.from_source(
+ """
[email protected]_func(private=True)
+def func[M: int](A: T.Buffer((M, M * 2), "float32")):
+ A[0, 0] = T.float32(1)
+"""
+ )
+ tvm.ir.assert_structural_equal(func, typed)
+ else:
+ assert "from __future__ import annotations" not in script
+ assert 'M = TypeVar("M")' in script
+
+ portable = func.script(extra_config={"script.use_pep695": False})
+ assert "from __future__ import annotations" not in portable
+ assert 'M = TypeVar("M")' in portable
+ assert 'T.Buffer((M, "M * T.int64(2)"), "float32")' in portable
+ assert "UNUSED" not in script
+ assert "M = T.int64()" not in script
+ assert len(func.params) == 1
+ assert not hasattr(func, "type_params")
+ assert func.attrs.get("tirx.type_vars") is None
+ tvm.ir.assert_structural_equal(func, tvm.script.from_source(script))
+ tvm.ir.assert_structural_equal(func, tvm.script.from_source(portable))
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/tvmscript/test_tvmscript_printer_doc.py
b/tests/python/tvmscript/test_tvmscript_printer_doc.py
index 18c3cec267..571f75e35a 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_doc.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_doc.py
@@ -478,6 +478,7 @@ def test_function_doc(args, decorators, return_type, body):
assert list(doc.decorators) == decorators
assert doc.return_type == return_type
assert list(doc.body) == body
+ assert list(doc.type_params) == []
@pytest.mark.parametrize(
diff --git
a/tests/python/tvmscript/test_tvmscript_printer_python_doc_printer.py
b/tests/python/tvmscript/test_tvmscript_printer_python_doc_printer.py
index 9cf7baa302..0fe90a1a83 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_python_doc_printer.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_python_doc_printer.py
@@ -804,6 +804,26 @@ def test_print_function_doc(args, decorators, body,
return_type, expected):
assert to_python_script(doc) == format_script(expected) # test
+def test_print_function_doc_with_type_params():
+ doc = FunctionDoc(
+ IdDoc("func"),
+ [],
+ [],
+ None,
+ [ReturnDoc(IdDoc("T"))],
+ type_params=[
+ IdDoc("T"),
+ AssignDoc(IdDoc("U"), rhs=None, annotation=IdDoc("int")),
+ ],
+ )
+ assert to_python_script(doc) == format_script(
+ """
+ def func[T, U: int]():
+ return T
+ """
+ )
+
+
def get_func_doc_for_class(name):
args = [
AssignDoc(IdDoc("x"), rhs=None, annotation=IdDoc("int")),
