This is an automated email from the ASF dual-hosted git repository.
yongwww 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 7a303d96c9 [Instrument] Add default instrument to print all passes
(#16497)
7a303d96c9 is described below
commit 7a303d96c9cb9752897ce9c1562b81d9ebe7fbc0
Author: Anirudh Sundar Subramaniam <[email protected]>
AuthorDate: Sat Feb 3 00:12:51 2024 +0530
[Instrument] Add default instrument to print all passes (#16497)
A small patch to add default instruments that prints all passes before
or after each pass. This is used very frequently when debugging and we
don't know which pass could cause a particular change
---
python/tvm/ir/instrument.py | 18 ++++++++++
tests/python/ir/test_pass_instrument.py | 63 +++++++++++++++++++++++++++++++++
2 files changed, 81 insertions(+)
diff --git a/python/tvm/ir/instrument.py b/python/tvm/ir/instrument.py
index 59b86d1bb5..4402d0fc22 100644
--- a/python/tvm/ir/instrument.py
+++ b/python/tvm/ir/instrument.py
@@ -273,3 +273,21 @@ class PassPrintingInstrument:
def run_after_pass(self, mod, pass_info):
if pass_info.name in self.print_after_pass_names:
print(f"Print IR after: {pass_info.name}\n{mod}\n\n")
+
+
+@pass_instrument
+class PrintAfterAll:
+ """Print the name of the pass, the IR, only after passes execute."""
+
+ def run_after_pass(self, mod, info):
+ print(f"After Running Pass: {info}")
+ print(mod)
+
+
+@pass_instrument
+class PrintBeforeAll:
+ """Print the name of the pass, the IR, only before passes execute."""
+
+ def run_before_pass(self, mod, info):
+ print(f"Before Running Pass: {info}")
+ print(mod)
diff --git a/tests/python/ir/test_pass_instrument.py
b/tests/python/ir/test_pass_instrument.py
new file mode 100644
index 0000000000..cfeb70b963
--- /dev/null
+++ b/tests/python/ir/test_pass_instrument.py
@@ -0,0 +1,63 @@
+# 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.
+""" Instrument test cases.
+"""
+
+import tvm
+from tvm import relax
+from tvm.ir.instrument import PrintAfterAll, PrintBeforeAll
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tir as T
+
+# pylint:
disable=invalid-name,missing-function-docstring,no-value-for-parameter
+
+
+def test_tir_print_all_passes(capsys):
+ @T.prim_func
+ def func(a: T.handle, b: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128, 128, 128))
+ B = T.match_buffer(b, (128, 128, 128, 128))
+ for i, j, k, l in T.grid(128, 128, 128, 128):
+ with T.block("B"):
+ vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
+ B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
+
+ with tvm.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(),
PrintAfterAll()]):
+ tvm.lower(func)
+ all_passes_output = capsys.readouterr().out
+ assert "Before Running Pass:" in all_passes_output
+ assert "After Running Pass:" in all_passes_output
+ assert "pass name: tir." in all_passes_output
+
+
+def test_relax_print_all_passes(capsys):
+ @I.ir_module
+ class Module:
+ @R.function
+ def func(x: R.Tensor((16,), "float32"), y: R.Tensor((16,), "float32")):
+ z = R.add(x, y)
+ y = z
+ return y
+
+ pipeline = relax.get_pipeline("default_build")
+ with tvm.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(),
PrintAfterAll()]):
+ pipeline(Module)
+ all_passes_output = capsys.readouterr().out
+ assert "Before Running Pass:" in all_passes_output
+ assert "After Running Pass:" in all_passes_output
+ assert "pass name: _pipeline" in all_passes_output
