This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tqchen/printer-refinement in repository https://gitbox.apache.org/repos/asf/tvm.git
commit c2f712c67eb8116223cbbc02a10a07d120058d92 Author: Tianqi Chen <[email protected]> AuthorDate: Thu Oct 1 19:49:22 2026 +0000 [REFACTOR][Script] Always emit eligible function type parameters --- python/tvm/runtime/script_printer.py | 9 ---- src/script/printer/ir/utils.cc | 3 +- .../script/test_relax_script_dynamic_shape.py | 45 ++++++++----------- .../s_tir/script/test_s_tir_script_printer.py | 8 ++-- .../python/tirx/script/test_tirx_script_buffers.py | 4 +- .../tirx/script/test_tirx_script_dynamic_shape.py | 52 +++++++++------------- 6 files changed, 48 insertions(+), 73 deletions(-) diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py index 4fe8739dd8..2fe649927e 100644 --- a/python/tvm/runtime/script_printer.py +++ b/python/tvm/runtime/script_printer.py @@ -17,7 +17,6 @@ """Configuration of TVMScript printer""" import os -import sys from collections.abc import Sequence from tvm_ffi import get_global_func, register_object @@ -203,10 +202,6 @@ class Scriptable: # Keep the historical interactive display header. Direct printer calls # retain executable imports for standalone parser round-trips. merged_extra.setdefault("ir.comment_imports", True) - if "ir.use_pep695" not in merged_extra: - merged_extra["ir.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: @@ -278,10 +273,6 @@ class Scriptable: obj_to_annotate: dict[Object, str] | None = None, ) -> str: merged_extra = dict(extra_config or {}) - if "ir.use_pep695" not in merged_extra: - merged_extra["ir.use_pep695"] = merged_extra.get( - "relax.use_pep695", sys.version_info >= (3, 12) - ) return _relax_script( self, PrinterConfig( diff --git a/src/script/printer/ir/utils.cc b/src/script/printer/ir/utils.cc index a2f01b05b0..4674395bcb 100644 --- a/src/script/printer/ir/utils.cc +++ b/src/script/printer/ir/utils.cc @@ -237,10 +237,9 @@ void FinalizeFunctionDefinitions(DocTranslatorObj* d, const ffi::Dict<Var, IdDoc auto module = d->GetOrCreateExtraState<ffi::Optional<ffi::Dict<ffi::String, ffi::List<GlobalInfo>>>>( "ir.global_info_map"); - bool use_pep695 = !module.has_value() && d->GetExtraConfig<bool>("ir.use_pep695", false); auto pending = d->GetImplicitDefs(); std::vector<Var> parameters; - if (use_pep695) { + if (!module.has_value()) { for (const auto& [var, id] : signature) { if (pending.count(var) && var->ty.as<PrimTypeNode>()) parameters.push_back(var); } diff --git a/tests/python/relax/script/test_relax_script_dynamic_shape.py b/tests/python/relax/script/test_relax_script_dynamic_shape.py index 095b64d50b..4051941cfb 100644 --- a/tests/python/relax/script/test_relax_script_dynamic_shape.py +++ b/tests/python/relax/script/test_relax_script_dynamic_shape.py @@ -19,8 +19,6 @@ from __future__ import annotations -import sys - import pytest import tvm @@ -46,31 +44,26 @@ def test_type_vars_roundtrip(): 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 - assert "M = T.int64()" not in script - typed = tvm.script.from_source( - """ + 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 + assert "M = T.int64()" not in script + typed = tvm.script.from_source( + """ @R.function(private=True) def func[M: int](x: R.Tensor((M, M * 2), "float32")): return x """, - extra_vars={"I": tvm.script.ir, "R": tvm.script.relax}, - ) - tvm.ir.assert_structural_equal(func, typed) - else: - assert "from __future__ import annotations" not in script - assert 'M = I.dynamic("M", dtype="int64")' in script - assert "M = T.int64()" not 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 = I.dynamic("M", dtype="int64")' in portable - assert 'R.Tensor((M, M * 2), dtype="float32")' in portable - assert "M = T.int64()" not in portable + extra_vars={"I": tvm.script.ir, "R": tvm.script.relax}, + ) + tvm.ir.assert_structural_equal(func, typed) + assert 'M = I.dynamic("M", dtype="int64")' not in script + + repeated = func.script() + assert "from __future__ import annotations" in repeated + assert "def main[M](" in repeated + assert 'R.Tensor((M, M * 2), dtype="float32")' in repeated + assert "M = T.int64()" not in repeated assert "UNUSED_GENERIC" not in script assert [param.name for param in func.params] == ["x"] assert not hasattr(func, "type_params") @@ -80,7 +73,7 @@ def func[M: int](x: R.Tensor((M, M * 2), "float32")): ) tvm.ir.assert_structural_equal( func, - tvm.script.from_source(portable, extra_vars={"I": tvm.script.ir, "R": tvm.script.relax}), + tvm.script.from_source(repeated, extra_vars={"I": tvm.script.ir, "R": tvm.script.relax}), ) @@ -229,11 +222,11 @@ def test_call_tir_with_tir_var(): Y[vi] = X[vi] _check(Module) - portable = Module.script(show_meta=True, extra_config={"ir.use_pep695": False}) + repeated = Module.script(show_meta=True) tvm.ir.assert_structural_equal( Module, tvm.script.from_source( - portable, + repeated, extra_vars={ "I": tvm.script.ir, "R": tvm.script.relax, diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py b/tests/python/s_tir/script/test_s_tir_script_printer.py index 2640bed191..89c5a9dcda 100644 --- a/tests/python/s_tir/script/test_s_tir_script_printer.py +++ b/tests/python/s_tir/script/test_s_tir_script_printer.py @@ -46,9 +46,9 @@ def test_prim_func_symbolic_buffer_param_roundtrip(): .with_attr("s_tir", True) ) - source = func.script(extra_config={"ir.use_pep695": False}) + source = func.script() assert "T.Buffer((n + 1, n)" in source - assert source.index('n = I.dynamic("n", dtype="int32")') < source.index("T.evaluate(n)") + assert source.index("def main[n: T.int32](") < source.index("T.evaluate(n)") tvm.ir.assert_structural_equal( tvm.script.from_source( source, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx, "Ts": tvm.script.s_tir} @@ -66,9 +66,9 @@ def test_prim_func_compound_buffer_shape_first_use_roundtrip(): .with_attr("s_tir", True) ) - source = func.script(extra_config={"ir.use_pep695": False}) + source = func.script() assert "T.Buffer((T.max(n, 1),)" in source - assert source.index('n = I.dynamic("n", dtype="int32")') < source.index("T.evaluate(n)") + assert source.index("def main[n: T.int32](") < source.index("T.evaluate(n)") tvm.ir.assert_structural_equal( tvm.script.from_source( source, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx, "Ts": tvm.script.s_tir} diff --git a/tests/python/tirx/script/test_tirx_script_buffers.py b/tests/python/tirx/script/test_tirx_script_buffers.py index df5d2a557c..152c94ddd6 100644 --- a/tests/python/tirx/script/test_tirx_script_buffers.py +++ b/tests/python/tirx/script/test_tirx_script_buffers.py @@ -325,8 +325,8 @@ def test_buffer_shape_repeated_var_prints_out_of_line(): buffer = tvm.tirx.decl_buffer((n + n,), name="A") func = tvm.tirx.PrimFunc([buffer], tvm.tirx.Evaluate(0)) - code = func.script(extra_config={"ir.use_pep695": False}) - assert 'n = I.dynamic("n", dtype="int32")' in code + code = func.script() + assert "def main[n: T.int32](" in code assert_structural_equal(func, from_source(code)) diff --git a/tests/python/tirx/script/test_tirx_script_dynamic_shape.py b/tests/python/tirx/script/test_tirx_script_dynamic_shape.py index 5fa9a11aef..3278d2758e 100644 --- a/tests/python/tirx/script/test_tirx_script_dynamic_shape.py +++ b/tests/python/tirx/script/test_tirx_script_dynamic_shape.py @@ -19,8 +19,6 @@ from __future__ import annotations -import sys - import pytest import tvm @@ -133,31 +131,27 @@ def test_type_vars_roundtrip(): 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", layout="default")' in script - assert "M = T.int64()" not in script - typed = tvm.script.from_source( - """ + assert script.startswith("from __future__ import annotations\n\n") + assert "def main[M](" in script + assert 'T.Buffer((M, M * T.int64(2)), "float32", layout="default")' in script + assert "M = T.int64()" not in script + typed = tvm.script.from_source( + """ @T.prim_func(private=True) def func[M: int](A: T.Buffer((M, M * 2), "float32")): A[0, 0] = T.float32(1) """, - extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}, - ) - tvm.ir.assert_structural_equal(func, typed) - else: - assert "from __future__ import annotations" not in script - assert 'M = I.dynamic("M", dtype="int64")' in script - assert "M = T.int64()" not in script - - portable = func.script(extra_config={"ir.use_pep695": False}) - assert "from __future__ import annotations" not in portable - assert 'M = I.dynamic("M", dtype="int64")' in portable - assert 'T.Buffer((M, M * T.int64(2)), "float32", layout="default")' in portable + extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}, + ) + tvm.ir.assert_structural_equal(func, typed) + assert 'M = I.dynamic("M", dtype="int64")' not in script + + repeated = func.script() + assert "from __future__ import annotations" in repeated + assert "def main[M](" in repeated + assert 'T.Buffer((M, M * T.int64(2)), "float32", layout="default")' in repeated assert "UNUSED" not in script - assert "M = T.int64()" not in portable + assert "M = T.int64()" not in repeated assert len(func.params) == 1 assert not hasattr(func, "type_params") assert func.attrs.get("tirx.type_vars") is None @@ -166,7 +160,7 @@ def func[M: int](A: T.Buffer((M, M * 2), "float32")): ) tvm.ir.assert_structural_equal( func, - tvm.script.from_source(portable, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}), + tvm.script.from_source(repeated, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}), ) @@ -178,18 +172,16 @@ def test_dynamic_int32_roundtrip(): A[0] = T.float32(1) source = func.script() - if sys.version_info >= (3, 12): - assert "def main[n: T.int32](" in source - else: - assert 'n = I.dynamic("n", dtype="int32")' in source - portable = func.script(extra_config={"ir.use_pep695": False}) - assert 'n = I.dynamic("n", dtype="int32")' in portable + assert "def main[n: T.int32](" in source + assert 'n = I.dynamic("n", dtype="int32")' not in source + repeated = func.script() + assert "def main[n: T.int32](" in repeated tvm.ir.assert_structural_equal( func, tvm.script.from_source(source, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}) ) tvm.ir.assert_structural_equal( func, - tvm.script.from_source(portable, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}), + tvm.script.from_source(repeated, extra_vars={"I": tvm.script.ir, "T": tvm.script.tirx}), )
