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 3ddd9c249e186a69bbe7c60b5291b6a75ddfc18b Author: Tianqi Chen <[email protected]> AuthorDate: Thu Oct 1 22:48:50 2026 +0000 Align printer expectations with explicit blocks and generic display policy --- .../relax/script/test_relax_script_printer.py | 50 +++++++++++++---- tests/python/relax/test_analysis_well_formed.py | 2 +- .../s_tir/script/test_s_tir_script_printer.py | 65 +++++++++++----------- .../python/tirx/script/test_tirx_script_buffers.py | 28 +++++----- .../python/tirx/script/test_tirx_script_printer.py | 22 ++++---- 5 files changed, 98 insertions(+), 69 deletions(-) diff --git a/tests/python/relax/script/test_relax_script_printer.py b/tests/python/relax/script/test_relax_script_printer.py index 173fbe23ba..a21098dc2b 100644 --- a/tests/python/relax/script/test_relax_script_printer.py +++ b/tests/python/relax/script/test_relax_script_printer.py @@ -116,6 +116,8 @@ def test_func(): _assert_print( TestModule["foo"], """ +from __future__ import annotations + # from tvm.script import relax as R @R.function @@ -133,10 +135,9 @@ def test_module(): from __future__ import annotations # from tvm.script import ir as I -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis -# from tvm.script import s_tir as Ts # from tvm.script import relax as R +# from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @I.ir_module class Module: @@ -181,6 +182,8 @@ def test_function(): _assert_print_lines( func, """ +from __future__ import annotations + # from tvm.script import relax as R @R.function @@ -230,7 +233,7 @@ def test_function_dependent_shape_source_spans(): definition, underline = render(cast_path.attr("dtype")) dtype_start = definition.index(dtype_literal) assert underline[dtype_start : dtype_start + len(dtype_literal)] == "^" * len(dtype_literal) - assert underline.strip() == "^" * len(dtype_literal) + assert underline.strip() == "^" * len(expression) definition, underline = render(cast_path.attr("value")) variable_start = definition.index(expression) + expression.rindex("n") @@ -248,6 +251,8 @@ def test_lone_private_function(): _assert_print_lines( func, """ +from __future__ import annotations + # from tvm.script import relax as R @R.function(private=True) @@ -271,6 +276,8 @@ def test_extern_func(): _assert_print_lines( obj, """ +from __future__ import annotations + # from tvm.script import ir as I # from tvm.script import relax as R @@ -342,6 +349,8 @@ def test_nested_function(): _assert_print_lines( NestedFunction, """ +from __future__ import annotations + # from tvm.script import ir as I # from tvm.script import relax as R @@ -498,9 +507,10 @@ def test_prim_value(): _assert_print_lines( func, """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis +from __future__ import annotations + # from tvm.script import relax as R +# from tvm.script import tirx as T @R.function def func() -> T.int64: @@ -871,11 +881,12 @@ def test_module_cross_func_call(): _assert_print_lines( TestModule, """ +from __future__ import annotations + # from tvm.script import ir as I -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis -# from tvm.script import s_tir as Ts # from tvm.script import relax as R +# from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @I.ir_module class Module: @@ -895,11 +906,12 @@ class Module: _assert_print_lines( module_str, """ +from __future__ import annotations + # from tvm.script import ir as I -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis -# from tvm.script import s_tir as Ts # from tvm.script import relax as R +# from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @I.ir_module class Module: @@ -926,6 +938,8 @@ def test_assert_op(): _assert_print_lines( AssertOpMod, """ +from __future__ import annotations + # from tvm.script import ir as I # from tvm.script import relax as R @@ -950,6 +964,8 @@ def test_print(): _assert_print_lines( PrintMod, """ +from __future__ import annotations + # from tvm.script import ir as I # from tvm.script import relax as R @@ -974,6 +990,8 @@ def test_private_function(): _assert_print_lines( AddMod, """ +from __future__ import annotations + # from tvm.script import ir as I # from tvm.script import relax as R @@ -1027,6 +1045,8 @@ def test_directly_construct_private_funcs(): _assert_print_lines( obj, """ +from __future__ import annotations + # from tvm.script import ir as I # from tvm.script import relax as R @@ -1072,6 +1092,8 @@ def test_reused_extern_func(): _assert_print_lines( func, """ +from __future__ import annotations + # from tvm.script import relax as R @R.function @@ -1100,6 +1122,8 @@ def test_inline_extern_func(): _assert_print_lines( func, """ +from __future__ import annotations + # from tvm.script import relax as R @R.function @@ -1148,6 +1172,8 @@ def test_hide_inferable_ty(): _assert_print_lines( func.script(show_all_ty=False), """ +from __future__ import annotations + # from tvm.script import relax as R @R.function diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py index daf834999a..a6fad5df7d 100644 --- a/tests/python/relax/test_analysis_well_formed.py +++ b/tests/python/relax/test_analysis_well_formed.py @@ -669,7 +669,7 @@ def test_impure_in_dataflow_block(): # The throwing form surfaces the offending impure call in its message. with pytest.raises(Exception) as excinfo: rx.analysis.well_formed(mod) - assert "R.print" in str(excinfo.value) + assert 'I.Call("relax.print", ["{}", x], ty=R.Tuple())' in str(excinfo.value) def test_well_formed_function(): 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 a6a2d00bca..bea1c3b4fb 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 @@ -279,9 +279,8 @@ def test_bind(): _assert_print( obj, """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func(private=True) def main(): @@ -315,9 +314,8 @@ def test_remap(): v4, v5 = Ts.axis.remap("RS", [i4, i5]) expected_output = """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func def main(): @@ -361,20 +359,21 @@ def test_root_block(): TB.evaluate(0) expected_output = """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func def main(): - # with Ts.sblock("root"): - buffer = Ts.sblock_alloc_buffer((128, 128), "float32") - for i in range(128): - for j in range(128): - with Ts.sblock(""): - Ts.reads() - Ts.writes() - T.evaluate(0)""" + with Ts.sblock("root"): + Ts.reads() + Ts.writes() + v = Ts.sblock_alloc_buffer((128, 128), "float32") + for i in range(128): + for j in range(128): + with Ts.sblock(""): + Ts.reads() + Ts.writes() + T.evaluate(0)""" _assert_print(root_block_implicitly.with_attr("global_symbol", "main"), expected_output) _assert_print(root_block_explicitly.with_attr("global_symbol", "main"), expected_output) @@ -409,9 +408,10 @@ def test_prim_func_different_symbol(): TB.evaluate(0) expected_output = """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis +from __future__ import annotations + # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")): @@ -504,9 +504,10 @@ def test_predicated_load_store(): ) expected_output = """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis +from __future__ import annotations + # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")): @@ -548,9 +549,10 @@ def test_predicated_buffer_load_store(): ).with_attr("s_tir", True) expected_output = """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis +from __future__ import annotations + # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func(private=True) def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")): @@ -577,9 +579,10 @@ def test_predicated_scalable_load_store(): ) expected_output = """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis +from __future__ import annotations + # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func def func(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")): @@ -617,9 +620,10 @@ def test_vload_with_explicit_scalable_data_type(): B[0 : TB.vscale() * 4] = A.vload([TB.Ramp(0, 1, TB.vscale() * 4)], dtype="float32xvscalex4") expected_output = """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis +from __future__ import annotations + # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")): @@ -637,9 +641,10 @@ def test_vectorize_llvm_pure_intrin(): ) expected_output = """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis +from __future__ import annotations + # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")): @@ -660,9 +665,10 @@ def test_func_with_loop_jumps(): break expected_output = """ -# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis +from __future__ import annotations + # from tvm.script import s_tir as Ts +# from tvm.script import tirx as T @Ts.prim_func def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")): @@ -2077,7 +2083,6 @@ def test_annotation_multi_access_paths(): assert ( result == """# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis # from tvm.script import s_tir as Ts @Ts.prim_func @@ -2105,7 +2110,6 @@ def test_annotate_from_multi_obj(): assert ( result == """# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis # from tvm.script import s_tir as Ts @Ts.prim_func @@ -2138,7 +2142,6 @@ def test_disable_concise_scoping_when_scope_annotated(): assert ( result == """# from tvm.script import tirx as T -# from tvm.tirx.layout import Axis # from tvm.script import s_tir as Ts @Ts.prim_func diff --git a/tests/python/tirx/script/test_tirx_script_buffers.py b/tests/python/tirx/script/test_tirx_script_buffers.py index 152c94ddd6..625c083c47 100644 --- a/tests/python/tirx/script/test_tirx_script_buffers.py +++ b/tests/python/tirx/script/test_tirx_script_buffers.py @@ -442,7 +442,7 @@ def test_buffer_local_ir(): # Round-trip code = func.script() assert ( - 'buffer_1 = T.decl_buffer((2,), "float16", data=buffer.data, ' + 'v_1 = T.decl_buffer((2,), "float16", data=v.data, ' 'scope="local", layout="default")' ) in code assert from_source(code).script() == code @@ -518,11 +518,11 @@ def test_buffer_local_physical_order(): code = func.script() assert ( - 'buffer_1 = T.decl_buffer((32,), "float32", data=buffer.data, ' + 'v_1 = T.decl_buffer((32,), "float32", data=v.data, ' 'scope="local", layout="default")' ) in code assert ( - 'buffer_2 = T.decl_buffer((4, 8), "float32", data=buffer.data, ' + 'v_2 = T.decl_buffer((4, 8), "float32", data=v.data, ' 'scope="local", layout="default")' ) in code assert from_source(code).script() == code @@ -553,13 +553,13 @@ def test_buffer_local_layout_overrides_roundtrip(): assert not b_custom.ty.layout.is_trivial() code = func.script() - storage_line = next(line for line in code.splitlines() if "buffer_1 =" in line) - custom_line = next(line for line in code.splitlines() if "buffer_2 =" in line) + storage_line = next(line for line in code.splitlines() if "v_1 =" in line) + custom_line = next(line for line in code.splitlines() if "v_2 =" in line) assert ( - 'T.decl_buffer((32,), "float32", data=buffer.data, scope="local", layout=' + 'T.decl_buffer((32,), "float32", data=v.data, scope="local", layout=' ) in storage_line assert ( - 'T.decl_buffer((2, 4), "float32", data=buffer.data, scope="local", ' + 'T.decl_buffer((2, 4), "float32", data=v.data, scope="local", ' 'layout=' ) in custom_line assert_structural_equal(func, from_source(code)) @@ -608,7 +608,7 @@ def test_buffer_local_compose_layout_printer_roundtrip(): assert [int(dim) for dim in b_buf.ty.shape] == [64] assert b_buf.ty.layout.is_trivial() code = func.script() - local_line = next(line for line in code.splitlines() if "buffer =" in line) + local_line = next(line for line in code.splitlines() if "v =" in line) assert 'T.decl_buffer((64,), "float32", data=A.data, scope="local", layout=' in local_line parsed = from_source(code) assert_structural_equal(func, parsed) @@ -661,8 +661,8 @@ def test_buffer_local_physical_span_includes_gaps_and_offset(): assert int(b_storage.ty.layout.apply(1, shape=list(b_storage.ty.shape))["m"]) == 5 code = func.script() - storage_line = next(line for line in code.splitlines() if "buffer_3 =" in line) - assert 'T.decl_buffer((2,), "float32", data=buffer.data, scope="local", layout=' in storage_line + storage_line = next(line for line in code.splitlines() if "v_3 =" in line) + assert 'T.decl_buffer((2,), "float32", data=v.data, scope="local", layout=' in storage_line assert_structural_equal(func, from_source(code)) assert from_source(code).script() == code @@ -685,16 +685,16 @@ def test_buffer_local_printer_is_stable_with_multiple_aliases(): expected = func.script() assert ( - 'buffer_1 = T.decl_buffer((32,), "float32", data=buffer.data, ' + 'v_1 = T.decl_buffer((32,), "float32", data=v.data, ' 'scope="local", layout="default")' ) in expected assert ( - 'buffer_2 = T.decl_buffer((4, 8), "float32", data=buffer.data, ' + 'v_2 = T.decl_buffer((4, 8), "float32", data=v.data, ' 'scope="local", layout="default")' ) in expected - storage_line = next(line for line in expected.splitlines() if "buffer_3 =" in line) + storage_line = next(line for line in expected.splitlines() if "v_3 =" in line) assert ( - 'T.decl_buffer((32,), "float32", data=buffer.data, scope="local", layout=' + 'T.decl_buffer((32,), "float32", data=v.data, scope="local", layout=' ) in storage_line for _ in range(20): parsed = from_source(expected) diff --git a/tests/python/tirx/script/test_tirx_script_printer.py b/tests/python/tirx/script/test_tirx_script_printer.py index 09d1e03ed7..6e455f934e 100644 --- a/tests/python/tirx/script/test_tirx_script_printer.py +++ b/tests/python/tirx/script/test_tirx_script_printer.py @@ -228,7 +228,7 @@ def test_allocate(): _assert_print( obj.body, """ -buffer = T.alloc_buffer((128, 128), "float32", layout="default") +v = T.alloc_buffer((128, 128), "float32", layout="default") T.evaluate(1) """, ) @@ -246,8 +246,8 @@ def test_allocate_with_decl_buffer_sugar(): _assert_print( obj.body, """ -buffer = T.alloc_buffer((128, 128), "float32", layout="default") -buffer_1 = T.decl_buffer((128, 128), "float32", data=buffer.data, layout="default") +v = T.alloc_buffer((128, 128), "float32", layout="default") +v_1 = T.decl_buffer((128, 128), "float32", data=v.data, layout="default") T.evaluate(1) """, ) @@ -265,9 +265,9 @@ def test_allocate_with_decl_buffer_sugar_multi_usage(): _assert_print( obj.body, """ -buffer = T.alloc_buffer((128, 128), "float32", layout="default") -buffer_1 = T.decl_buffer((128, 128), "float32", data=buffer.data, layout="default") -T.evaluate(buffer.data) +v = T.alloc_buffer((128, 128), "float32", layout="default") +v_1 = T.decl_buffer((128, 128), "float32", data=v.data, layout="default") +T.evaluate(v.data) """, ) @@ -283,9 +283,9 @@ def test_allocate_with_decl_buffer_no_sugar_mismatch(): _assert_print( obj.body, """ -buffer = T.alloc_buffer((128, 128), "float32", layout="default") -buffer_1 = T.decl_buffer((256, 256), "float32", data=buffer.data, layout="default") -T.evaluate(buffer.data) +v = T.alloc_buffer((128, 128), "float32", layout="default") +v_1 = T.decl_buffer((256, 256), "float32", data=v.data, layout="default") +T.evaluate(v.data) """, ) @@ -302,8 +302,8 @@ def test_decl_buffer(): _assert_print( obj.body, """ -v: T.handle("float32", "global") = T.handle("float32", "global") -buffer = T.decl_buffer((10, 10), "float32", data=v, layout="default") +v_1: T.handle("float32", "global") = T.handle("float32", "global") +v = T.decl_buffer((10, 10), "float32", data=v_1, layout="default") T.evaluate(1) """, )
