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 b65b88a91 [DOCS] Polish IR Overview (#120)
b65b88a91 is described below
commit b65b88a91b44ba3b1d1438b6fd4f30b1a3dc4687
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 28 22:22:17 2026 -0400
[DOCS] Polish IR Overview (#120)
Rewrite the IR Overview as a short narrative introduction to TIRx IR.
- **IR Overview:**
- Opens with the design goal: developers and agents can access every
stage of compilation in the same way.
- Introduces IR nodes and TVMScript, and shows the CUDA `add_one`
example.
- Lists the main elements (types, expressions, statements, functions,
modules) with links to their API classes.
- Class hierarchy: compact layout that fits the page width, `prim.*`
labels that give each operation's Python operator, and a paragraph on
the `ir`, `ir.prim`, and `tirx` namespaces and how downstream projects
extend the IR.
- A short note on source locations, and a bullet list for What to Read
Next.
- **Source locations:** the detailed walkthrough moves from Types and
Expressions to Printing and Parsing, merged with the parser `filename`
section there.
- **Quick Start:** the see-also link now says "IR Overview", matching
the page title.
Validation:
- The docs build passes with `-W`.
- pre-commit passes.
- The source-location output was checked against the examples.
---
docs/get_started/quick_start.rst | 2 +-
docs/ir/overview.rst | 304 +++++++++++++++++++-------------------
docs/ir/printing_and_parsing.rst | 49 +++++-
docs/ir/types_and_expressions.rst | 45 ------
4 files changed, 198 insertions(+), 202 deletions(-)
diff --git a/docs/get_started/quick_start.rst b/docs/get_started/quick_start.rst
index bc574a8f1..a7490f386 100644
--- a/docs/get_started/quick_start.rst
+++ b/docs/get_started/quick_start.rst
@@ -186,7 +186,7 @@ program, while the original ``add_one`` stays unchanged.
.. seealso::
- :doc:`IR Walkthrough </ir/overview>` for a more in-depth introduction to the
+ :doc:`IR Overview </ir/overview>` for a more in-depth introduction to the
main IR objects and their roles in the overall program structure.
diff --git a/docs/ir/overview.rst b/docs/ir/overview.rst
index 1d371dfbf..b8ace5c16 100644
--- a/docs/ir/overview.rst
+++ b/docs/ir/overview.rst
@@ -18,164 +18,164 @@
IR Overview
===========
-An intermediate representation (IR) describes a program in a form that compiler
-components can inspect and transform. Throughout this walkthrough, we'll use
-the CPU and GPU examples from :doc:`/get_started/quick_start` to introduce
-the main IR objects and show how they fit together.
-
-TVMScript constructs IR objects to describe the computation in these examples.
-PyTorch tensors hold the input and output data, and the compiled functions
-perform the computation when called.
-
-TIRx shares its IR object system across C++, Python, and Rust through TVM FFI,
-so compiler components written in different languages can work together. The
-same infrastructure supports custom types and operators and can be reused to
-build analyses, transformation passes, and backends.
-
-The main components of TIRx IR are:
-
-.. list-table::
- :header-rows: 1
- :widths: 15 25 30 30
-
- * - Component
- - Main classes
- - Purpose
- - Example
- * - Expressions
- - ``ir.Expr``
- - Represent values and computations.
- - ``A[i] + 1``
- * - Types
- - ``ir.Type``
- - Describe value properties, such as dtypes and tensor shapes.
- - ``T.float32``
- * - Statements
- - ``ir.Stmt``
- - Describe execution order and control flow.
- - ``B[i] = A[i] + 1``
- * - Source locations
- - ``ir.Location``
- - Record where IR nodes originate in the source.
- - The filename and line number of ``A[i] + 1``.
- * - Functions
- - ``ir.BaseFunc``, ``tirx.Function``
- - Define parameters and a computation body.
- - ``add_one`` defined with ``@T.function``.
- * - Modules
- - ``ir.IRModule``
- - Organize named functions into a program.
- - ``ir.IRModule({"add_one": add_one})``
-
-The hierarchy below shows the inheritance relationships between these classes:
+One primary design goal of TIRx is to allow developers and agents to quickly
+access every stage of compilation, not only its input and output, and to
+work with every stage in the same way, so they can 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 and
+maps it to the target program. At every stage, the tree can be printed as
+Python code, in a syntax we call TVMScript (provided by the
+``tvm_tirx.script`` module). For example, the following TVMScript program
+defines ``add_one`` for a GPU:
+
+.. literalinclude:: ../../examples/quick_start/add_one_cuda.py
+ :language: python
+ :start-at: from tvm_tirx.script import tirx as T
+ :end-at: B[i] = A[i] + 1
+
+``add_one`` adds one to every element of a tensor ``A`` of length ``n`` and
+writes the result to ``B``. Each GPU thread computes its index ``i`` and
+handles one element. At a high level, TIRx IR
+contains the following elements:
+
+- **Types** (:py:class:`ir.Type <tvm_tirx.ir.Type>`) describe the values that
+ expressions produce, such as the scalar type ``T.float32`` and the tensor
+ type ``T.Tensor((n,), T.float32)``.
+- **Expressions** (:py:class:`ir.Expr <tvm_tirx.ir.Expr>`) compute values,
+ such as ``A[i] + 1`` and the call ``T.cuda.thread_idx("x")`` to a registered
+ operator (:py:class:`ir.Op <tvm_tirx.ir.Op>`).
+- **Statements** (:py:class:`ir.Stmt <tvm_tirx.ir.Stmt>`) describe what the
+ program does and in what order, such as the ``if`` branch and the write
+ ``B[i] = A[i] + 1``.
+- **Functions** (:py:class:`ir.BaseFunc <tvm_tirx.ir.BaseFunc>`) wrap a body
+ of statements with parameters, such as ``add_one``.
+- **Modules** (:py:class:`ir.IRModule <tvm_tirx.ir.IRModule>`) group named
+ functions into a complete program.
+
+Class Hierarchy
+---------------
+
+The following diagram shows the class hierarchy of the IR nodes:
.. code-block:: text
- ffi.Object TVM FFI base object
- ├── ir.Expr Expressions
- │ ├── ir.Constant Literal constants
- │ │ ├── ir.IntImm Integer constants
- │ │ ├── ir.FloatImm Floating-point constants
- │ │ ├── ir.StringImm String constants
- │ │ ├── ir.DataTypeImm Dtype constants
- │ │ └── ir.GenericConst Constants with
arbitrary payloads
- │ ├── ir.Var Local variables
- │ ├── ir.GlobalVar Global function
references
- │ ├── ir.Tuple Tuple construction
- │ ├── ir.TupleGetItem Tuple field access
- │ ├── ir.Call Calls
- │ ├── ir.Op Registered operators
- │ ├── ir.TensorLoad Tensor element reads
- │ ├── ir.TensorRegion Tensor regions
- │ ├── ir.BaseFunc Function definitions
- │ │ └── tirx.Function Tensor-program functions
- │ ├── ir.StagingExpr Expressions removed
during compilation
- │ │ └── ir.LambdaExpr Construction-time
lambdas
- │ ├── ir.OpaqueExpr Opaque
construction-time expressions
- │ ├── prim.Cast Dtype conversion
- │ ├── prim.Add Addition
- │ ├── prim.Sub Subtraction
- │ ├── prim.Mul Multiplication
- │ ├── prim.Div Division
- │ ├── prim.Mod Remainder
- │ ├── prim.FloorDiv Floor division
- │ ├── prim.FloorMod Floor remainder
- │ ├── prim.Min Minimum
- │ ├── prim.Max Maximum
- │ ├── prim.LShift Left bit shift
- │ ├── prim.RShift Right bit shift
- │ ├── prim.BitwiseAnd Bitwise AND
- │ ├── prim.BitwiseOr Bitwise OR
- │ ├── prim.BitwiseXor Bitwise XOR
- │ ├── prim.BitwiseNot Bitwise NOT
- │ ├── prim.EQ Equal
- │ ├── prim.NE Not equal
- │ ├── prim.LT Less than
- │ ├── prim.LE Less than or equal
- │ ├── prim.GT Greater than
- │ ├── prim.GE Greater than or equal
- │ ├── prim.And Logical AND
- │ ├── prim.Or Logical OR
- │ ├── prim.Not Logical NOT
- │ ├── prim.Select Conditional value
selection
- │ ├── prim.Let Expression-local binding
- │ ├── prim.Ramp Vector arithmetic
progression
- │ ├── prim.Broadcast Scalar broadcast to
vector lanes
- │ └── prim.Shuffle Vector lane selection
- ├── ir.Type Types
- │ ├── ir.MissingType Unresolved type
information
- │ ├── ir.NoneType The type of None
- │ ├── ir.AnyType Any value
- │ ├── ir.OpaqueType Opaque
construction-time types
- │ ├── ir.PrimType Scalar and vector dtypes
- │ ├── ir.StringType String values
- │ ├── ir.PtrType Pointer types
- │ ├── ir.TupleType Tuple field types
- │ ├── ir.FuncType Function signatures
- │ ├── ir.TensorRegionType Tensor region values
- │ ├── tirx.TensorType Tensor types
- │ └── tirx.TensorMapType Tensor map types
- ├── ir.Stmt Statements
- │ ├── ir.SeqStmt Statement sequences
- │ ├── ir.Bind Variable definitions
- │ ├── ir.RegionStmt Operator-defined
lexical regions
- │ ├── ir.AssertStmt Program assertions
- │ ├── ir.TensorStore Tensor element writes
- │ ├── ir.Evaluate Evaluation with an
unused result
- │ ├── ir.If Conditional branches
- │ ├── ir.For Counted loops
- │ ├── ir.While Conditional loops
- │ ├── ir.Return Function returns
- │ ├── ir.Break Loop exits
- │ └── ir.Continue Next loop iterations
- ├── ir.Location Source locations
- │ ├── ir.UnknownLoc Unavailable source
information
- │ ├── ir.SourceLoc Source ranges
- │ └── ir.CallSiteLoc Callee and caller
locations
- └── ir.IRModule Global function
definitions
-
-The hierarchy shows inheritance between classes. The smaller diagrams in
-the following chapters show selected fields of individual objects. Each
-branch names a field, list entry, or mapping entry; for example,
-``seq[0]`` can refer to an ``ir.Bind`` stored inside an ``ir.SeqStmt``.
+ ffi.Object TVM-FFI base object
+ ├── ir.Expr Expressions
+ │ ├── ir.Constant Literal constants
+ │ │ ├── ir.IntImm Integer constants
+ │ │ ├── ir.FloatImm Floating-point constants
+ │ │ ├── ir.StringImm String constants
+ │ │ ├── ir.DataTypeImm Dtype constants
+ │ │ └── ir.GenericConst Constants with arbitrary payloads
+ │ ├── ir.Var Local variables
+ │ ├── ir.GlobalVar Global function references
+ │ ├── ir.Tuple Tuple construction
+ │ ├── ir.TupleGetItem Tuple field access
+ │ ├── ir.Call Calls
+ │ ├── ir.Op Registered operators
+ │ ├── ir.TensorLoad Tensor element reads
+ │ ├── ir.TensorRegion Tensor regions
+ │ ├── ir.BaseFunc Function definitions
+ │ │ └── tirx.Function Tensor-program functions
+ │ ├── ir.StagingExpr Expressions removed during compilation
+ │ │ └── ir.LambdaExpr Construction-time lambdas
+ │ ├── ir.OpaqueExpr Opaque construction-time expressions
+ │ ├── prim.Cast Primitive dtype conversion
+ │ ├── prim.Add Primitive addition, a + b
+ │ ├── prim.Sub Primitive subtraction, a - b
+ │ ├── prim.Mul Primitive multiplication, a * b
+ │ ├── prim.Div Primitive division in C semantics
+ │ ├── prim.Mod Primitive remainder in C semantics
+ │ ├── prim.FloorDiv Primitive floor division, a // b
+ │ ├── prim.FloorMod Primitive floor remainder, a % b
+ │ ├── prim.Min Primitive minimum, min(a, b)
+ │ ├── prim.Max Primitive maximum, max(a, b)
+ │ ├── prim.LShift Primitive left shift, a << b
+ │ ├── prim.RShift Primitive right shift, a >> b
+ │ ├── prim.BitwiseAnd Primitive bitwise AND, a & b
+ │ ├── prim.BitwiseOr Primitive bitwise OR, a | b
+ │ ├── prim.BitwiseXor Primitive bitwise XOR, a ^ b
+ │ ├── prim.BitwiseNot Primitive bitwise NOT, ~a
+ │ ├── prim.EQ Primitive equal, a == b
+ │ ├── prim.NE Primitive not equal, a != b
+ │ ├── prim.LT Primitive less than, a < b
+ │ ├── prim.LE Primitive less or equal, a <= b
+ │ ├── prim.GT Primitive greater than, a > b
+ │ ├── prim.GE Primitive greater or equal, a >= b
+ │ ├── prim.And Primitive logical AND
+ │ ├── prim.Or Primitive logical OR
+ │ ├── prim.Not Primitive logical NOT
+ │ ├── prim.Select Primitive conditional value selection
+ │ ├── prim.Let Primitive expression-local binding
+ │ ├── prim.Ramp Primitive vector arithmetic progression
+ │ ├── prim.Broadcast Primitive scalar broadcast to vector
+ │ └── prim.Shuffle Primitive vector lane selection
+ ├── ir.Type Types
+ │ ├── ir.MissingType Unresolved type information
+ │ ├── ir.NoneType The type of None
+ │ ├── ir.AnyType Any value
+ │ ├── ir.OpaqueType Opaque construction-time types
+ │ ├── ir.PrimType Scalar and vector dtypes
+ │ ├── ir.StringType String values
+ │ ├── ir.PtrType Pointer types
+ │ ├── ir.TupleType Tuple field types
+ │ ├── ir.FuncType Function signatures
+ │ ├── ir.TensorRegionType Tensor region values
+ │ ├── tirx.TensorType Tensor types
+ │ └── tirx.TensorMapType Tensor map types
+ ├── ir.Stmt Statements
+ │ ├── ir.SeqStmt Statement sequences
+ │ ├── ir.Bind Variable definitions
+ │ ├── ir.RegionStmt Operator-defined lexical regions
+ │ ├── ir.AssertStmt Program assertions
+ │ ├── ir.TensorStore Tensor element writes
+ │ ├── ir.Evaluate Evaluation with an unused result
+ │ ├── ir.If Conditional branches
+ │ ├── ir.For Counted loops
+ │ ├── ir.While Conditional loops
+ │ ├── ir.Return Function returns
+ │ ├── ir.Break Loop exits
+ │ └── ir.Continue Next loop iterations
+ ├── ir.Location Source locations
+ │ ├── ir.UnknownLoc Unavailable source information
+ │ ├── ir.SourceLoc Source ranges
+ │ └── 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.
+
+Every IR node can also carry a source location
+(:py:class:`ir.Location <tvm_tirx.ir.Location>`) in its ``loc`` field, which
+records where the node comes from in the source code and can be used for
+diagnostics.
.. _choose-a-chapter:
What to Read Next
-----------------
-The CPU example provides the tensor types, arithmetic, and loop; the GPU
-example adds a variable binding, a conditional, and calls for thread
-coordinates. Short additional examples introduce scalar and vector
-values through the IR API.
-
-* The :doc:`types_and_expressions` chapter explains how values, computations,
- and tensor shapes are represented.
-* The :doc:`statements` chapter covers statement bodies with bindings, tensor
- writes, branches, and loops.
-* The :doc:`operators_and_calls` chapter introduces registered operators and
- function calls.
-* The :doc:`functions_and_modules` chapter shows how functions are defined
- and organized in modules.
-* The :doc:`printing_and_parsing` chapter covers TVMScript definitions,
- program generation, and conversion between source text and IR.
+- :doc:`types_and_expressions` explains how values, computations, and tensor
+ shapes are represented.
+- :doc:`statements` covers how a body binds variables, writes tensors,
+ branches, and loops.
+- :doc:`operators_and_calls` introduces registered operators and function
+ calls.
+- :doc:`functions_and_modules` shows how functions are defined and organized
+ in modules.
+- :doc:`printing_and_parsing` covers utilities for parsing and printing
+ TVMScript.
diff --git a/docs/ir/printing_and_parsing.rst b/docs/ir/printing_and_parsing.rst
index 7fb4ce2ab..0713f5ecd 100644
--- a/docs/ir/printing_and_parsing.rst
+++ b/docs/ir/printing_and_parsing.rst
@@ -261,8 +261,50 @@ into a constant. The addition in ``add_two_variant`` is
therefore
``A[i] + T.float32(2.0)``; both functions still take only ``A`` and ``B``
as parameters.
+.. _user-ir-source-locations:
+
Source Locations
-~~~~~~~~~~~~~~~~
+----------------
+
+IR nodes can carry source locations to help diagnostics point back to
+the original code. The ``loc`` field stores this information, and
+:class:`ir.SourceLoc <tvm_tirx.ir.SourceLoc>` records the source file
+and range.
+
+The parameter ``A`` is declared on this line of the CPU program:
+
+.. literalinclude:: ../../examples/quick_start/add_one_cpu.py
+ :language: python
+ :start-at: def add_one(
+ :end-at: def add_one(
+ :lineno-match:
+
+The parameter's ``loc`` points back to this declaration. We can read
+``start_line`` to report its line number:
+
+.. code-block:: python
+
+ parameter = add_one.params[0]
+ location = parameter.loc
+ line = location.start_line
+ print(f"Parameter {parameter.name} is declared on line {line}.")
+
+The output is:
+
+.. code-block:: text
+
+ Parameter A is declared on line 29.
+
+``source_name.name`` also records the path to the CPU program's source
+file:
+
+.. code-block:: text
+
+ parameter: ir.Var A
+ └── loc: ir.SourceLoc
+ ├── source_name: ir.SourceName
+ │ └── name path to the source file
+ └── start_line line of A's declaration
The ``filename`` argument labels generated source in parser diagnostics.
The following example removes the required type annotation from ``A``
@@ -287,9 +329,8 @@ The output is:
Parameter 'A' requires an annotation (generated_add_constant.py, line 3)
The program comes from ``invalid_source``; ``filename`` supplies the
-name shown in the message. Successfully parsed
-nodes also record their positions through the
-:ref:`source-location fields <user-ir-source-locations>` introduced earlier.
+name shown in the message, and successfully parsed nodes record it in
+their ``loc`` fields.
.. _tvmscript-region-syntax:
diff --git a/docs/ir/types_and_expressions.rst
b/docs/ir/types_and_expressions.rst
index 44be5f866..687186457 100644
--- a/docs/ir/types_and_expressions.rst
+++ b/docs/ir/types_and_expressions.rst
@@ -434,48 +434,3 @@ call is therefore also a primitive expression. The Python
helper
The :doc:`operators_and_calls` chapter also covers calls to IR functions
and explains how registered operators provide result-type inference
and argument validation.
-
-.. _user-ir-source-locations:
-
-Source Locations
-----------------
-
-IR nodes can carry source locations to help diagnostics point back to
-the original code. The ``loc`` field stores this information, and
-:class:`ir.SourceLoc <tvm_tirx.ir.SourceLoc>` records the source file
-and range.
-
-The parameter ``A`` is declared on this line of the CPU program:
-
-.. literalinclude:: ../../examples/quick_start/add_one_cpu.py
- :language: python
- :start-at: def add_one(
- :end-at: def add_one(
- :lineno-match:
-
-The parameter's ``loc`` points back to this declaration. We can read
-``start_line`` to report its line number:
-
-.. code-block:: python
-
- parameter = add_one.params[0]
- location = parameter.loc
- line = location.start_line
- print(f"Parameter {parameter.name} is declared on line {line} of the CPU
program.")
-
-The output is:
-
-.. code-block:: text
-
- Parameter A is declared on line 29 of the CPU program.
-
-``source_name.name`` also records the path to the CPU program's source
-file:
-
-.. code-block:: text
-
- parameter: ir.Var A
- └── loc: ir.SourceLoc
- ├── source_name: ir.SourceName
- │ └── name path to the CPU program's source file
- └── start_line line of A's declaration
