This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 34b8f2211 [FEAT] Add reflected IR tree inspection (#119)
34b8f2211 is described below
commit 34b8f22118228ce9264bf42883ac907e8281ca8b
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 28 22:22:08 2026 -0400
[FEAT] Add reflected IR tree inspection (#119)
Add `tvm_tirx.ir.format_tree` to inspect IR and extension objects
through inherited FFI reflection. Class-aware fold, expand, and hide
rules partially override a small default profile; folded callbacks
customize leaf text while expansion exposes structure. Definition
metadata distinguishes variable declarations from references, and
structural comparison omits reflected defaults.
Declare typed loop defaults in native reflection, expose separate
default selector collections and folded formatters, and use the API in
Quick Start.
---
docs/get_started/quick_start.rst | 31 +--
docs/reference/python/ir.rst | 22 ++
include/tvm/ir/stmt.h | 8 +-
python/tvm_tirx/ir/__init__.py | 7 +
python/tvm_tirx/ir/format_tree.py | 399 +++++++++++++++++++++++++++++++++
tests/python/ir/test_ir_format_tree.py | 339 ++++++++++++++++++++++++++++
6 files changed, 790 insertions(+), 16 deletions(-)
diff --git a/docs/get_started/quick_start.rst b/docs/get_started/quick_start.rst
index cd8cb41ff..bc574a8f1 100644
--- a/docs/get_started/quick_start.rst
+++ b/docs/get_started/quick_start.rst
@@ -124,23 +124,28 @@ Inspect and Transform the Program
Besides compiling and running ``add_one``, we can also look inside it. The
``add_one`` returned by ``@T.function`` is a tree of program elements,
-commonly called an intermediate representation (IR). Roughly, the tree for
-the CPU version looks like this:
+commonly called an intermediate representation (IR). To see how the program
+is structured, we can use ``format_tree`` to display its IR as a tree:
+
+.. code-block:: python
+
+ import tvm_tirx
+
+ print(tvm_tirx.ir.format_tree(add_one))
.. code-block:: text
Function
- ├── params
- │ ├── Var(name="A", ty=TensorType(shape=[n], dtype=float32))
- │ └── Var(name="B", ty=TensorType(shape=[n], dtype=float32))
- └── body
- └── For(var=i, begin=0, end=n)
- └── body
- └── TensorStore(dest=B, indices=[i])
- └── value
- └── Add
- ├── a: TensorLoad(source=A, indices=[i])
- └── b: FloatImm(value=1.0)
+ ├── params[0]: A = Var(ty=TensorType(shape=[n], dtype=float32))
+ ├── params[1]: B = Var(ty=TensorType(shape=[n], dtype=float32))
+ └── body: SeqStmt
+ └── seq[0]: For(begin=IntImm(value=0), end=n)
+ ├── var: i = Var(ty=int32)
+ └── body: SeqStmt
+ └── seq[0]: TensorStore(dest=B, indices=[i])
+ └── value: Add
+ ├── a: TensorLoad(source=A, indices=[i])
+ └── b: FloatImm(value=1.0)
Every element is accessible from Python and other languages, so we can walk
this tree directly in Python:
diff --git a/docs/reference/python/ir.rst b/docs/reference/python/ir.rst
index 28215a22b..4926f3465 100644
--- a/docs/reference/python/ir.rst
+++ b/docs/reference/python/ir.rst
@@ -219,6 +219,7 @@ IR helpers
:toctree: generated/
const
+ format_tree
make_node
assert_structural_equal
load_json
@@ -237,3 +238,24 @@ Supporting functions
~base.register_script_printer
~base.register_dialect_script_method
+
+Tree Formatting Defaults
+------------------------
+
+.. data:: DEFAULT_FOLD
+
+ Tuple of field selectors to inline with their parent by default.
+
+.. data:: DEFAULT_EXPAND
+
+ Tuple of field selectors to expand explicitly by default. This is empty
+ because objects already expand unless a rule folds or hides them.
+
+.. data:: DEFAULT_HIDE
+
+ Tuple of field selectors to omit by default.
+
+.. data:: DEFAULT_FOLD_FORMATTERS
+
+ Standard folded representations, keyed by class. Callbacks take an object
+ and return single-line text.
diff --git a/include/tvm/ir/stmt.h b/include/tvm/ir/stmt.h
index 9d8a081c4..34777ce30 100644
--- a/include/tvm/ir/stmt.h
+++ b/include/tvm/ir/stmt.h
@@ -506,10 +506,12 @@ class ForObj : public StmtObj {
.def_ro("var", &ForObj::var, refl::AttachFieldFlag::SEqHashDefSimple())
.def_ro("begin", &ForObj::begin)
.def_ro("end", &ForObj::end)
- .def_ro("step", &ForObj::step)
- .def_ro("kind", &ForObj::kind)
+ // Reflection stores one typed default. Wider unit steps remain
distinct.
+ .def_ro("step", &ForObj::step, refl::DefaultValue(IntImm::Int32(1)))
+ .def_ro("kind", &ForObj::kind, refl::DefaultValue(ForKind::kDefault))
.def_ro("body", &ForObj::body)
- .def_ro("annotations", &ForObj::annotations);
+ .def_ro("annotations", &ForObj::annotations,
+ refl::DefaultValue(ffi::Map<ffi::String, ffi::Any>()));
}
/*! \brief Check it is a loop without nontrivial loop step. */
diff --git a/python/tvm_tirx/ir/__init__.py b/python/tvm_tirx/ir/__init__.py
index 6364296e1..f58ac43b5 100644
--- a/python/tvm_tirx/ir/__init__.py
+++ b/python/tvm_tirx/ir/__init__.py
@@ -90,6 +90,13 @@ from .stmt import (
# Load passes after the common IR classes used by transformations are
registered.
from . import transform
+from .format_tree import (
+ DEFAULT_FOLD,
+ DEFAULT_EXPAND,
+ DEFAULT_HIDE,
+ DEFAULT_FOLD_FORMATTERS,
+ format_tree,
+)
# Connect exposure after the shared IR wrappers are ready.
from ._expose_api import _install_exposure_callbacks
diff --git a/python/tvm_tirx/ir/format_tree.py
b/python/tvm_tirx/ir/format_tree.py
new file mode 100644
index 000000000..1b59a6990
--- /dev/null
+++ b/python/tvm_tirx/ir/format_tree.py
@@ -0,0 +1,399 @@
+# 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.
+"""Text trees for inspecting IR and extension objects."""
+
+from collections.abc import Callable, Iterable, Mapping
+from typing import Any
+
+import tvm_ffi
+
+from .expr import Expr, Var
+from .function import BaseFunc
+from .op import Op
+from .stmt import For
+from .type import PrimType
+
+FieldSelector = str | tuple[type, str]
+
+# Keep the public defaults in the same selector form as the keyword options.
+# Normalize them privately before applying exact-key caller overrides.
+DEFAULT_FOLD: tuple[FieldSelector, ...] = (
+ (Var, "ty"),
+ "source",
+ "dest",
+ "indices",
+ "op",
+ (For, "begin"),
+ (For, "end"),
+ (For, "step"),
+)
+"""Fields inlined with their parent by default."""
+
+DEFAULT_EXPAND: tuple[FieldSelector, ...] = ()
+"""Fields explicitly expanded by default; other objects expand
automatically."""
+
+DEFAULT_HIDE: tuple[FieldSelector, ...] = (
+ "loc",
+ (BaseFunc, "attrs"),
+ (BaseFunc, "ret_type"),
+ (Expr, "ty"),
+)
+"""Fields omitted by default."""
+
+DEFAULT_FOLD_FORMATTERS = {
+ PrimType: lambda node: str(node.dtype),
+ Op: lambda node: node.name,
+ Var: lambda node: node.name if node.name.isidentifier() else
repr(node.name),
+}
+"""Standard folded representations, keyed by class."""
+
+_SCALARS = (bool, int, float, complex, str, bytes)
+_SEQUENCES = (list, tuple, tvm_ffi.Array, tvm_ffi.List)
+_MAPPINGS = (Mapping, tvm_ffi.Map, tvm_ffi.Dict)
+_FFI_VALUES = (tvm_ffi.Object, tvm_ffi.Array, tvm_ffi.List, tvm_ffi.Map,
tvm_ffi.Dict)
+
+
+def _scalar(value):
+ if isinstance(value, tvm_ffi.dtype):
+ return str(value)
+ if isinstance(value, str):
+ return repr(str(value))
+ if value is None or type(value) in _SCALARS:
+ return repr(value)
+ return None
+
+
+def _identity(value):
+ if isinstance(value, _FFI_VALUES):
+ return ("ffi", value.__chandle__())
+ return ("python", id(value))
+
+
+def _fields(value):
+ # dataclasses.fields excludes legacy register_object classes. Keep this
+ # adapter local until FFI exposes enumeration for both class kinds.
+ info = type(value).__tvm_ffi_type_info__
+ chain = []
+ while info is not None:
+ chain.append(info)
+ info = info.parent_type_info
+ for info in reversed(chain):
+ yield from info.fields
+
+
+def _selector(selector):
+ if isinstance(selector, str) and selector:
+ return tvm_ffi.Object, selector
+ if (
+ isinstance(selector, tuple)
+ and len(selector) == 2
+ and isinstance(selector[0], type)
+ and issubclass(selector[0], tvm_ffi.Object)
+ and isinstance(selector[1], str)
+ and selector[1]
+ ):
+ return selector
+ raise TypeError("selectors must be nonempty field names or (FFI class,
field name) pairs")
+
+
+class _TreeFormatter:
+ def __init__(self, max_depth, show_defaults, fold, expand, hide,
fold_formatters):
+ self.max_depth = max_depth
+ self.show_defaults = show_defaults
+ self.rules = {
+ _selector(selector): action
+ for action, selectors in (
+ ("hide", DEFAULT_HIDE),
+ ("fold", DEFAULT_FOLD),
+ ("expand", DEFAULT_EXPAND),
+ )
+ for selector in selectors
+ }
+ self.explicit = set()
+ # Normalize first: a bare name and its Object pair are the same key.
+ # Applying collections in this order gives expand > fold > hide.
+ for action, selectors in (("hide", hide), ("fold", fold), ("expand",
expand)):
+ if selectors is None:
+ continue
+ if isinstance(selectors, (str, bytes)):
+ raise TypeError(f"{action} must be a collection of field
selectors")
+ for selector in selectors:
+ key = _selector(selector)
+ self.rules[key] = action
+ self.explicit.add(key)
+ self.fold_formatters = dict(DEFAULT_FOLD_FORMATTERS)
+ if fold_formatters is not None:
+ if not isinstance(fold_formatters, Mapping):
+ raise TypeError("fold_formatters must be a mapping of FFI
classes to callables")
+ for cls, callback in fold_formatters.items():
+ if not isinstance(cls, type) or not issubclass(cls,
tvm_ffi.Object):
+ raise TypeError("fold_formatters keys must be FFI classes")
+ if not callable(callback):
+ raise TypeError("fold_formatters values must be callable")
+ self.fold_formatters[cls] = callback
+
+ def _rule(self, value, field):
+ # Python MRO also includes unregistered IR family bases and downstream
+ # Python classes; native type-info ancestry alone would miss those.
+ for cls in type(value).__mro__:
+ key = (cls, field)
+ if key in self.rules:
+ return self.rules[key], key in self.explicit
+ return "auto", False
+
+ def _parts(self, value, region):
+ """Return a node name, layout kind, and ordered labelled children."""
+ if isinstance(value, _SEQUENCES):
+ return (
+ type(value).__name__,
+ "sequence",
+ [(i, child, region, "auto") for i, child in enumerate(value)],
+ )
+ if isinstance(value, _MAPPINGS):
+ items = list(value.items())
+ if all(_scalar(key) is not None for key, _ in items):
+ items.sort(key=lambda item: (type(item[0]).__name__,
_scalar(item[0])))
+ return (
+ type(value).__name__,
+ "mapping",
+ [(_scalar(key), child, region, "auto") for key, child in
items],
+ )
+ # Object keys have structure too; never invoke a dialect's repr.
+ return (
+ type(value).__name__,
+ "sequence",
+ [(i, pair, region, "auto") for i, pair in enumerate(items)],
+ )
+ if isinstance(value, tvm_ffi.Object):
+ name = type(value).__tvm_ffi_type_info__.type_key.rsplit(".",
1)[-1]
+ if isinstance(value, Var):
+ var_name = value.name if value.name.isidentifier() else
repr(value.name)
+ name = f"{var_name} = {name}"
+ entries = []
+ for field in _fields(value):
+ if isinstance(value, Var) and field.name == "name":
+ continue # The declaration name is already to the left of
'='.
+ action, explicit = self._rule(value, field.name)
+ if action == "hide":
+ continue
+ child = getattr(value, field.name)
+ if (
+ not self.show_defaults
+ and not explicit
+ and field.c_has_default
+ and tvm_ffi.structural_equal(child, field.c_default,
map_free_vars=False)
+ ):
+ continue
+ role = field.c_structural_eq
+ if role not in ("def-simple", "def-pattern"):
+ # Pattern roles propagate into types/containers, but stop
+ # at Var declarations. Their own type fields set new roles.
+ role = (
+ "def-pattern"
+ if region == "def-pattern"
+ and not isinstance(value, Var)
+ and role != "ignore"
+ else None
+ )
+ entries.append((field.name, child, role, action))
+ return name, "object", entries
+ return type(value).__name__, "opaque", []
+
+ def render(
+ self,
+ value,
+ depth=0,
+ ancestors=frozenset(),
+ region=None,
+ folded=False,
+ expand_reference=False,
+ ):
+ scalar = _scalar(value)
+ if scalar is not None:
+ return [scalar]
+ if (
+ isinstance(value, Var)
+ and depth
+ and region is None
+ and not expand_reference
+ and not folded
+ ):
+ return [value.name if value.name.isidentifier() else
repr(value.name)]
+ identity = _identity(value)
+ if identity in ancestors:
+ return ["<cycle>"]
+ name, kind, entries = self._parts(value, region)
+ if kind == "opaque":
+ return [f"<{name}>"]
+ if self.max_depth is not None and depth >= self.max_depth and entries:
+ return [f"{name}(...)"]
+ if folded:
+ for cls in type(value).__mro__:
+ if cls in self.fold_formatters:
+ text = self.fold_formatters[cls](value)
+ if not isinstance(text, str) or "\n" in text or "\r" in
text:
+ raise TypeError("fold_formatters callbacks must return
single-line text")
+ return [text]
+ ancestors = ancestors | {identity}
+ if folded or not entries:
+ opening, closing = {"sequence": ("[", "]"), "mapping": ("{",
"}")}.get(
+ kind, (f"{name}(", ")")
+ )
+ parts = []
+ for label, child, role, _ in entries:
+ # An enclosing fold is sticky: descendants may hide, but their
+ # expansion rules cannot reopen this inline subtree.
+ text = self.render(child, depth + 1, ancestors, role,
folded=True)[0]
+ if kind == "object":
+ text = f"{label}={text}"
+ elif kind == "mapping":
+ text = f"{label}: {text}"
+ parts.append(text)
+ return [opening + ", ".join(parts) + closing]
+ header, edges = [], []
+ for label, child, role, action in entries:
+ if kind == "object" and (
+ action == "fold" or (action == "auto" and _scalar(child) is
not None)
+ ):
+ text = self.render(child, depth + 1, ancestors, role,
folded=True)[0]
+ header.append(f"{label}={text}")
+ else:
+ label = f"[{label}]" if kind == "sequence" else str(label)
+ # Only explicit expansion opens a Var use into a declaration;
+ # the normal structural fallback keeps references name-only.
+ edges.extend(
+ self._edges(label, child, depth + 1, ancestors, role,
action == "expand")
+ )
+ lines = [name + ("(" + ", ".join(header) + ")" if header else "")]
+ for index, (label, child_lines) in enumerate(edges):
+ branch, padding = ("└── ", " ") if index == len(edges) - 1 else
("├── ", "│ ")
+ lines.append(f"{branch}{label}: {child_lines[0]}")
+ lines.extend(padding + line for line in child_lines[1:])
+ return lines
+
+ def _edges(self, label, value, depth, ancestors, region, expand_reference):
+ identity = _identity(value)
+ if (
+ isinstance(value, _SEQUENCES)
+ and len(value)
+ and (self.max_depth is None or depth < self.max_depth)
+ and identity not in ancestors
+ ):
+ for index, child in enumerate(value):
+ yield from self._edges(
+ f"{label}[{index}]",
+ child,
+ depth + 1,
+ ancestors | {identity},
+ region,
+ expand_reference,
+ )
+ else:
+ yield (
+ label,
+ self.render(value, depth, ancestors, region,
expand_reference=expand_reference),
+ )
+
+
+def format_tree(
+ value: Any,
+ *,
+ fold: Iterable[FieldSelector] | None = None,
+ expand: Iterable[FieldSelector] | None = None,
+ hide: Iterable[FieldSelector] | None = None,
+ fold_formatters: Mapping[type, Callable[[Any], str]] | None = None,
+ max_depth: int | None = None,
+ show_defaults: bool = False,
+) -> str:
+ """Display an object's fields as a text tree.
+
+ Objects expand into labelled edges, arrays into indexed edges such as
+ ``seq[0]``, and scalar fields stay in headers. SeqStmt wrappers remain.
+ Expanded variable definitions display ``name = Var(...)`` at each
+ occurrence, even when several definitions refer to the same variable.
+ Variable uses and folded variables display names. Explicit expansion
+ reveals the details of a variable use.
+
+ Field selectors are bare names that match every object, or ``(class,
name)``
+ pairs that match a class and its subclasses. More-specific classes take
+ precedence. Caller selections replace the default action for the exact
+ same selector; if it appears in multiple caller options, expand wins over
+ fold, which wins over hide. Explicit fold or expand selections also expose
+ default-valued fields. Container keys are not field selectors. None and
+ empty collections retain the standard defaults. Formatting never changes
the supplied
+ collections or the shared defaults.
+
+ Parameters
+ ----------
+ value : Any
+ IR object, extension object, container, or scalar to inspect.
+ fold : iterable of str or (class, str), optional
+ Fields to inline with their parent. Standard fold selections are
+ listed in :data:`~tvm_tirx.ir.DEFAULT_FOLD`. By default, variable
types,
+ source/destination, indices, operation references, and loop bounds
+ are folded. Folding
+ includes the visible subtree; descendants stay folded or hidden even
+ if an inner field requests expansion.
+ expand : iterable of str or (class, str), optional
+ Fields to show on separate tree edges. Standard expansion selections
+ are listed in :data:`~tvm_tirx.ir.DEFAULT_EXPAND`. Objects expand
automatically
+ unless another rule folds or hides them. Expanding a variable use
shows its declaration.
+ hide : iterable of str or (class, str), optional
+ Fields to omit. Standard hide selections are listed in
+ :data:`~tvm_tirx.ir.DEFAULT_HIDE`: source locations, function
attributes/return
+ types, and expression types except variable types.
+ fold_formatters : mapping of class to callable, optional
+ Partial overrides of :data:`~tvm_tirx.ir.DEFAULT_FOLD_FORMATTERS`. The
most specific
+ matching callback takes an object and returns single-line text. These
+ callbacks apply only inside folds; expansion shows the object's fields.
+ max_depth : int, optional
+ Maximum number of object/container edges from the root (depth zero).
+ Deeper children display ``...``. Indexed arrays still count toward
+ depth. None means no depth limit. Cycles display ``<cycle>``.
+ show_defaults : bool
+ Include fields structurally equal to their declared defaults. Different
+ types and distinct free variables remain distinguishable. By default,
+ matching fields are omitted; fields without declared defaults remain
+ visible. Hidden fields stay hidden.
+ For loops omit default sequential kind, empty annotations, and an
+ ``int32`` unit step; differently typed or nonunit steps remain visible.
+
+ Returns
+ -------
+ text : str
+ A tree without a trailing newline. Folded text is not width-truncated;
+ expand selected fields to split long headers. Shared children can
+ appear repeatedly. Maps with scalar keys sort by key type and text;
+ other maps show key/value pairs in iteration order. This inspection
+ text is not a serialization format.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ from tvm_tirx import ir
+
+ print(ir.format_tree(ir.Var("i", "int32")))
+ # i = Var(ty=int32)
+ print(ir.format_tree(ir.Var("i", "int32"), expand=[(ir.Var, "ty")]))
+ # i = Var
+ # └── ty: PrimType(dtype=int32)
+ """
+ if max_depth is not None and (type(max_depth) is not int or max_depth < 0):
+ raise ValueError("max_depth must be a nonnegative integer or None")
+ formatter = _TreeFormatter(max_depth, show_defaults, fold, expand, hide,
fold_formatters)
+ return "\n".join(formatter.render(value))
diff --git a/tests/python/ir/test_ir_format_tree.py
b/tests/python/ir/test_ir_format_tree.py
new file mode 100644
index 000000000..36b6ee2d1
--- /dev/null
+++ b/tests/python/ir/test_ir_format_tree.py
@@ -0,0 +1,339 @@
+# 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.
+"""Test reflected IR tree inspection."""
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+import tvm_ffi
+from tvm_ffi.dataclasses import field, py_class
+
+import tvm_tirx.testing
+from tvm_tirx import ir, tirx
+from tvm_tirx.script import from_source
+
+
+@py_class("testing.tree.Base")
+class TreeBase(tvm_ffi.Object):
+ label: str
+
+
+@py_class("testing.tree.Extension")
+class TreeExtension(TreeBase):
+ child: Any
+ count: Any = 0
+
+
+@py_class("testing.tree.Binding")
+class TreeBinding(tvm_ffi.Object):
+ declaration: Any = field(structural_eq="def-simple")
+ value: Any
+
+
[email protected]("device", ["cpu", "cuda"])
+def test_quick_start(device):
+ # Parse the actual definitions without compilation, device allocation, or
export.
+ root = Path(__file__).resolve().parents[3]
+ source = (root / f"examples/quick_start/add_one_{device}.py").read_text()
+ source = source.split("# [run.begin]")[0]
+ source = source[source.index("from tvm_tirx.script") :]
+ function = from_source(source)
+ text = ir.format_tree(function)
+ assert "params[0]: A = Var(ty=TensorType" in text
+ assert "params[1]: B = Var(ty=TensorType" in text
+ assert text.count("shape=[n]") == 2 # Folded definitions use names too.
+ expanded = ir.format_tree(function, expand=[(ir.Var, "ty")])
+ assert expanded.count("shape[0]: n = Var") == 2 # Repeated expanded
pattern occurrences.
+ assert "SeqStmt" in text
+ assert "TensorStore(dest=B, indices=[i])" in text
+ assert "├── a: TensorLoad(source=A, indices=[i])" in text
+ assert "└── b: FloatImm(value=1.0)" in text
+ assert "loc:" not in text
+ assert "global_symbol" not in text
+ if device == "cpu":
+ assert "For(begin=IntImm(value=0), end=n)" in text
+ assert "var: i = Var(ty=int32)" in text
+ assert "var=i" not in text
+ assert "step=" not in text and "kind=" not in text and "annotations"
not in text
+ else:
+ assert "RegionStmt(op=tirx.device_entry)" in text and "If" in text and
"cuda" in text
+ assert "index=" not in text and "args_info" not in text
+ assert "op: Op(index=" in ir.format_tree(function, expand=["op"])
+ ir.assert_structural_equal(function, from_source(source))
+ assert "def add_one" in function.script()
+
+
+def test_var_definition_and_type_expansion():
+ var = ir.Var("i", "int32")
+ assert ir.format_tree(var) == "i = Var(ty=int32)"
+ assert ir.format_tree(var, expand=[(ir.Var, "ty")]) == (
+ "i = Var\n└── ty: PrimType(dtype=int32)"
+ )
+ assert ir.format_tree(var, hide=[(ir.Var, "ty")]) == "i = Var()"
+ assert ir.format_tree(ir.Var("a long name", "int32")) == "'a long name' =
Var(ty=int32)"
+
+
+def test_extension_fields_and_defaults():
+ extension = TreeExtension(label="extension", child=ir.Var("x", "int32"))
+ text = ir.format_tree(extension)
+ assert "Extension(label='extension')" in text # Includes inherited fields.
+ assert "child: x" in text
+ assert "count=" not in text
+ assert "count=0" in ir.format_tree(extension, show_defaults=True)
+ assert "count: 0" in ir.format_tree(extension, expand=["count"])
+ extension.count = 2
+ assert "count=2" in ir.format_tree(extension)
+ extension.count = False
+ assert "count=False" in ir.format_tree(extension) # 0 is not structurally
False.
+ assert "value=0" in ir.format_tree(ir.IntImm("int32", 0)) # No declared
default.
+
+
+def test_rule_overlay_specificity_and_nonmutation():
+ var = ir.Var("i", "int32")
+ defaults = (ir.DEFAULT_FOLD, ir.DEFAULT_EXPAND, ir.DEFAULT_HIDE)
+ assert all(isinstance(selectors, tuple) for selectors in defaults)
+ formatters = dict(ir.DEFAULT_FOLD_FORMATTERS)
+ baseline = ir.format_tree(var)
+ assert ir.format_tree(var, fold=None, expand=None, hide=None,
fold_formatters=None) == baseline
+ assert ir.format_tree(var, fold=[], expand=[], hide=[],
fold_formatters={}) == baseline
+ # A broad override cannot erase the more-specific default Var.ty rule.
+ assert ir.format_tree(var, expand=[(ir.Expr, "ty")]) == baseline
+ assert ir.format_tree(var, hide=["ty"]) == baseline
+ assert "ty: PrimType" in ir.format_tree(var, expand=[(ir.Var, "ty")])
+ extension = TreeExtension(label="extension", child=ir.PrimType("int32"))
+ fold, expand, hide = ["child"], [(TreeBase, "child")], [(TreeExtension,
"child")]
+ assert "child" not in ir.format_tree(extension, fold=fold, expand=expand,
hide=hide)
+ assert (fold, expand, hide) == (["child"], [(TreeBase, "child")],
[(TreeExtension, "child")])
+ # Bare and Object forms normalize to one key before caller precedence.
+ text = ir.format_tree(
+ extension, hide=["child"], fold=[(tvm_ffi.Object, "child")],
expand=["child"]
+ )
+ assert "child: PrimType(dtype=int32)" in text
+ assert "child=int32" in ir.format_tree(extension, hide=["child"],
fold=["child"])
+ assert (ir.DEFAULT_FOLD, ir.DEFAULT_EXPAND, ir.DEFAULT_HIDE) == defaults
+ assert ir.DEFAULT_FOLD_FORMATTERS == formatters
+
+
+def test_bare_selector_overrides_default_fold():
+ load = ir.TensorLoad(ir.Var("A", tirx.TensorType([4], "float32")),
[ir.IntImm("int32", 0)])
+ assert "source=A" in ir.format_tree(load)
+ assert "source: A = Var(ty=TensorType" in ir.format_tree(load,
expand=["source"])
+ assert "source: A = Var(ty=TensorType" in ir.format_tree(
+ load, expand=[(ir.TensorLoad, "source")]
+ )
+ assert "source: A = Var()" in ir.format_tree(load, expand=["source"],
hide=[(ir.Var, "ty")])
+ expanded = ir.format_tree(load, expand=["source", (ir.Var, "ty")])
+ assert "ty: TensorType" in expanded and "shape[0]" in expanded
+ assert "dtype: PrimType(dtype=float32)" in expanded
+ extension = TreeExtension(label="load", child=load)
+ assert "child=TensorLoad(source=A," in ir.format_tree(
+ extension, fold=["child"], expand=["source"]
+ )
+ assert "source" not in ir.format_tree(load, hide=["source"])
+
+
+def test_recursive_fold_is_sticky():
+ inner = TreeExtension(label="inner", child=ir.PrimType("int32"))
+ outer = TreeExtension(label="outer", child=inner)
+ text = ir.format_tree(outer, fold=[(TreeExtension, "child")],
expand=[(TreeBase, "label")])
+ assert "child=Extension(label='inner', child=int32)" in text
+ assert "label: 'outer'" in text
+ assert "inner" in text and "..." not in text
+ assert "child=Extension(child=int32)" in ir.format_tree(outer,
fold=["child"], hide=["label"])
+ seq = ir.SeqStmt([ir.Evaluate(ir.IntImm("int32", 1))])
+ assert ir.format_tree(seq, fold=["seq"], expand=[(ir.Evaluate, "value")])
== (
+ "SeqStmt(seq=[Evaluate(value=IntImm(value=1))])"
+ )
+
+
+def test_fold_formatters_inheritance_and_expansion():
+ inner = TreeExtension(label="inner", child=None)
+ outer = TreeExtension(label="outer", child=inner)
+ callbacks = {TreeBase: lambda obj: obj.label.upper()}
+ assert "child=INNER" in ir.format_tree(outer, fold=["child"],
fold_formatters=callbacks)
+ assert "child: Extension(label='inner')\n └── child: None" in
ir.format_tree(
+ outer, expand=["child"], fold_formatters=callbacks
+ )
+ callbacks[TreeExtension] = lambda obj: "specific"
+ assert "child=specific" in ir.format_tree(outer, fold=["child"],
fold_formatters=callbacks)
+ assert len(callbacks) == 2
+ # Most-specific standard PrimType callback wins over a broad caller one.
+ var = ir.Var("i", "int32")
+ assert ir.format_tree(var, fold_formatters={ir.Type: lambda obj: "type"})
== "i = Var(ty=int32)"
+ assert (
+ ir.format_tree(var, fold_formatters={ir.PrimType: lambda obj:
"custom"})
+ == "i = Var(ty=custom)"
+ )
+ assert ir.format_tree(
+ ir.PrimType("int32"), fold_formatters={ir.PrimType: lambda obj:
"custom"}
+ ) == ("PrimType(dtype=int32)")
+
+
+def test_reflected_binding_sites_and_indexed_edges():
+ var = ir.Var("item", "int32")
+ # Binding roles come from reflection, including on downstream arbitrary
names.
+ binding = TreeBinding(declaration=[var, var], value=var)
+ text = ir.format_tree(binding)
+ assert "declaration[0]: item = Var(ty=int32)" in text
+ assert "declaration[1]: item = Var(ty=int32)" in text
+ assert "value: item" in text
+ assert "declaration=[item, item]" in ir.format_tree(binding,
fold=["declaration"])
+ seq = ir.SeqStmt([ir.Evaluate(var), ir.Evaluate(var)])
+ text = ir.format_tree(seq)
+ assert "seq[0]: Evaluate" in text and "seq[1]: Evaluate" in text
+ assert text.count("value: item") == 2
+
+
+def make_loop(dtype="int32", step=1, kind=ir.ForKind.DEFAULT,
annotations=None):
+ return ir.For(
+ ir.Var("i", dtype),
+ ir.IntImm(dtype, 0),
+ ir.IntImm(dtype, 4),
+ ir.IntImm(dtype, step),
+ kind,
+ [],
+ annotations=annotations,
+ )
+
+
+def test_native_loop_defaults_and_structural_comparison():
+ fields = {f.name: f for f in ir.For.__tvm_ffi_type_info__.fields}
+ loop = make_loop()
+ for name in ("kind", "step", "annotations"):
+ assert fields[name].c_has_default
+ assert tvm_ffi.structural_equal(
+ getattr(loop, name), fields[name].c_default, map_free_vars=False
+ )
+ assert not loop.step.same_as(fields["step"].c_default) # Equal,
separately allocated.
+ assert not loop.annotations.same_as(fields["annotations"].c_default)
+ text = ir.format_tree(loop)
+ assert "var: i = Var(ty=int32)" in text
+ assert "For(var=i," in ir.format_tree(loop, fold=[(ir.For, "var")])
+ assert "step=" not in text and "kind=" not in text and "annotations" not
in text
+ assert "step=IntImm(value=1)" in ir.format_tree(loop, show_defaults=True)
+ assert "kind=0" in ir.format_tree(loop, show_defaults=True)
+ assert "annotations: {}" in ir.format_tree(loop, expand=[(ir.For,
"annotations")])
+ assert "step=IntImm(value=1)" in ir.format_tree(loop, fold=[(ir.For,
"step")])
+ assert "step=IntImm(value=2)" in ir.format_tree(make_loop(step=2))
+ assert "step=IntImm(value=1)" in ir.format_tree(make_loop(dtype="int64"))
+ assert "kind=1" in ir.format_tree(make_loop(kind=ir.ForKind.PARALLEL))
+ assert "'tag': 1" in ir.format_tree(make_loop(annotations={"tag": 1}))
+ # Metadata defaults also participate in generic reflected construction.
+ constructed = ir.make_node(
+ "ir.For", var=loop.var, begin=loop.begin, end=loop.end,
body=loop.body, loc=loop.loc
+ )
+ ir.assert_structural_equal(loop, constructed)
+ ir.assert_structural_equal(loop, ir.load_json(ir.save_json(loop)))
+ ir.assert_structural_equal(
+ make_loop(dtype="int64"),
ir.load_json(ir.save_json(make_loop(dtype="int64")))
+ )
+
+
+@py_class("testing.tree.DefaultVar")
+class TreeDefaultVar(tvm_ffi.Object):
+ value: Any = ir.Var("default", "int32")
+
+
+def test_default_free_variables_are_not_remapped():
+ assert "value" not in ir.format_tree(TreeDefaultVar())
+ other = ir.Var("default", "int32")
+ assert "value: default" in ir.format_tree(TreeDefaultVar(value=other))
+
+
+def test_explicit_expansion_through_indexed_edges():
+ var = ir.Var("item", "int32")
+ extension = TreeExtension(label="uses", child=[var, var])
+ assert "child[0]: item" in ir.format_tree(extension)
+ expanded = ir.format_tree(extension, expand=["child"])
+ assert expanded.count("item = Var(ty=int32)") == 2
+ assert "child[1]: item = Var" in expanded
+
+
+def test_locations_and_containers():
+ var = ir.Var("x", "int32")
+ assert "loc: UnknownLoc()" in ir.format_tree(var, expand=["loc"])
+ assert "loc" not in ir.format_tree(var, show_defaults=True)
+ assert "loc" not in ir.format_tree(var, hide=[])
+ assert ir.format_tree({"z": 1, "a": 2}) == "dict\n├── 'a': 2\n└── 'z': 1"
+ assert ir.format_tree(tvm_ffi.convert({"z": 1, "a": 2})) == "Map\n├── 'a':
2\n└── 'z': 1"
+ assert "'loc': x" in ir.format_tree({"loc": var}, hide=["loc"])
+ assert ir.format_tree([]) == "[]" and ir.format_tree({}) == "{}"
+ assert ir.format_tree(object()) == "<object>"
+ assert "'x\\ny'" in ir.format_tree([0, False, None, "x\ny"])
+ text = ir.format_tree(tvm_ffi.convert({var: 7}))
+ assert "[0][0]: x" in text and "[0][1]: 7" in text
+ text = ir.format_tree([var, var])
+ assert text.count("x") == 2 and "<cycle>" not in text
+
+
[email protected]("folded", [False, True])
+def test_depth_and_cycles(folded):
+ options = {"fold": ["child"]} if folded else {}
+ nested = TreeExtension(label="outer", child=[[[1]]])
+ assert ir.format_tree([[[1]]], max_depth=0) == "list(...)"
+ text = ir.format_tree(nested, max_depth=1, **options)
+ assert "Array(...)" in text and "1" not in text
+ assert ir.format_tree(1, max_depth=0) == "1"
+ assert ir.format_tree(ir.Var("i", "int32"), max_depth=0) == "i = Var(...)"
+ cycle = []
+ cycle.append(cycle)
+ assert "<cycle>" in ir.format_tree(cycle)
+ extension = TreeExtension(label="cycle", child=None)
+ extension.child = extension
+ try:
+ assert "<cycle>" in ir.format_tree(extension, **options)
+ finally:
+ extension.child = None
+
+
[email protected]("depth", [-1, 1.5, True])
+def test_invalid_depth(depth):
+ with pytest.raises(ValueError, match="max_depth"):
+ ir.format_tree(None, max_depth=depth)
+
+
[email protected](
+ "option,selectors",
+ [
+ ("fold", "ty"),
+ ("expand", [""]),
+ ("hide", [(ir.Var,)]),
+ ("fold", [(int, "ty")]),
+ ("expand", [(ir.Var, 1)]),
+ ("hide", [1]),
+ ],
+)
+def test_invalid_selectors(option, selectors):
+ with pytest.raises(TypeError):
+ ir.format_tree(None, **{option: selectors})
+
+
[email protected]("callbacks", [[], {int: str}, {ir.Type: None}])
+def test_invalid_callbacks(callbacks):
+ with pytest.raises(TypeError, match="fold_formatters"):
+ ir.format_tree(None, fold_formatters=callbacks)
+
+
[email protected]("result", [1, "line\nline"])
+def test_invalid_callback_output(result):
+ with pytest.raises(TypeError, match="single-line text"):
+ ir.format_tree(ir.Var("i", "int32"), fold_formatters={ir.PrimType:
lambda obj: result})
+
+
+if __name__ == "__main__":
+ tvm_tirx.testing.main()
