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
The following commit(s) were added to refs/heads/tqchen/printer-refinement by
this push:
new 6ffc67eedd Keep complete scripts executable and parameter occurrences
precise
6ffc67eedd is described below
commit 6ffc67eedda97270912c4d2d7abe4ac6a39c1471
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 1 23:18:20 2026 +0000
Keep complete scripts executable and parameter occurrences precise
---
src/script/printer/printer.cc | 33 +++++++++++++++---
.../relax/script/test_relax_script_printer.py | 20 ++++++-----
.../s_tir/script/test_s_tir_script_printer.py | 40 ++++++++++++----------
.../tirx/script/test_tirx_script_basic_usage.py | 6 ++--
4 files changed, 64 insertions(+), 35 deletions(-)
diff --git a/src/script/printer/printer.cc b/src/script/printer/printer.cc
index 8269f0745e..a8ca135d9f 100644
--- a/src/script/printer/printer.cc
+++ b/src/script/printer/printer.cc
@@ -235,9 +235,34 @@ void RecoverDocPaths(
ffi::reflection::ForEachFieldInfo(
TVMFFIGetTypeInfo(object->type_index()), [&](const TVMFFIFieldInfo*
field) {
if (ffi::String(field->name) == "source_paths") return;
- RecoverDocPaths(ffi::reflection::FieldGetter(field)(object),
context, origins, ir_paths,
- active,
- value.as<FunctionDocNode>() &&
ffi::String(field->name) == "args");
+ ffi::Any children = ffi::reflection::FieldGetter(field)(object);
+ bool parameters = value.as<FunctionDocNode>() &&
ffi::String(field->name) == "args";
+ if (parameters) {
+ auto function_origin = origins.find(object);
+ auto args = children.as<ffi::Array<AssignDoc>>();
+ auto params =
+ function_origin == origins.end()
+ ? std::nullopt
+ : ReflectedField(function_origin->second,
"params").as<ffi::Array<ffi::Any>>();
+ bool matched = args && params && args->size() == params->size();
+ for (size_t i = 0; matched && i < args->size(); ++i) {
+ auto origin = origins.find((*args)[i].get());
+ matched =
+ origin != origins.end() && origin->second.get() ==
(*params)[i].as<ffi::Object>();
+ }
+ if (matched) {
+ // Formal argument order identifies its occurrence even when the
+ // same object is also referenced by another parameter's type.
+ for (size_t i = 0; i < args->size(); ++i) {
+ std::vector<AccessPath> parameter_paths;
+ for (const AccessPath& path : context)
+
parameter_paths.push_back(path->Attr("params")->ArrayItem(i));
+ RecoverDocPaths((*args)[i], parameter_paths, origins,
ir_paths, active, true);
+ }
+ return;
+ }
+ }
+ RecoverDocPaths(children, context, origins, ir_paths, active,
parameters);
});
}
active->erase(object);
@@ -388,7 +413,7 @@ ffi::String Script(const ffi::ObjectRef& obj, const
PrinterConfig& config) {
ffi::String alias = canonical == "ir" ? config->ir_prefix
:
config->GetExtraConfig<ffi::String>(key, fallback);
ffi::String import = "from tvm.script import " + canonical + " as " +
std::string(alias);
- if (config->GetExtraConfig<bool>("ir.comment_imports", false)) {
+ if (config->GetExtraConfig<bool>("ir.comment_imports", false) &&
!config->show_meta) {
header.push_back(CommentDoc(import));
} else {
header.push_back(ffi::String(std::string(import) + "\n"));
diff --git a/tests/python/relax/script/test_relax_script_printer.py
b/tests/python/relax/script/test_relax_script_printer.py
index 08b12bc2e8..4a8ab35f95 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -146,15 +146,17 @@ class Module:
@Ts.prim_func
def tir_func(x: T.Buffer((T.int64(128), T.int64(128)), "float32"), y:
T.Buffer((T.int64(128), T.int64(128)), "float32")):
T.func_attr({"tirx.noalias": True})
- # with Ts.sblock("root"):
- for i in range(T.int64(0), T.int64(128)):
- for j in range(T.int64(0), T.int64(128)):
- with Ts.sblock(""):
- v = Ts.axis.spatial(T.int64(128), i, dtype="int64")
- v_1 = Ts.axis.spatial(T.int64(128), j, dtype="int64")
- Ts.reads(x[v, v_1])
- Ts.writes(y[v, v_1])
- y[v, v_1] = x[v, v_1] + T.float32(1.0)
+ with Ts.sblock("root"):
+ Ts.reads()
+ Ts.writes()
+ for i in range(T.int64(0), T.int64(128)):
+ for j in range(T.int64(0), T.int64(128)):
+ with Ts.sblock(""):
+ v = Ts.axis.spatial(T.int64(128), i, dtype="int64")
+ v_1 = Ts.axis.spatial(T.int64(128), j, dtype="int64")
+ Ts.reads(x[v, v_1])
+ Ts.writes(y[v, v_1])
+ y[v, v_1] = x[v, v_1] + T.float32(1.0)
@R.function
def foo(x: R.DTensor((128, 128), "float32", "mesh[0]", "S[0], R")) ->
R.DTensor((128, 128), "float32", "mesh[0]", "S[0], R"):
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 5d9c9b07a8..ae3b0f806f 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
@@ -322,23 +322,25 @@ def test_remap():
@Ts.prim_func
def main():
- # with Ts.sblock("root"):
- for i0 in range(128):
- for i1 in range(128):
- for i2 in range(128):
- for i3 in range(128):
- for i4 in range(128):
- for i5 in range(128):
- with Ts.sblock("update"):
- v = Ts.axis.spatial(128, i0 + 1)
- v_1 = Ts.axis.spatial(128, i1)
- v_2 = Ts.axis.reduce(128, i2)
- v_3 = Ts.axis.spatial(128, i3 - 1)
- v_4 = Ts.axis.reduce(128, i4)
- v_5 = Ts.axis.spatial(128, i5)
- Ts.reads()
- Ts.writes()
- T.evaluate(0)"""
+ with Ts.sblock("root"):
+ Ts.reads()
+ Ts.writes()
+ for i0 in range(128):
+ for i1 in range(128):
+ for i2 in range(128):
+ for i3 in range(128):
+ for i4 in range(128):
+ for i5 in range(128):
+ with Ts.sblock("update"):
+ v = Ts.axis.spatial(128, i0 + 1)
+ v_1 = Ts.axis.spatial(128, i1)
+ v_2 = Ts.axis.reduce(128, i2)
+ v_3 = Ts.axis.spatial(128, i3 - 1)
+ v_4 = Ts.axis.reduce(128, i4)
+ v_5 = Ts.axis.spatial(128, i5)
+ Ts.reads()
+ Ts.writes()
+ T.evaluate(0)"""
_assert_print(block_with_remap_explicitly.with_attr("global_symbol",
"main"), expected_output)
_assert_print(block_with_remap_implicitly.with_attr("global_symbol",
"main"), expected_output)
@@ -1278,8 +1280,8 @@ def test_same_name_var():
},
)
tvm.ir.assert_structural_equal(func, rt_func)
- assert out_str.count("for i in range(16):\n for j in range(16):")
== 1
- assert out_str.count("for i_1 in range(16):\n for j_1 in
range(16):") == 1
+ assert out_str.count("for i in range(16):\n for j in
range(16):") == 1
+ assert out_str.count("for i_1 in range(16):\n for j_1 in
range(16):") == 1
assert out_str.find("i_2") == -1
assert out_str.find("j_2") == -1
diff --git a/tests/python/tirx/script/test_tirx_script_basic_usage.py
b/tests/python/tirx/script/test_tirx_script_basic_usage.py
index d9910eb366..fccdf56306 100644
--- a/tests/python/tirx/script/test_tirx_script_basic_usage.py
+++ b/tests/python/tirx/script/test_tirx_script_basic_usage.py
@@ -732,7 +732,7 @@ def test_roundtrip_serial_unroll_false():
code = test.script()
assert "unroll=False" in code, f"printer should emit unroll=False,
got:\n{code}"
- assert "annotations" not in code, "printer should NOT emit annotations
dict"
+ assert "annotations=" not in code, "printer should NOT emit annotations
dict"
assert from_source(code).script() == code
assert_structural_equal(test, from_source(code))
@@ -754,7 +754,7 @@ def test_roundtrip_serial_unroll_true():
code = test.script()
assert "unroll=True" in code, f"printer should emit unroll=True,
got:\n{code}"
- assert "annotations" not in code, "printer should NOT emit annotations
dict"
+ assert "annotations=" not in code, "printer should NOT emit annotations
dict"
assert from_source(code).script() == code
assert_structural_equal(test, from_source(code))
@@ -776,7 +776,7 @@ def test_roundtrip_serial_unroll_count():
code = test.script()
assert "unroll=2" in code, f"printer should emit unroll=2, got:\n{code}"
- assert "annotations" not in code, "printer should NOT emit annotations
dict"
+ assert "annotations=" not in code, "printer should NOT emit annotations
dict"
assert from_source(code).script() == code
assert_structural_equal(test, from_source(code))
