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 82bf83eb8 [DOCS] Restructure IR Guides for scanning: tables, notes,
shorter paragraphs (#151)
82bf83eb8 is described below
commit 82bf83eb84eb6152ff6b7151402d27427045cd35
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Nov 1 12:50:01 2026 -0500
[DOCS] Restructure IR Guides for scanning: tables, notes, shorter
paragraphs (#151)
This PR restructures the IR Guides for scanning without removing
narrative content: long paragraphs are split to one idea each, parallel
items become tables, subtleties move into notes, and key points are
bolded sparingly.
- **Structural Traversal, Rewriting, and Comparison**: condensed
Definitions and Uses opening; definition region kinds as a table (Kind /
Meaning / In `add_one`); `DEF_PATTERN` details in a note; a bold-led
paragraph on how the utilities keep definitions and uses consistent. A
table of walk callback return values; "a map callback must return a
node" and free variables (`map_free_vars`) as notes; a shorter Equality
opening; Visit and Mutate splits the why from the how, with the "return
`default_visit`'s result" rule as a note.
- **Printing and Parsing**: the `[n]` / Python 3.12 type-parameter
detail, `save_json`, and `line_trace_scope` become notes; the section
that duplicated the page title is renamed; the predictability principle
moves to the start of How the Parser Works as a list; a
TVMScript-to-builder table replaces the mapping prose; region syntax
gets its own subsection.
- **Symbolic Integer Analysis**: the "False means not proved" and "last
synchronized context" points become notes; one table lists the four
`*_with_constraint_context` utilities.
- **IR Overview**: split the intro paragraph; collapse the `prim.*`
lines of the class hierarchy into categories; a namespace table (`ir` /
`ir.prim` / `tirx`); source locations as a note; the chapter map
mentions comparing IR.
- **Functions and Modules**: output blocks for `add_one.ty` and the
printed `T.func_attr`; a shorter Attributes paragraph; `GlobalVar` name
vs. `global_symbol` as a table.
- **Types and Expressions, Statements**: remove forward cross-references
to later chapters; keep a one-sentence version of "rewrites keep
dependent types consistent".
- **Consistency**: "chapter" instead of "guide" in body pages; remove
unreferenced labels; `@T.ir_module` in the `ir.transform` example;
Structural Hooks and pass selection rules as tables in the reference.
---
docs/ir_guides/functions_and_modules.rst | 71 ++++++++--
docs/ir_guides/operators_calls_and_regions.rst | 2 +-
docs/ir_guides/overview.rst | 77 ++++++-----
docs/ir_guides/printing_and_parsing.rst | 124 +++++++++++------
docs/ir_guides/statements.rst | 10 +-
docs/ir_guides/structural_utilities.rst | 178 ++++++++++++++-----------
docs/ir_guides/symbolic_integer_analysis.rst | 59 +++++---
docs/ir_guides/types_and_expressions.rst | 30 ++---
docs/reference/python/ffi.structural.rst | 38 ++++--
docs/reference/python/ir.transform.rst | 7 +-
examples/ir_guides/functions_and_modules.py | 4 +-
11 files changed, 373 insertions(+), 227 deletions(-)
diff --git a/docs/ir_guides/functions_and_modules.rst
b/docs/ir_guides/functions_and_modules.rst
index ff2d9bd9d..f7d10b2d7 100644
--- a/docs/ir_guides/functions_and_modules.rst
+++ b/docs/ir_guides/functions_and_modules.rst
@@ -66,21 +66,34 @@ A ``tirx.Function`` has three fields of its own:
- The statements that run when the function is called.
It also has the fields every ``ir.BaseFunc`` shares: ``attrs``, described
-below; ``ty``, the function's signature, built from its parameter and
-return types; and ``loc``, its source location.
+below; ``ty``, the function's signature; and ``loc``, its source location.
+``ty`` is built from the parameter and return types:
+
+.. literalinclude:: ../../examples/ir_guides/functions_and_modules.py
+ :language: python
+ :start-after: # [ty.begin]
+ :end-before: # [ty.end]
+
+The output is:
+
+.. code-block:: text
+
+ T.Func([T.Tensor((n,), T.float32), T.Tensor((n,), T.float32)], T.Tuple())
+
+The two tensor parameters come first, and ``T.Tuple()`` is the empty
+return type.
Attributes
~~~~~~~~~~
-``attrs`` allows us to store additional attribute information about a
-function. For example, ``global_symbol`` locks in the symbol name of a
-public function: ``add_one`` is exported as ``"add_one"`` in the compiled
-library. ``@T.function`` sets it from the Python name. A function
-marked ``@T.function(private=True)`` has no ``global_symbol`` and stays
-internal to the module.
+``attrs`` stores additional information about a function. The most
+common attribute is ``global_symbol``, the name a public function is
+exported under in the compiled library. ``@T.function`` sets it from the
+Python name, so ``add_one`` is exported as ``"add_one"``; a function marked
+``@T.function(private=True)`` has no ``global_symbol`` and stays internal
+to the module.
-In TVMScript, ``T.func_attr`` adds attributes in the function body, and
-the printer writes them back the same way:
+In TVMScript, ``T.func_attr`` adds attributes in the function body:
.. literalinclude:: ../../examples/ir_guides/functions_and_modules.py
:language: python
@@ -88,6 +101,20 @@ the printer writes them back the same way:
:end-before: # [func_attr.end]
:dedent: 4
+The printer writes them back the same way. ``add_one.script()`` gives:
+
+.. code-block:: python
+
+ from __future__ import annotations
+
+ # from tvm_tirx.script import tirx as T
+
+ @T.function
+ def add_one[n](A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
+ T.func_attr({"note": "example"})
+ for i in range(n):
+ B[i] = A[i] + T.float32(1.0)
+
From Python, ``with_attr`` returns a new function with an attribute added
and leaves the original unchanged:
@@ -142,7 +169,23 @@ replace the definition without touching the places that
refer to it.
In practice, references between functions mostly connect host and device
functions, such as a host function and the kernels it launches.
-The ``GlobalVar``'s name is the function's name within the module. It is
-separate from ``global_symbol``, the name a public function is exported
-under. For ``add_one``, both are ``"add_one"``; a private function has a
-module name but no exported symbol.
+A function has two names, which coincide for ``add_one``:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 22 42 16 20
+
+ * - Name
+ - Meaning
+ - ``add_one``
+ - Private function
+ * - ``GlobalVar`` name
+ - The function's name within the module.
+ - ``"add_one"``
+ - Present
+ * - ``global_symbol``
+ - The name a public function is exported under.
+ - ``"add_one"``
+ - Absent
+
+A private function has a module name but no exported symbol.
diff --git a/docs/ir_guides/operators_calls_and_regions.rst
b/docs/ir_guides/operators_calls_and_regions.rst
index f722b11f0..1a6a9078f 100644
--- a/docs/ir_guides/operators_calls_and_regions.rst
+++ b/docs/ir_guides/operators_calls_and_regions.rst
@@ -19,7 +19,7 @@
Operators, Calls, and Regions
=============================
-This guide introduces **operators**, the main extension point of TIRx.
+This chapter introduces **operators**, the main extension point of TIRx.
An operator, ``ir.Op``, names a computation such as sigmoid or a
thread-index read, and carries attributes that tell the compiler how to
treat it. A call applies an operator as an expression, and a region
diff --git a/docs/ir_guides/overview.rst b/docs/ir_guides/overview.rst
index 89eac22e4..a8935fcbc 100644
--- a/docs/ir_guides/overview.rst
+++ b/docs/ir_guides/overview.rst
@@ -23,16 +23,19 @@ to every stage of compilation, not only its input and
output. By providing
a consistent way to inspect and modify programs at each stage, TIRx makes
it easier to understand and customize the compiler.
-TIRx represents a program as a tree of intermediate representation (IR)
-nodes, such as functions, statements, and values. The tree is similar to the
-abstract syntax tree (AST) of Python code, but in a more structured, typed
-form. A compiler pipeline analyzes and transforms this tree step by step
-to produce target code. IR objects are immutable: an analysis reads the
-tree, and a transformation builds a new tree instead of modifying the
-existing one, sharing the parts that did not change. The
-``tvm_tirx.script`` interfaces print the IR as TVMScript, a Python-based
-notation for tensor programs. For example, this TVMScript definition
-describes an add-one computation on a GPU:
+TIRx represents a program as a **tree of intermediate representation (IR)
+nodes**, such as functions, statements, and values. The tree is similar to
+the abstract syntax tree (AST) of Python code, but in a more structured,
+typed form. A compiler pipeline analyzes and transforms this tree step by
+step to produce target code.
+
+IR objects are **immutable**: an analysis reads the tree, and a
+transformation builds a new tree instead of modifying the existing one,
+sharing the parts that did not change.
+
+The ``tvm_tirx.script`` interfaces print the IR as TVMScript, a Python-based
notation for tensor programs. For
+example, this TVMScript definition describes an add-one computation on a
+GPU:
.. literalinclude:: ../../examples/quick_start/add_one_cuda.py
:language: python
@@ -223,24 +226,38 @@ The following diagram shows the class hierarchy of the IR
nodes:
│ └── ir.CallSiteLoc Callee and caller locations
└── ir.IRModule Global function definitions
-These IR nodes are designed to be generic, so that new IR variants can be built
-on top of them, and their namespaces show their roles. The ``ir`` namespace
-holds the common nodes shared by all variants. Within it, ``ir.prim`` (shown as
-``prim``) holds primitive expressions, which compute scalar or vector values of
-a primitive type (:py:class:`ir.PrimType <tvm_tirx.ir.PrimType>`), such as
-``float32`` or ``int32``. These expressions concisely represent dynamic indices
-and element computations, such as ``A[i] + 1`` in ``add_one``. The ``tirx``
-namespace builds on ``ir`` with the types and operations specific to tensor
-programs, such as ``tirx.Function`` and ``tirx.TensorType``. More generally,
-TIRx is designed for extension: downstream projects can build on the IR in the
-same way, registering new operators, defining new types, and adding their own
-TVMScript namespaces from Python, C++, or Rust, while reusing TIRx as a
-foundation.
-
-IR nodes can also carry source locations
-(:py:class:`ir.Location <tvm_tirx.ir.Location>`). For example, an expression's
-``loc`` field records where it came from in the source, so a diagnostic
-can point back to the original code.
+These IR nodes are designed to be generic, so that new IR variants can be
+built on top of them, and their namespaces show their roles:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 15 85
+
+ * - Namespace
+ - Holds
+ * - ``ir``
+ - The common nodes shared by all variants.
+ * - ``ir.prim``
+ - Primitive expressions (shown as ``prim``), which compute scalar or
+ vector values of a primitive type
+ (:py:class:`ir.PrimType <tvm_tirx.ir.PrimType>`), such as ``float32``
+ or ``int32``. They concisely represent dynamic indices and element
+ computations, such as ``A[i] + 1`` in ``add_one``.
+ * - ``tirx``
+ - The types and operations specific to tensor programs, built on
+ ``ir``, such as ``tirx.Function`` and ``tirx.TensorType``.
+
+TIRx is **designed for extension**: downstream projects build on the IR
+the same way ``tirx`` builds on ``ir``, registering new operators, defining
new types, and adding
+their own TVMScript namespaces from Python, C++, or Rust, while reusing TIRx
+as a foundation.
+
+.. note::
+
+ IR nodes can also carry source locations
+ (:py:class:`ir.Location <tvm_tirx.ir.Location>`). For example, an
+ expression's ``loc`` field records where it came from in the source, so
+ a diagnostic can point back to the original code.
How This Guide Is Organized
---------------------------
@@ -263,8 +280,8 @@ and the second covers the tools we use to inspect and
transform it.
**Working with IR**
-- :doc:`structural_utilities` introduces structural walks, maps, visitors, and
- mutators for inspecting and rewriting IR.
+- :doc:`structural_utilities` introduces structural walks, maps, visitors,
+ mutators, and equality for traversing, rewriting, and comparing IR.
- :doc:`symbolic_integer_analysis` covers symbolic analysis and simplification
of integer indices and bounds, which compiler analyses and transformations
can use.
diff --git a/docs/ir_guides/printing_and_parsing.rst
b/docs/ir_guides/printing_and_parsing.rst
index dc05775e5..acdda3dc9 100644
--- a/docs/ir_guides/printing_and_parsing.rst
+++ b/docs/ir_guides/printing_and_parsing.rst
@@ -15,8 +15,6 @@
specific language governing permissions and limitations
under the License.
-.. _ir-construction-and-printing:
-
Printing and Parsing
====================
@@ -39,7 +37,7 @@ three roles:
`examples/ir_guides
<https://github.com/apache/tvm-tirx/tree/main/examples/ir_guides>`_
in the repository.
-Printing and Parsing
+Print and Parse Back
--------------------
We will use a CPU function that adds one to each tensor element:
@@ -69,16 +67,21 @@ The output is:
for i in range(n):
B[i] = A[i] + T.float32(1.0)
-The printer writes out everything the IR records. The constant ``1``
-becomes ``T.float32(1.0)``, its value with the type the addition uses. The
-symbolic dimension ``n`` is declared in ``[n]`` after the function name,
-using the type-parameter syntax introduced in Python 3.12. Each name in
-the brackets declares a symbolic variable as if by ``T.dynamic``, but
-scoped to this function, so the printed function is self-contained. On
-Python 3.12 and newer, we can also write this form ourselves. A symbol with a
dtype other than ``int32`` carries an
-annotation, as in ``[m: T.int64]``. The same method works on smaller
-objects. For example, ``add_one.params[0].ty.script()`` returns
-``T.Tensor((n,), T.float32)``.
+The printer writes out **everything the IR records**. The constant ``1``
+becomes ``T.float32(1.0)``, its value with the type the addition uses, and
+the symbolic dimension ``n`` is declared in ``[n]`` after the function
+name, so the printed function is self-contained.
+
+.. note::
+
+ ``[n]`` uses the type-parameter syntax introduced in Python 3.12. Each
+ name in the brackets declares a symbolic variable as if by
+ ``T.dynamic``, but scoped to this function. On Python 3.12 and newer, we
+ can also write this form ourselves. A symbol with a dtype other than
+ ``int32`` carries an annotation, as in ``[m: T.int64]``.
+
+The same method works on smaller objects. For example,
+``add_one.params[0].ty.script()`` returns ``T.Tensor((n,), T.float32)``.
A module prints the same way, as a class marked with ``@T.ir_module`` with
one ``@T.function`` per member:
@@ -124,12 +127,28 @@ IR again:
Printing and then parsing gives back the same IR. This is what lets a test
write its expected IR as TVMScript and compare it with the output of a
-pass. To store or transfer IR, use ``ir.save_json`` and ``ir.load_json``,
-which serialize the IR objects directly.
+pass.
+
+.. note::
+
+ To store or transfer IR, use ``ir.save_json`` and ``ir.load_json``,
+ which serialize the IR objects directly.
How the Parser Works
--------------------
+We value **predictability** over syntactic sugar. Each syntax form has one
+fixed meaning:
+
+- an annotation is a type;
+- an assignment binds a variable;
+- a subscript store writes a tensor;
+- a call is an ordinary Python call.
+
+Types such as ``T.Tensor`` and operators such as ``T.cuda.thread_idx`` are
+plain functions that return IR, so the parser has no rules specific to
+them.
+
The parser is a thin layer over an IR builder, the Python API that
constructs IR. The builder works with frames: entering a function, loop,
or region opens a frame, and the statements built inside it become its
@@ -167,17 +186,40 @@ becomes this builder program:
_result0.__name__ = 'add_one'
T.check_well_formed_(_result0)
-Each syntax form maps to builder calls: the definition becomes a
-``function_`` frame, parameter annotations become ``arg_`` calls with
-their types, the loop becomes a ``for_`` frame,
-and the store becomes ``setitem_``. The signature is declared before the
-body is built, and the finished function is checked for well-formedness.
-Expressions such as ``A[i] + 1`` stay as written: ``A`` is the IR
-variable returned by ``arg_``, so Python's ``[]`` and ``+`` build the IR
-expression directly.
-Because the program runs as normal Python, ``n`` and other names come from
-the surrounding scope. The real output also passes source positions to
-each call, which is how parsed IR gets its ``loc``.
+Concretely, the parser rewrites each syntax form into builder calls. The
+table also lists forms that this example does not use:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 45 55
+
+ * - TVMScript
+ - Builder program
+ * - ``@T.function`` definition
+ - A ``T.function_()`` frame, named by ``T.func_name_('add_one')``
+ * - Parameter ``A: T.Tensor((n,), T.float32)``
+ - ``T.arg_('A', T.Tensor((n,), T.float32))``
+ * - ``for i in range(n):``
+ - A ``T.for_(T.range_(n), names='i')`` frame
+ * - ``x = value``
+ - ``x = T.bind_(value, name='x')``
+ * - ``if cond:`` and ``else:``
+ - A ``T.if_`` frame with ``T.then_`` and ``T.else_``
+ * - ``B[i] = value``
+ - ``T.setitem_(value=value, target=B, key=i)``
+ * - Comparison ``i < n``
+ - ``T.lt_(i, n)``
+ * - Arithmetic ``A[i] + 1``
+ - Unchanged Python
+
+The signature is declared before the body is built, and the finished
+function is checked for well-formedness.
+
+Arithmetic such as ``A[i] + 1`` stays as written: ``A`` is the IR variable
+returned by ``arg_``, so Python's ``[]`` and ``+`` build the IR expression
+directly. Because the program runs as normal Python, ``n`` and other names
+come from the surrounding scope. The real output also passes source
+positions to each call, which is how parsed IR gets its ``loc``.
Because the builder program runs as ordinary Python, the parser supports
metaprogramming out of the box. A TVMScript function can read any Python
@@ -205,16 +247,25 @@ directly, without the parser:
``builder.get()`` returns a ``tirx.Function`` that is structurally equal
to ``add_one``. Hand-written builder code uses the user-facing API, such as
``T.serial`` and ``T.assign``, rather than the protocol hooks the parser
-emits. Writing builder code directly is sometimes helpful because
-it stays unambiguously in the Python domain. A plain Python ``for`` or
-``if`` runs while the IR is built, and IR control flow comes only from
-explicit frames such as ``T.serial`` and ``T.if_`` with ``T.then_``.
Developers and agents
-therefore know exactly what each line executes. To record source
-locations, wrap the construction in ``builder.line_trace_scope()``, and
-each statement then records its source line in ``loc``.
+emits.
+
+Writing builder code directly is sometimes helpful because it stays
+**unambiguously in the Python domain**. A plain Python ``for`` or ``if``
+runs while the IR is built, and IR control flow comes only from explicit
+frames such as ``T.serial`` and ``T.if_`` with ``T.then_``. Developers and
+agents therefore know exactly what each line executes.
+
+.. note::
+
+ To record source locations, wrap the construction in
+ ``builder.line_trace_scope()``, and each statement then records its
+ source line in ``loc``.
.. _tvmscript-region-syntax:
+Region Syntax
+~~~~~~~~~~~~~
+
Regions are built with frames in the same way. A region operator such as
``T.device_entry`` opens a frame, and the statements inside its ``with``
block become the region's body:
@@ -231,13 +282,6 @@ frame for the rest of its enclosing block, so
``T.device_entry(...)``
followed by the same statements at the same indentation produces the same
``ir.RegionStmt``.
-We value **predictability** over syntactic sugar. Each syntax form has one
-fixed meaning: an annotation is a type, an assignment binds a variable, a
-subscript store writes a tensor, and a call is an ordinary Python call.
-Types such as ``T.Tensor`` and operators such as ``T.cuda.thread_idx`` are
-plain functions that return IR, so the parser has no rules specific to
-them.
-
Extending the Syntax
--------------------
diff --git a/docs/ir_guides/statements.rst b/docs/ir_guides/statements.rst
index 9fb2aaab6..a100e9941 100644
--- a/docs/ir_guides/statements.rst
+++ b/docs/ir_guides/statements.rst
@@ -15,12 +15,10 @@
specific language governing permissions and limitations
under the License.
-.. _statements-and-control-flow:
-
Statements
==========
-This guide introduces **statements**, the IR nodes that describe
+This chapter introduces **statements**, the IR nodes that describe
actions: binding a value to a variable, writing to a tensor, repeating
work in a loop, or choosing between branches. TIRx provides a generic
set of statements that appear in most high-level kernel programs.
@@ -450,9 +448,3 @@ statements inside the region.
Regions and calls are the two ways operators extend TIRx: a call
applies an operator as an expression, and a region applies one as a
statement.
- :doc:`operators_calls_and_regions` covers how operators are registered and
- how calls and regions use them.
- :ref:`Region Syntax <tvmscript-region-syntax>` explains how TVMScript
- writes a region, and :doc:`../advanced/defining_ir_extensions` walks
- through a complete extension with its own type, operators, and
- lowering pass.
diff --git a/docs/ir_guides/structural_utilities.rst
b/docs/ir_guides/structural_utilities.rst
index 33ce82235..776897087 100644
--- a/docs/ir_guides/structural_utilities.rst
+++ b/docs/ir_guides/structural_utilities.rst
@@ -204,6 +204,22 @@ The walk ends at the first write, and ``found.value``
holds it. A walk that
finishes without an interrupt returns ``None``, so the result also tells us
whether anything was found.
+In summary, a walk callback steers the traversal with its return value:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 35 65
+
+ * - Return value
+ - Effect on the walk
+ * - ``None``
+ - Continue.
+ * - ``WalkResult.SKIP``
+ - In pre-order, do not descend into the node's children. In
+ post-order, the children are already visited, so it has no effect.
+ * - ``VisitInterrupt(payload)``
+ - Stop the whole walk; ``structural_walk`` returns the interrupt.
+
Structural Map
--------------
@@ -237,10 +253,14 @@ The output is:
Every occurrence of ``n``, in both parameter shapes and the loop bound,
becomes ``128``. Because IR is immutable, the map returns a new function
-and leaves ``add_one`` unchanged. Note that a map callback must always
-return a node, even when it has nothing to change: returning its input
-keeps that node as it is, which is what ``specialize`` does for every other
-variable.
+and leaves ``add_one`` unchanged.
+
+.. note::
+
+ A map callback must always return a node, even when it has nothing to
+ change: returning its input keeps that node as it is, which is what
+ ``specialize`` does for every other variable. Returning ``None``,
+ ``WalkResult.SKIP``, or a ``VisitInterrupt`` raises a ``TypeError``.
Immutability also keeps rewrites cheap: a map does not copy the whole
program. It rebuilds only the nodes on the path from each replacement up to
@@ -251,12 +271,10 @@ or compare.
Structural Equality and Hashing
-------------------------------
-Once we can rewrite programs, we need a way to check the result. To test a
-transformation, we write down the program we expect and compare it with the
-one the transformation produced. The two programs are built separately, so
-they are different objects even when they describe the same computation,
-and comparing them by identity cannot work. Instead, we compare them by
-structure. Let's build the same function twice with a small helper:
+To test a transformation, we compare its result with a program we write
+down ourselves. The two are different objects even when they describe the
+same computation, so we compare them **by structure**, not by identity.
+Let's build the same function twice with a small helper:
.. literalinclude:: ../../examples/ir_guides/structural_utilities.py
:language: python
@@ -279,21 +297,27 @@ gives them the same hash. Each call to ``make_add`` also
creates its own
in one program with its counterpart in the other, so these different
variable objects do not make the programs unequal.
-A variable that is not defined anywhere in the compared values is free and
-has no definition to match. By default, a free variable is equal only to
-itself. ``map_free_vars=True`` pairs free variables up as well:
+.. note::
-.. literalinclude:: ../../examples/ir_guides/structural_utilities.py
- :language: python
- :start-after: # [free_vars.begin]
- :end-before: # [free_vars.end]
+ A variable that is not defined anywhere in the compared values is
+ **free** and has no definition to match. By default, a free variable is
+ equal only to itself. ``map_free_vars=True`` pairs free variables up as
+ well:
-The output is:
+ .. literalinclude:: ../../examples/ir_guides/structural_utilities.py
+ :language: python
+ :start-after: # [free_vars.begin]
+ :end-before: # [free_vars.end]
-.. code-block:: text
+ The output is:
- False
- True
+ .. code-block:: text
+
+ False
+ True
+
+ ``x + 1`` and ``y + 1`` differ by default; with the flag, ``x`` pairs
+ with ``y``.
In a test, we usually want more than a yes-or-no answer.
:py:func:`ir.assert_structural_equal <tvm_tirx.ir.assert_structural_equal>`
@@ -347,55 +371,53 @@ The two ``make_add(1)`` functions collapse into one
entry, and
Definitions and Uses
--------------------
-A variable, represented by :py:class:`ir.Var <tvm_tirx.ir.Var>`, can
-appear in a program for one of two purposes. A **definition site**
-introduces the variable, and a **use site** reads its value. In
-``add_one``, the parameters are the definition sites of ``A``, ``B``, and
-the length ``n``, and the loop is the definition site of ``i``. The loop
-bound ``range(n)`` is a use site of ``n``, and ``B[i] = A[i] + 1``
-contains use sites of ``A``, ``B``, and ``i``. Each variable has exactly
-one definition site and can have any number of use sites.
-
-In the IR, every site refers to the same ``ir.Var`` object. Many tasks need
-to tell the two purposes apart. To rename ``i``, we create the new variable
-once at its definition site and point every use site at it. To find an
-unused variable, we look for a definition site with no use sites. To check
-whether two programs are the same, we match each definition site in one
-with its counterpart in the other. The variable object alone does not say
-which purpose it serves at a given site, but its position in the program
-does.
+A variable (:py:class:`ir.Var <tvm_tirx.ir.Var>`) appears in a program in
+two roles: a **definition site** introduces it, and a **use site** reads
+it. In ``add_one``, the parameters define ``A``, ``B``, and ``n``, and the
+loop defines ``i``; ``range(n)`` and ``B[i] = A[i] + 1`` use them. Each
+variable has one definition site and any number of use sites.
+
+**Every site refers to the same ir.Var object**, so the object alone cannot
+tell us its role; its position can. Renaming, finding unused variables, and
+comparing programs depend on that position.
TIRx IR nodes use their structural hooks to mark which of their fields are
definition sites. As the utilities traverse the program, they report the
role at each position as a **definition region kind**:
-- ``NONE`` marks a use site, such as ``n`` in ``range(n)`` and ``A``,
- ``B``, and ``i`` in the body.
-- ``DEF_SIMPLE`` defines exactly the variable at the site, such as the
- loop variable ``i``. A ``Bind`` statement defines its variable the same
- way.
-- ``DEF_PATTERN`` defines every variable in the site by matching,
- including those in its type.
-
-A ``DEF_PATTERN`` site is a pattern that is matched against the values
-passed in. In ``add_one``, the parameter ``A: T.Tensor((n,), T.float32)``
-is matched against the tensor passed for ``A``. This defines both ``A`` and
-its length ``n``, so ``n`` needs no separate declaration. A variable can
-appear more than once in a pattern. The first appearance defines it, and
-each later appearance adds a constraint that is checked at runtime. Here,
-``n`` also appears in ``B``'s shape, so ``B`` must have the same length.
-When we call ``add_one`` with an ``A`` of length 100, ``n`` is bound to
-100, and passing a ``B`` of length 50 raises an error.
-
-Importantly, the structural utilities are themselves aware of definitions
-and uses. When :py:func:`structural_map <tvm_ffi.structural_map>` rewrites
-the type of a variable at its definition site, it creates a new variable
-with the new type and remaps every use site to it. For example, replacing
-``n`` with 128 in ``add_one`` changes the type of ``A`` from
-``T.Tensor((n,), T.float32)`` to ``T.Tensor((128,), T.float32)``, so the
-result has a new ``A``, and ``A[i]`` in the body refers to that new ``A``.
-Our callback only replaced ``n``; the mapper kept every definition and its
-uses consistent.
+.. list-table::
+ :header-rows: 1
+ :widths: 20 45 35
+
+ * - Kind
+ - Meaning
+ - In ``add_one``
+ * - ``NONE``
+ - A use site.
+ - ``n`` in ``range(n)``; ``A``, ``B``, and ``i`` in the body
+ * - ``DEF_SIMPLE``
+ - Defines exactly the variable at the site.
+ - The loop variable ``i``; a ``Bind`` works the same way
+ * - ``DEF_PATTERN``
+ - Defines every variable in the site by matching, including those in
+ its type.
+ - The parameters ``A`` and ``B``, with ``n`` in their shapes
+
+.. note::
+
+ A ``DEF_PATTERN`` site is a pattern matched against the values passed
+ in. The parameter ``A: T.Tensor((n,), T.float32)`` defines both ``A``
+ and its length ``n``, so ``n`` needs no separate declaration. A later
+ appearance of ``n``, as in ``B``'s shape, adds a runtime check: calling
+ ``add_one`` with an ``A`` of length 100 binds ``n`` to 100, so a ``B`` of
+ length 50 raises an error.
+
+**The utilities keep definitions and uses consistent for us.** When
+:py:func:`structural_map <tvm_ffi.structural_map>` changes the type of a
+variable at its definition site, it creates a new variable with the new
+type and repoints every use site to it. Replacing ``n`` with 128, for
+example, changes ``A``'s type to ``T.Tensor((128,), T.float32)``, so the
+result has a new ``A``, and ``A[i]`` in the body refers to it.
To receive the kind in our own function, we register it through
``with_def_region_kind`` instead of ``callbacks``, and the function is
@@ -420,7 +442,7 @@ The output is:
i DEF_SIMPLE
``n`` appears twice, once for each shape: the first defines it, and the
-second is the constraint. Every other site is a use site, reported as
+second is the runtime check. Every other site is a use site, reported as
``NONE``.
``structural_map`` accepts the same form. A rename leaves each variable's
@@ -454,13 +476,15 @@ Advanced Control: Visit and Mutate
----------------------------------
Walk and map recurse for us, and that covers most tasks. Some analyses,
-however, need to know where they are in the program, not just which node
-they are looking at. To report the loops around each write, for example, a
-callback must know when the traversal enters a loop and when it leaves
-again. :py:func:`tvm_ffi.structural_visit <tvm_ffi.structural_visit>` and
+however, need to know **where they are** in the program, not just which
+node they are looking at. To report the loops around each write, for
+example, a callback must know when the traversal enters a loop and when it
+leaves again.
+
+:py:func:`tvm_ffi.structural_visit <tvm_ffi.structural_visit>` and
:py:func:`tvm_ffi.structural_mutate <tvm_ffi.structural_mutate>` hand that
-control to us: a matching callback also receives the visitor or mutator and
-decides when, or whether, to process the node's children. If it never does,
+control to us. A matching callback also receives the visitor or mutator and
+decides when, or whether, to process the node's children; if it never does,
the children are skipped. Let's record the loops around each write:
.. literalinclude:: ../../examples/ir_guides/structural_utilities.py
@@ -482,10 +506,12 @@ their usual traversal, without calling the node's own
callback again;
:py:meth:`default_mutate <tvm_ffi.StructuralMutator.default_mutate>` does
the same for a mutator.
-Because the callback now drives the recursion, it also takes on what the
-walker used to do for us. ``visit_for`` returns the result of
-``default_visit``, so that a ``VisitInterrupt`` from a child reaches the
-caller.
+.. note::
+
+ Because the callback now drives the recursion, it also takes on what the
+ walker used to do for us. Return the result of ``default_visit``, as
+ ``visit_for`` does, so that a ``VisitInterrupt`` from a child reaches the
+ caller.
Mutate brings the same control to rewriting. A mutate callback returns the
new node, as in a map, but it can also set up state before its children are
diff --git a/docs/ir_guides/symbolic_integer_analysis.rst
b/docs/ir_guides/symbolic_integer_analysis.rst
index e7efebcb2..35917814b 100644
--- a/docs/ir_guides/symbolic_integer_analysis.rst
+++ b/docs/ir_guides/symbolic_integer_analysis.rst
@@ -106,10 +106,14 @@ The output is:
False
-A ``False`` result means the analyzer could not establish the condition, not
-that the condition is false. Here the analyzer knows nothing about ``i``. As
-far as it can tell, ``i`` could be any integer. The bound depends on the
-loop, and the analyzer needs that information as a fact.
+Here the analyzer knows nothing about ``i``. As far as it can tell, ``i``
+could be any integer. The bound depends on the loop, and the analyzer needs
+that information as a fact.
+
+.. note::
+
+ A ``False`` result means the analyzer could not establish the condition,
+ not that the condition is false.
Constraint Context
------------------
@@ -163,21 +167,44 @@ The output is:
Inside the branch, the analyzer proves ``i % 4 < 3``. Synchronizing with
``loop_context`` again models leaving the branch, and the fact no longer
-applies. An analyzer uses only the facts of the context it was last
-synchronized with, which lets the same analyzer serve different points in a
-program.
+applies.
+
+.. note::
+
+ An analyzer uses only the facts of the context it was last synchronized
+ with, which lets the same analyzer serve different points in a program.
Structural Walk and Map with Constraint Context
-----------------------------------------------
Building contexts by hand works for one query, but a pass over a whole
function would have to rebuild the facts at every node. The structural
-utilities can collect them for us.
-:py:func:`tirx.structural_walk_with_constraint_context
<tvm_tirx.tirx.structural.structural_walk_with_constraint_context>`
-and
-:py:func:`tirx.structural_map_with_constraint_context
<tvm_tirx.tirx.structural.structural_map_with_constraint_context>`
-work like walk and map, and also pass the facts valid at each node to our
-callback.
+utilities in ``tvm_tirx.tirx`` **collect the facts for us**: each works like
+its TVM-FFI counterpart and also gives our callback the facts valid at each
+node.
+
+.. list-table::
+ :header-rows: 1
+ :widths: 46 22 32
+
+ * - Utility
+ - Works like
+ - Facts reach the callback as
+ * - :py:func:`structural_walk_with_constraint_context
<tvm_tirx.tirx.structural.structural_walk_with_constraint_context>`
+ - ``structural_walk``
+ - Callback is ``f(value, constraints)``
+ * - :py:func:`structural_map_with_constraint_context
<tvm_tirx.tirx.structural.structural_map_with_constraint_context>`
+ - ``structural_map``
+ - Callback is ``f(value, constraints)``
+ * - :py:func:`structural_visit_with_constraint_context
<tvm_tirx.tirx.structural.structural_visit_with_constraint_context>`
+ - ``structural_visit``
+ - ``visitor.constraints``
+ * - :py:func:`structural_mutate_with_constraint_context
<tvm_tirx.tirx.structural.structural_mutate_with_constraint_context>`
+ - ``structural_mutate``
+ - ``mutator.constraints``
+
+Here ``constraints`` is a ``ConstraintContext``, which our callback reads but
+must not add facts to.
Compilers often insert bounds checks to stay safe, such as the
``if i // 4 < n`` below. Here the loop already guarantees the condition, so
@@ -228,12 +255,6 @@ The output is:
for i in range(n * 4):
B[i // 4, i % 4] = A[i // 4, i % 4] + T.float32(1.0)
-Visit and mutate variants,
-:py:func:`structural_visit_with_constraint_context
<tvm_tirx.tirx.structural.structural_visit_with_constraint_context>`
-and
-:py:func:`structural_mutate_with_constraint_context
<tvm_tirx.tirx.structural.structural_mutate_with_constraint_context>`,
-provide the same facts to callbacks that control their own traversal.
-
A transformation like ``remove_redundant_guard`` is an ordinary function, and
we can call it directly. When it helps, we can also wrap it as a pass with
:py:func:`~tvm_tirx.tirx.transform.function_pass` and compose it with other
diff --git a/docs/ir_guides/types_and_expressions.rst
b/docs/ir_guides/types_and_expressions.rst
index 07d00a988..0a9142ab0 100644
--- a/docs/ir_guides/types_and_expressions.rst
+++ b/docs/ir_guides/types_and_expressions.rst
@@ -15,8 +15,6 @@
specific language governing permissions and limitations
under the License.
-.. _expressions-and-types:
-
Types and Expressions
=====================
@@ -299,7 +297,8 @@ The following table lists the primitive expression node
classes in
- TVMScript syntax
* - Arithmetic
- ``Add``, ``Sub``, ``Mul``, ``Div``, ``Mod``, ``FloorDiv``,
- ``FloorMod``, ``Min``, ``Max``
+ ``FloorMod``, ``Min``, ``Max`` (``Div`` and ``Mod`` follow C
+ semantics; ``FloorDiv`` and ``FloorMod`` round down)
- ``+``, ``-``, ``*``, ``/`` (float), ``//``, ``%`` (int)
* - Comparison
- ``EQ``, ``NE``, ``LT``, ``LE``, ``GT``, ``GE``
@@ -564,16 +563,15 @@ The output is:
└── args[0]: TensorLoad(source=A, indices=[i])
└── ty: PrimType(dtype=float32)
-The call's ``op`` is an ``ir.GlobalVar``, the module's name for the
-helper, explained in :ref:`Global Variables <ir-global-var>`. Its type is
-the helper's ``FuncType``, and the call takes its result type from that
-signature.
+The call's ``op`` is an ``ir.GlobalVar``, an expression that refers to
+the helper by its name in the module. Its type is the helper's
+``FuncType``, and the call takes its result type from that signature.
.. note::
- Operators are the main extension point of TIRx.
- :doc:`operators_calls_and_regions` covers how operators are registered and
- how their attributes define type rules, effects, and lowering.
+ Operators are the main extension point of TIRx. A library registers new
+ operators, and their attributes define type rules, effects, and
+ lowering.
.. _user-ir-dependent-type:
@@ -628,13 +626,11 @@ called, ``n`` takes the length of ``A``, and ``B`` must
have the same
length. The loop body can also use ``n`` as an ordinary int32 value,
here as the loop bound.
-TIRx analysis and transformation utilities are designed with dependent
-types in mind. When a transformation replaces a variable, it updates
-every type that refers to that variable, not only the expressions. For
-example, replacing ``n`` with ``16`` in the add-one function changes the
-loop bound to ``16`` and also changes the types of ``A`` and ``B`` to
-``T.Tensor((16,), T.float32)``. Types and expressions therefore stay
-consistent after each rewrite.
+TIRx rewriting utilities are designed with dependent types in mind: when
+a rewrite replaces a variable, it updates every type that refers to it,
+not only the expressions. Replacing ``n`` with ``16`` in the add-one
+function changes the loop bound and also the types of ``A`` and ``B`` to
+``T.Tensor((16,), T.float32)``.
Customization
-------------
diff --git a/docs/reference/python/ffi.structural.rst
b/docs/reference/python/ffi.structural.rst
index 4106f63b0..c96cc94a4 100644
--- a/docs/reference/python/ffi.structural.rst
+++ b/docs/reference/python/ffi.structural.rst
@@ -68,18 +68,28 @@ Structural Hooks
An object type tells these utilities how to handle it through structural
hooks, which it can register from C++ or Python:
-- ``__s_visit__`` visits the object's children and marks which fields
- define variables.
-- ``__s_mutate__`` rewrites the object's children. An optional
- ``__s_maybe_inplace_mutate__`` reuses the object's storage when it is
- safe to do so.
-- ``__s_equal__`` and ``__s_hash__`` define structural equality and
- hashing.
-
-Every TIRx IR node registers compiled hooks. Other objects fall back to
-their reflected fields. TIRx marks binding variables, such as the variable
-of a ``For`` loop or a ``Bind`` statement, as ``DEF_SIMPLE`` definitions,
-and function parameters as ``DEF_PATTERN`` definitions, which also define
-the variables in their types. See `Custom Visit and Mutation Hooks
+.. list-table::
+ :header-rows: 1
+ :widths: 35 65
+
+ * - Hook
+ - Purpose
+ * - ``__s_visit__``
+ - Visits the object's children and marks which fields define
+ variables.
+ * - ``__s_mutate__``
+ - Rewrites the object's children.
+ * - ``__s_maybe_inplace_mutate__``
+ - Optional. Reuses the object's storage when it is safe to do so.
+ * - ``__s_equal__`` and ``__s_hash__``
+ - Define structural equality and hashing.
+
+Every TIRx IR node registers compiled hooks; other objects fall back to
+their reflected fields.
+
+TIRx marks loop and ``Bind`` variables as ``DEF_SIMPLE``; it marks function
+parameters, including the variables in their types, as ``DEF_PATTERN``.
+
+See `Custom Visit and Mutation Hooks
<https://tvm.apache.org/ffi/concepts/structural_eq_hash.html#custom-visit-and-mutation-hooks>`_
-in the TVM-FFI documentation for how to write them.
+in the TVM-FFI documentation for how to write hooks.
diff --git a/docs/reference/python/ir.transform.rst
b/docs/reference/python/ir.transform.rst
index 512a71da2..cab17b655 100644
--- a/docs/reference/python/ir.transform.rst
+++ b/docs/reference/python/ir.transform.rst
@@ -47,13 +47,12 @@ the module before the first pass and after the last one:
from tvm_tirx import ir, tirx
from tvm_tirx.ir.transform import PassContext, Sequential
from tvm_tirx.ir.transform.instrument import PassPrintingInstrument
- from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
n = T.dynamic("n")
- @I.ir_module
+ @T.ir_module
class Module:
@T.function
def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
@@ -108,10 +107,6 @@ The output is:
B[i] = A[i] + T.float32(1.0)
``Module`` itself is unchanged, and ``result`` holds the transformed module.
-Inside a :py:class:`Sequential`, a pass runs only when its optimization level
-is at most the context's ``opt_level`` or its name is in ``required_pass``,
-and never when its name is in ``disabled_pass``. Calling a pass directly, as
-in ``specialize(Module)``, always runs it.
Classes
-------
diff --git a/examples/ir_guides/functions_and_modules.py
b/examples/ir_guides/functions_and_modules.py
index 8064c26c7..3f35573e3 100644
--- a/examples/ir_guides/functions_and_modules.py
+++ b/examples/ir_guides/functions_and_modules.py
@@ -37,10 +37,12 @@ class Module:
# [module.end]
+# [ty.begin]
add_one = Module["add_one"]
+print(add_one.ty)
+# [ty.end]
assert [p.name for p in add_one.params] == ["A", "B"]
assert len(add_one.ret_type.fields) == 0
-print(add_one.ty)
assert add_one.attrs["global_symbol"] == "add_one"
