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 37ac4abdf [DOCS] Add an IR walkthrough with runnable examples (#24)
37ac4abdf is described below
commit 37ac4abdfa4f3f37fb85883e8bfe498947017397
Author: Shushi Hong <[email protected]>
AuthorDate: Sat Oct 24 18:29:59 2026 -0400
[DOCS] Add an IR walkthrough with runnable examples (#24)
This pr adds an IR Walkthrough covering the shared IR object hierarchy,
types and expressions, statements, operators and calls, functions and
modules, and TVMScript printing and parsing.
Include minimal examples explaining how the IR objects relate and how to
use them, with downloadable Python scripts and links to the API
reference.
---------
Signed-off-by: Yuchuan <[email protected]>
Signed-off-by: tlopex <[email protected]>
Co-authored-by: Yuchuan <[email protected]>
---
docs/index.rst | 12 ++
docs/ir/examples/operators_and_calls.py | 56 ++++++
docs/ir/examples/printing_and_parsing.py | 132 +++++++++++++
docs/ir/examples/statements_and_functions.py | 78 ++++++++
docs/ir/examples/types_and_expressions.py | 100 ++++++++++
docs/ir/functions_and_modules.rst | 134 +++++++++++++
docs/ir/operators_and_calls.rst | 195 +++++++++++++++++++
docs/ir/overview.rst | 166 ++++++++++++++++
docs/ir/printing_and_parsing.rst | 139 ++++++++++++++
docs/ir/statements.rst | 190 +++++++++++++++++++
docs/ir/types_and_expressions.rst | 274 +++++++++++++++++++++++++++
python/tvm_tirx/ir/expr.py | 26 ++-
python/tvm_tirx/ir/module.py | 19 +-
python/tvm_tirx/ir/stmt.py | 11 +-
python/tvm_tirx/tirx/function.py | 7 +-
15 files changed, 1528 insertions(+), 11 deletions(-)
diff --git a/docs/index.rst b/docs/index.rst
index cf288fe58..5f5d9967d 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -21,6 +21,18 @@ tvm-tirx
Table of Contents
-----------------
+.. toctree::
+ :maxdepth: 1
+ :caption: IR Walkthrough
+ :titlesonly:
+
+ ir/overview
+ ir/types_and_expressions
+ ir/statements
+ ir/operators_and_calls
+ ir/functions_and_modules
+ ir/printing_and_parsing
+
.. toctree::
:maxdepth: 1
:caption: Reference
diff --git a/docs/ir/examples/operators_and_calls.py
b/docs/ir/examples/operators_and_calls.py
new file mode 100644
index 000000000..385c4ce53
--- /dev/null
+++ b/docs/ir/examples/operators_and_calls.py
@@ -0,0 +1,56 @@
+# 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.
+"""Inspect an operator, construct calls, and check result types."""
+
+from tvm_tirx import ir
+
+# begin operator lookup
+op = ir.Op.get("prim.exp")
+argument_names = [entry.name for entry in op.args_info]
+effect = op.get_attr("TCallEffectKind")
+# end operator lookup
+
+# begin call contract
+value = ir.const(1.0, "float32")
+call = ir.Call(op, [value])
+# end call contract
+
+# begin function call type
+callee = ir.Var("f", ir.FuncType([value.ty], value.ty))
+function_call = ir.Call(callee, [value])
+# end function call type
+
+# begin validation failure
+invalid = ir.Call(op, [], ty="float32")
+# end validation failure
+try:
+ invalid.validate()
+except TypeError as error:
+ assert "expected 1 argument, got 0" in str(error)
+else:
+ raise AssertionError("Expected an argument-count error")
+
+# begin provisional call
+provisional = ir.Call(op, [value], ty=ir.Type.missing())
+inferred = ir.reinfer_type(provisional)
+# end provisional call
+
+assert argument_names == ["x"] and effect == ir.CallEffectKind.Pure
+assert call.ty == function_call.ty == inferred == value.ty
+assert call.validate() is None
+assert isinstance(provisional.ty, ir.MissingType)
+print("Operator and call examples passed.")
diff --git a/docs/ir/examples/printing_and_parsing.py
b/docs/ir/examples/printing_and_parsing.py
new file mode 100644
index 000000000..1c5ed9738
--- /dev/null
+++ b/docs/ir/examples/printing_and_parsing.py
@@ -0,0 +1,132 @@
+# 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.
+"""Generate tensor programs, inspect their IR, and check source round trips."""
+
+# TVMScript resolves dependent parameter annotations when constructing IR.
+
+# begin tvmscript example
+from __future__ import annotations
+
+from tvm_tirx import ir, tirx
+from tvm_tirx.script import from_source
+from tvm_tirx.script import ir as I
+from tvm_tirx.script import tirx as T
+
+
[email protected]
+def f(n: T.int64, A: T.Tensor((2 * n, 128), "float32")) -> T.void:
+ for i, j in T.grid(2 * n, T.int64(128)):
+ A[i, j] = A[i, j] + T.float32(1)
+
+
+# end tvmscript example
+
+
+# begin script helpers
[email protected]
+def increment(value):
+ return value + ir.const(1, value.ty)
+
+
[email protected]
+def increment_at(A, i, j):
+ A[i, j] = increment(A[i, j])
+
+
+# end script helpers
+
+
+# begin program factory
+def make_increment(width, dtype="float32"):
+ if not isinstance(width, int) or width <= 0:
+ raise ValueError("width must be a positive Python integer")
+
+ @T.function
+ def generated(n: T.int64, A: T.Tensor((2 * n, width), dtype)) -> T.void:
+ for i, j in T.grid(2 * n, T.int64(width)):
+ increment_at(A, i, j)
+
+ return generated
+
+
+wide = make_increment(128)
+narrow = make_increment(64, "float16")
+# end program factory
+
+# begin print generated function
+source = wide.script()
+print(source)
+# end print generated function
+
+# begin diagnostic display
+wide.show(
+ print_line_numbers=True,
+ obj_to_underline=[wide.body[0]],
+)
+# end diagnostic display
+
+# begin printing and parsing
+restored = from_source(source, extra_vars={"T": T})
+ir.assert_structural_equal(wide, restored)
+# end printing and parsing
+
+
+# begin module round trip
[email protected]_module
+class Kernels:
+ wide = wide.with_attr("global_symbol", "wide")
+ narrow = narrow.with_attr("global_symbol", "narrow")
+
+
+module_source = Kernels.script()
+restored_module = from_source(module_source, extra_vars={"I": I, "T": T})
+ir.assert_structural_equal(Kernels, restored_module)
+# end module round trip
+
+# begin standalone source
+standalone_source = Kernels.script(extra_config={"ir.comment_imports": False})
+standalone_module = from_source(standalone_source)
+ir.assert_structural_equal(Kernels, standalone_module)
+# end standalone source
+
+assert isinstance(f, tirx.Function)
+ir.assert_structural_equal(f.with_attr("global_symbol", ""),
wide.with_attr("global_symbol", ""))
+for program, width, dtype in ((wide, 128, "float32"), (narrow, 64, "float16")):
+ n, tensor = program.params
+ assert tensor.ty.shape[0].b.same_as(n)
+ assert program.body[0].end.b.same_as(n)
+ assert int(tensor.ty.shape[1]) == width
+ assert tensor.ty.dtype == ir.PrimType(dtype)
+ assert tirx.analysis.check_well_formed(program)
+assert tirx.analysis.check_well_formed(Kernels)
+assert "increment_at(" not in source and "increment(" not in source
+assert wide.params[1].ty.script().startswith("T.Tensor(")
+assert "for " in wide.body[0].script()
+assert source == wide.script()
+try:
+ from_source(source)
+except SyntaxError:
+ pass
+else:
+ raise AssertionError("Source text must supply its namespace through
imports or extra_vars")
+try:
+ make_increment(0)
+except ValueError:
+ pass
+else:
+ raise AssertionError("The factory should reject a zero width")
+print("Printing and parsing example passed.")
diff --git a/docs/ir/examples/statements_and_functions.py
b/docs/ir/examples/statements_and_functions.py
new file mode 100644
index 000000000..1126b04aa
--- /dev/null
+++ b/docs/ir/examples/statements_and_functions.py
@@ -0,0 +1,78 @@
+# 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.
+"""Assemble an integer increment function from expressions and statements."""
+
+from tvm_tirx import ir, tirx
+
+# begin statements
+x = ir.Var("x", "int32")
+added = x + 1
+result = ir.Var("result", added.ty)
+body = ir.SeqStmt(
+ [
+ ir.Bind(result, added),
+ ir.Return(result),
+ ]
+)
+# end statements
+
+# begin branch
+branch = ir.If(x > 0, ir.Return(added), ir.Return(x))
+# end branch
+
+# begin loop
+i = ir.Var("i", "int32")
+loop = ir.For(i, 0, 4, 1, ir.ForKind.DEFAULT, ir.Evaluate(i))
+# end loop
+
+# begin function construction
+value = ir.Var("value", "int32")
+function = tirx.Function([value], ir.Return(value + 1), ret_type=value.ty)
+# end function construction
+
+# begin module construction
+function = function.with_attr("global_symbol", "main")
+entry = ir.GlobalVar("main")
+module = ir.IRModule({entry: function})
+# end module construction
+
+# begin global call
+global_call = ir.Call(entry, [ir.const(3, "int32")], ty=function.ret_type)
+# end global call
+
+# begin module serialization
+restored = ir.load_json(ir.save_json(module))
+ir.assert_structural_equal(module, restored)
+# end module serialization
+
+assert global_call.ty == x.ty
+assert module[entry].same_as(module["main"])
+assert tirx.analysis.check_well_formed(module)
+assert tirx.analysis.check_well_formed(tirx.Function([x], body, ret_type=x.ty))
+assert tirx.analysis.check_well_formed(tirx.Function([x], branch,
ret_type=x.ty))
+# begin module update
+variant = module.clone()
+variant["main"] = function.with_attr("example.note", "clone")
+assert "example.note" not in module["main"].attrs
+# end module update
+
+escaped = tirx.Function([], [loop, ir.Evaluate(i)])
+assert not tirx.analysis.check_well_formed(escaped, assert_mode=False)
+executable = tirx.compile(module, target="llvm")
+for value in (-2, 0, 3):
+ assert executable(value) == value + 1
+print("Statement, function, and execution examples passed.")
diff --git a/docs/ir/examples/types_and_expressions.py
b/docs/ir/examples/types_and_expressions.py
new file mode 100644
index 000000000..4f8d7b766
--- /dev/null
+++ b/docs/ir/examples/types_and_expressions.py
@@ -0,0 +1,100 @@
+# 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.
+"""Build and inspect typed scalar and tensor expressions.
+
+Run ``python types_and_expressions.py`` with TIRx installed.
+
+The four steps show scalar expressions, symbolic tensor shapes, operator
+calls, and a custom type. Each step prints the constructed IR and its type.
+For example, the tensor step prints::
+
+ 2. Tensor expressions
+ Fixed tensor type: T.Tensor((8, 128), T.float32)
+ Symbolic tensor type: T.Tensor((T.int64(2) * n, 128), T.float32)
+ A[0, 0] : T.float32
+
+The printed forms use TVMScript notation, such as T.int32 and T.Tensor(...).
+"""
+
+from tvm_ffi.dataclasses import py_class
+
+from tvm_tirx import ir, tirx
+
+# 1. Describe an integer computation, then construct its input and addition.
+# begin composite type
+integer_type = ir.PrimType("int32")
+signature = ir.FuncType([integer_type], integer_type)
+# end composite type
+
+# begin scalar expressions
+x = ir.Var("x", integer_type)
+added = x + ir.const(1, integer_type)
+# end scalar expressions
+
+print("1. Scalar expressions")
+print(f" {added.script()} : {added.ty.script()}")
+print(f" Function signature: {signature.script()}")
+
+# 2. A tensor type can use fixed dimensions or reuse integer expressions.
+# begin constant tensor type
+tile_type = tirx.TensorType((8, 128), "float32")
+# end constant tensor type
+
+# begin tensor expressions
+n = ir.Var("n", "int64")
+tensor_type = tirx.TensorType((2 * n, 128), "float32")
+tensor = ir.Var("A", tensor_type)
+element = tensor[0, 0]
+# end tensor expressions
+
+# The shape retains the same n object; no numeric value for n is needed here.
+assert tensor.ty.shape[0].b.same_as(n)
+print("\n2. Tensor expressions")
+print(f" Fixed tensor type: {tile_type.script()}")
+print(f" Symbolic tensor type: {tensor_type.script()}")
+print(f" {element.script()} : {element.ty.script()}")
+
+# 3. Retrieve one operator and use it in calls with different input
expressions.
+# begin operator expressions
+exp_op = ir.Op.get("prim.exp")
+value = ir.Var("value", "float32")
+exp_call = ir.Call(exp_op, [value])
+# end operator expressions
+
+tensor_exp = ir.Call(exp_op, [element])
+assert tensor_exp.op.same_as(exp_call.op)
+print("\n3. Operator calls")
+print(f" Registered operator: {exp_op.name}")
+print(f" {exp_call.script()} : {exp_call.ty.script()}")
+print(f" {tensor_exp.script()} : {tensor_exp.ty.script()}")
+
+
+# 4. Define a type for a pair of values using existing Type objects as fields.
+# begin custom type
+@py_class("my_ext.PairType")
+class PairType(ir.Type):
+ first: ir.Type
+ second: ir.Type
+
+
+pair = ir.Var("p", PairType(ir.PrimType("int32"), ir.PrimType("float32")))
+# end custom type
+
+print("\n4. A custom type")
+print(f" {pair.name} : {type(pair.ty).__name__}")
+print(f" Component types: {pair.ty.first.script()},
{pair.ty.second.script()}")
+print(f" Inherited source location: {type(pair.ty.loc).__name__}")
diff --git a/docs/ir/functions_and_modules.rst
b/docs/ir/functions_and_modules.rst
new file mode 100644
index 000000000..36d6d5dc1
--- /dev/null
+++ b/docs/ir/functions_and_modules.rst
@@ -0,0 +1,134 @@
+.. 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.
+
+Functions and Modules
+=====================
+
+Functions define computations through parameters, result types, and statement
+bodies, while modules collect named function definitions. Extensions can build
+programs with custom types and operators using these same structures, and
+passes can inspect or replace definitions through the existing interfaces.
+
+Functions
+---------
+
+``ir.BaseFunc`` is the shared base class for function definitions and is itself
+an expression. It inherits ``ty`` and ``loc`` from ``ir.Expr`` and adds a field
+shared by function definitions:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 15 25 60
+
+ * - Field
+ - Value type
+ - Meaning
+ * - ``attrs``
+ - ``ir.DictAttrs``
+ - Named function properties.
+
+``tirx.Function`` adds these fields:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 25 75
+
+ * - Field
+ - Meaning
+ * - ``params``
+ - Parameter variables in calling order.
+ * - ``ret_type``
+ - The declared result type; an omitted type becomes an empty
``TupleType``.
+ * - ``body``
+ - A ``SeqStmt``, or None for a declaration without an implementation.
+
+For example, this function takes an int32 input and returns that value plus
+one:
+
+.. literalinclude:: examples/statements_and_functions.py
+ :language: python
+ :start-after: # begin function construction
+ :end-before: # end function construction
+ :prepend: from tvm_tirx import ir, tirx
+
+``value`` is defined as a parameter and used by the expression in the return
+statement. The function's ``ty`` is ``ir.FuncType([value.ty], value.ty)``:
+construction derives this signature from the parameter types and ``ret_type``.
+The return type is declared explicitly; it is not inferred from the body.
+Passing None for ``body`` creates a declaration with a signature but no
+implementation.
+
+Modules
+-------
+
+``ir.IRModule`` maps ``ir.GlobalVar`` references to ``ir.BaseFunc``
definitions.
+Calls can use these references to identify a function, allowing a pass to
+replace its definition while retaining the reference. Place the increment
+function in a module under the name ``main``:
+
+.. literalinclude:: examples/statements_and_functions.py
+ :language: python
+ :start-after: # begin module construction
+ :end-before: # end module construction
+
+``entry`` is the reference; ``function`` contains the parameters and body.
+``module[entry]`` and ``module["main"]`` retrieve that same definition.
+``global_symbol`` sets the exported function name used by code generation.
+
+Use the stored reference to construct a call representing ``main(3)``:
+
+.. literalinclude:: examples/statements_and_functions.py
+ :language: python
+ :start-after: # begin global call
+ :end-before: # end global call
+
+This creates a call expression with an int32 result type. The signature lives
+on ``function.ty``; insertion does not copy it to ``entry.ty``, so the call
+supplies its result type explicitly. Reuse ``entry`` or retrieve it with
+``module.get_global_var("main")`` when referring to this definition.
+
+To update a definition while keeping the original module, clone the module
+and replace the definition in the clone. This example adds an attribute to
+the cloned module's ``main`` function:
+
+.. literalinclude:: examples/statements_and_functions.py
+ :language: python
+ :start-after: # begin module update
+ :end-before: # end module update
+
+``with_attr`` returns a new function, and assignment updates ``variant``'s
+function map. The original ``module["main"]`` keeps its attributes. Use
+:meth:`~tvm_tirx.tirx.Function.with_body` to replace a function body, and see
+:class:`~tvm_tirx.ir.IRModule` for lookup, copying, and update APIs.
+
+A module also carries metadata: ``attrs`` stores named attributes, and
+``global_infos`` stores named arrays of ``ir.GlobalInfo``. Save a module as
JSON
+to reconstruct its IR later:
+
+.. literalinclude:: examples/statements_and_functions.py
+ :language: python
+ :start-after: # begin module serialization
+ :end-before: # end module serialization
+
+Structural comparison checks the representation, including corresponding bound
+variables. Loading requires the graph's object types to be registered in the
+receiving process.
+
+See :class:`~tvm_tirx.tirx.Function` and :class:`~tvm_tirx.ir.IRModule` for the
+complete APIs. The :download:`statements_and_functions.py
<examples/statements_and_functions.py>`
+example checks scopes and updates, compiles with LLVM, and verifies that an
+input of three returns four.
diff --git a/docs/ir/operators_and_calls.rst b/docs/ir/operators_and_calls.rst
new file mode 100644
index 000000000..db33168b8
--- /dev/null
+++ b/docs/ir/operators_and_calls.rst
@@ -0,0 +1,195 @@
+.. 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.
+
+Operators and Calls
+===================
+
+Operators describe computations, while calls apply them to particular inputs.
+Custom operators can be registered alongside built-in ones, with their own
+argument requirements, type rules, and effects. Their calls use the existing
+IR representation, so extensions can build analyses, transformations, and
+lowering passes around the same objects and interfaces.
+
+Operators
+---------
+
+``ir.Op`` is an expression that identifies a registered operator. The registry
+stores shared operator definitions, which can be retrieved by name with
+``ir.Op.get(name)``. An operator's ``args_info`` describes its required value
+arguments. A validator is a registered function that checks whether a call's
+arguments meet the operator's requirements.
+
+Operator attributes describe properties shared by all calls to that operator.
+The ``TCallEffectKind`` attribute records whether a call only computes a value
+or can read or change program state.
+
+``prim.exp`` computes the exponential of one input. It declares an argument
+named ``x`` and has the effect kind ``Pure``: it computes a value without
+reading or changing external state. Inspect these properties on the registered
+object:
+
+.. literalinclude:: examples/operators_and_calls.py
+ :language: python
+ :start-after: # begin operator lookup
+ :end-before: # end operator lookup
+ :prepend: from tvm_tirx import ir
+
+Here, ``argument_names`` is ``["x"]``, and ``effect`` equals
+``ir.CallEffectKind.Pure``. See :class:`~tvm_tirx.ir.Op` for the registration
+and attribute APIs.
+
+Calls
+-----
+
+``ir.Call`` represents an invocation of an operator or function. The expression
+being called is its *callee*. A call inherits ``ty`` and ``loc`` from
+``ir.Expr`` and adds these fields:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 25 75
+
+ * - Field
+ - Meaning
+ * - ``op``
+ - The callee: a registered operator or another function-valued expression.
+ * - ``args``
+ - Value arguments in calling order.
+ * - ``ty_args``
+ - Explicit type arguments, empty by default.
+ * - ``attrs``
+ - Attributes for this invocation, or None. A dictionary becomes
+ ``ir.DictAttrs``.
+
+``ty_args`` supplies types as inputs to the operation; ``ty`` describes its
+result. Call attributes belong to one invocation, while the operator's
+registration supplies the rules used to interpret them.
+
+For example, this call applies the registered exponential operator to a float32
+constant:
+
+.. literalinclude:: examples/operators_and_calls.py
+ :language: python
+ :start-after: # begin call contract
+ :end-before: # end call contract
+
+This constructs the IR for ``exp(1.0)``; its result type is
+``ir.PrimType("float32")``. The computation is performed when a program
+containing the call is compiled and executed.
+
+A function-valued variable can also be the callee. Here, ``f`` has a signature
+that takes one float32 value and returns a float32 value:
+
+.. literalinclude:: examples/operators_and_calls.py
+ :language: python
+ :start-after: # begin function call type
+ :end-before: # end function call type
+
+``function_call`` represents ``f(1.0)``. The variable supplies a signature;
+the example does not provide an implementation of ``f``.
+
+.. _types-and-validation:
+
+Type Inference and Validation
+-----------------------------
+
+To avoid repeating type information at each call, ``ir.Call`` can infer its
+result type from the operator's registered rule or the callee's function
+signature. Omitting ``ty`` asks ``ir.Call`` to infer it; passing an explicit
+type stores that type instead. ``MissingType`` marks a result whose type has
+not yet been determined.
+
+Validation checks whether the arguments meet the operator's requirements.
+It is separate from construction, so tools can assemble a call before
+checking its arguments with ``call.validate()``. Type inference can also
+reject inputs while computing the result type.
+
+For ``prim.exp``, an explicit result type lets the constructor record a call
+with no arguments. Validation then rejects the argument count:
+
+.. literalinclude:: examples/operators_and_calls.py
+ :language: python
+ :start-after: # begin validation failure
+ :end-before: # end validation failure
+ :append: invalid.validate() # raises TypeError: expected one argument
+
+When constructing or rewriting IR, a call can temporarily carry a missing
+result type. ``ir.reinfer_type`` derives the type from its current inputs:
+
+.. literalinclude:: examples/operators_and_calls.py
+ :language: python
+ :start-after: # begin provisional call
+ :end-before: # end provisional call
+
+Here, ``inferred`` is ``ir.PrimType("float32")``. The function returns a type;
+it leaves ``provisional.ty`` unchanged. See :class:`~tvm_tirx.ir.Call` and
+:func:`~tvm_tirx.ir.reinfer_type` for the inference rules, and
+:meth:`~tvm_tirx.ir.Call.validate` for the validation contract.
+
+Register New Operators
+----------------------
+
+Registering an operator gives a custom computation an argument and result-type
+contract that calls and passes can share. New operators can be added from a
+Rust extension without modifying the C++ core.
+
+Use a ``register_ops!`` block with one ``OpDef`` chain per operator. Each chain
+supplies the operator's name, argument requirements, type-inference rule, and
+attributes.
+The operator joins the same registry as the built-in operators, so Python,
+C++, and Rust code can retrieve it.
+
+For example, this block registers the interface of ``my_ext.clamp``, an
+operation for clamping an integer index to ``[0, bound)``. A call has the type
+of its index, and ``validate()`` rejects arguments that are not integer
+expressions:
+
+.. code-block:: rust
+
+ use tvm_tirx::ir::op_attr::CALL_EFFECT_KIND;
+ use tvm_tirx::ir::sig::{self, IntExpr};
+ use tvm_tirx::ir::{CallEffectKind, CallObj, OpDef, Type};
+ use tvm_tirx::tvm_ffi::Result;
+
+ fn index_type(call: &CallObj) -> Result<Type> {
+ Ok(call.args.get(0)?.ty.clone())
+ }
+
+ tvm_tirx::register_ops! {
+ OpDef::new("my_ext.clamp", "Clamp an index to the range `[0, bound)`.")
+ .signature((
+ sig::arg::<IntExpr>("index", "The index."),
+ sig::arg::<IntExpr>("bound", "The exclusive upper bound."),
+ ))
+ .set_infer_type(index_type)
+ .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
+ }
+
+Here, ``.signature`` generates a validator from the ``IntExpr`` argument
+constraints. ``.set_infer_type`` registers ``index_type``, which returns the
+first argument's type. ``.set_attr`` records that the operation is pure.
+
+Once the extension is loaded and the operator is registered, Python can
+retrieve it with ``ir.Op.get("my_ext.clamp")`` and construct calls in the same
+way as for a built-in operator.
+
+To make these calls executable, supply a lowering that replaces them with
+existing IR or implement the operation in a backend.
+
+See :doc:`/reference/python/ir` for the complete operator and call APIs. The
+:download:`operators_and_calls.py <examples/operators_and_calls.py>` example
includes runnable
+calls, re-inference, and validation failures.
diff --git a/docs/ir/overview.rst b/docs/ir/overview.rst
new file mode 100644
index 000000000..db392288b
--- /dev/null
+++ b/docs/ir/overview.rst
@@ -0,0 +1,166 @@
+.. 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.
+
+IR Overview
+===========
+
+An intermediate representation (IR) describes a program in a form that compiler
+components can inspect and transform. It makes values, computations, control
+flow, and types explicit, allowing analyses and transformations to work with
the
+program's structure independently of the syntax used to construct it.
+
+TIRx's C++, Python, and Rust interfaces share a common IR object system through
+TVM FFI, allowing compiler components written in different languages to work
+together. The IR can be extended with custom operators and types, and its
+existing interfaces can be used to build custom analyses, transformation
passes,
+and backends.
+
+The following tree shows the inheritance relationships between IR object types.
+Registered names can differ from Python API names. For example, ``prim.Add``
+is available as ``ir.prim.Add`` in Python.
+
+.. 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
+ │ ├── sym.IterMapExpr Iterator mapping
expressions
+ │ │ ├── sym.IterSplitExpr Iterator splitting
+ │ │ └── sym.IterSumExpr Sums of iterator
expressions
+ │ └── sym.CanonicalExpr Internal simplifier
expressions
+ │ ├── sym.SplitExpr Internal split
expressions
+ │ └── sym.SumExpr Internal sum expressions
+ ├── 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.IRModule Global function
definitions
+ ├── ir.Attrs Structured attributes
+ │ ├── ir.DictAttrs Named attribute maps
+ │ ├── tirx.CallFFIKernelAttr FFI kernel call
attributes
+ │ ├── tirx.tensor.ScopeAttrs Tensor instruction
scopes
+ │ ├── tirx.tensor.TrnAttrs Tensor instruction
attributes
+ │ ├── tirx.cuda.TensorMapEncodeTiledAttr CUDA tensor map
encoding attributes
+ │ ├── tirx.cuda.TCGen05InstrDescriptorAttrs CUDA instruction
descriptors
+ │ └── tirx.cuda.TCGen05InstrDescriptorBlockScaledAttrs Block-scaled CUDA
descriptors
+ ├── ir.Location Source locations
+ │ ├── ir.UnknownLoc Unavailable source
information
+ │ ├── ir.SourceLoc Source ranges
+ │ └── ir.CallSiteLoc Callee and caller
locations
+ ├── ir.SourceName Source identifiers
+ ├── ir.Range Begin and end bounds
+ ├── ir.ArgumentInfo Operator argument
descriptions
+ ├── ir.GlobalInfo Module-level metadata
+ ├── ir.EnvFunc Named global function
references
+ ├── ir.UniqueNameSupply Unique-name generation
+ ├── tirx.Layout Index-to-axis mappings
+ │ ├── tirx.TileLayout Tile layouts
+ │ └── tirx.ComposeLayout Composed layouts
+ ├── tirx.Iter Layout iterators
+ ├── tirx.IndexMap Index transformations
+ ├── sym.IterMark Iterator source and
extent
+ └── ffi.Enum Registered enumeration
objects
+ ├── tirx.Axis Layout axis identities
+ └── tirx.ThreadScope Execution scope
identities
+
+.. _choose-a-chapter:
+
+What to Read Next
+-----------------
+
+The following chapters explain how to work with each part of the IR:
+
+* To understand how values, computations, and tensor shapes are represented,
+ see :doc:`types_and_expressions`.
+* To build statement bodies with bindings, tensor writes, branches, and loops,
+ see :doc:`statements`.
+* To work with registered operators and function calls, see
+ :doc:`operators_and_calls`.
+* To define functions and organize them into modules, see
+ :doc:`functions_and_modules`.
+* To write or generate TVMScript and print or parse IR, see
+ :doc:`printing_and_parsing`.
diff --git a/docs/ir/printing_and_parsing.rst b/docs/ir/printing_and_parsing.rst
new file mode 100644
index 000000000..df2cb8566
--- /dev/null
+++ b/docs/ir/printing_and_parsing.rst
@@ -0,0 +1,139 @@
+.. 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.
+
+.. _ir-construction-and-printing:
+
+Printing and Parsing
+====================
+
+Printing renders IR as readable TVMScript, and parsing constructs IR from
+Python source. Compiler extensions can reuse these interfaces to generate
+programs, inspect the resulting IR, and reconstruct programs from source.
+Extensions can customize how source builds IR through builder namespaces
+and how IR is rendered through printer hooks.
+
+For example, this function increments every element of a float32 tensor with
+shape ``(2*n, 128)``. Save the code in a Python file so the decorator can
inspect
+its source:
+
+.. literalinclude:: examples/printing_and_parsing.py
+ :language: python
+ :start-after: # begin tvmscript example
+ :end-before: # end tvmscript example
+
+``@T.function`` parses the definition into a ``tirx.Function``. The future
+import postpones annotation evaluation until the parser reads the definition,
+so ``n`` in the tensor annotation can resolve to the first IR parameter. The
+tensor's shape and the loop bound refer to that same variable, keeping ``2*n``
+symbolic in both places. The loops and indexing construct IR; the tensor
+computation runs when the compiled function is called.
+
+Script helpers let generated functions reuse parts of a computation.
+An ``@T.inline`` helper can return an expression, while ``@T.macro`` inserts
+statements into the function being built. These helpers separate the element
+addition from the tensor write:
+
+.. literalinclude:: examples/printing_and_parsing.py
+ :language: python
+ :start-after: # begin script helpers
+ :end-before: # end script helpers
+
+``increment`` constructs the addition, and ``increment_at`` writes its result
+back to the tensor. ``ir.const(1, value.ty)`` gives the constant the element's
+type, so the helpers also work with other element types.
+
+An ordinary Python function can choose a shape and dtype, construct a
+``@T.function`` definition, and return the resulting ``tirx.Function``.
+Here, ``make_increment`` builds a function with the requested column count
+and element type:
+
+.. literalinclude:: examples/printing_and_parsing.py
+ :language: python
+ :start-after: # begin program factory
+ :end-before: # end program factory
+
+``width`` and ``dtype`` are Python values read while the factory runs. ``wide``
+has a fixed second dimension of 128 and float32 elements; ``narrow`` has a
second
+dimension of 64 and float16 elements. Their first dimension remains ``2*n``,
+where ``n`` is an IR parameter supplied when the compiled function is called.
+The width check and function construction run while ``make_increment``
+executes; the generated loops describe the later tensor computation.
+Generating IR with Python logic in this way is metaprogramming.
+
+Printing lets you inspect what this Python code generated:
+
+.. literalinclude:: examples/printing_and_parsing.py
+ :language: python
+ :start-after: # begin print generated function
+ :end-before: # end print generated function
+
+``script()`` returns source text. The printed ``wide`` function contains these
+loops, using the selected width and element type:
+
+.. code-block:: text
+
+ for i in range(T.int64(0), T.int64(2) * n):
+ for j in range(T.int64(0), T.int64(128)):
+ A[i, j] = A[i, j] + T.float32(1.0)
+
+The helper calls have expanded into an addition and a tensor write. The source
+also declares ``n`` with ``T.dynamic``, which creates a symbolic ``ir.Var`` for
+reuse in the signature and tensor shape. You can also print the result of a
+pass or smaller parts of a program:
+``wide.params[1].ty.script()`` renders its tensor type, and
+``wide.body[0].script()`` renders its outer loop.
+
+For interactive inspection, ``wide.show()`` displays highlighted source.
+See :meth:`~tvm_tirx.tirx.Function.show` for options such as line numbers and
+underlining IR objects.
+
+``from_source`` parses a source string into a new IR object. Names used in
+the text must be supplied by its imports or the ``extra_vars`` dictionary.
+The printed source refers to ``T``, but its import is a comment by default,
+so this example supplies ``T`` through ``extra_vars``:
+
+.. literalinclude:: examples/printing_and_parsing.py
+ :language: python
+ :start-after: # begin printing and parsing
+ :end-before: # end printing and parsing
+
+The structural comparison checks that the reconstructed program has the same
+operations, types, and variable bindings. Its printed syntax can differ from
+the original: for example, ``T.grid`` becomes nested loops.
+
+The same process works for modules. ``@I.ir_module`` turns a class containing
+function definitions into an ``ir.IRModule``. Here it collects both generated
+variants, which can then be printed and parsed together:
+
+.. literalinclude:: examples/printing_and_parsing.py
+ :language: python
+ :start-after: # begin module round trip
+ :end-before: # end module round trip
+
+To include import statements in the printed text, set ``ir.comment_imports``
+to False. The parser can then resolve ``I`` and ``T`` from those imports,
+without ``extra_vars``. This is useful when saving the source for another
+process:
+
+.. literalinclude:: examples/printing_and_parsing.py
+ :language: python
+ :start-after: # begin standalone source
+ :end-before: # end standalone source
+
+Download the :download:`complete example <examples/printing_and_parsing.py>`
+for runnable code. See the :doc:`parser </reference/python/script.parser>` and
+:doc:`printer </reference/python/script.printer>` API references for details.
diff --git a/docs/ir/statements.rst b/docs/ir/statements.rst
new file mode 100644
index 000000000..091c6c1d9
--- /dev/null
+++ b/docs/ir/statements.rst
@@ -0,0 +1,190 @@
+.. 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.
+
+.. _statements-and-control-flow:
+
+Statements
+==========
+
+Statements describe variable bindings, tensor writes, and control flow.
+Extensions can construct and rewrite program bodies with the existing
+statement nodes. Custom operators can use ``ir.RegionStmt`` to define how
+a group of statements executes.
+
+.. _statements-and-sequences:
+
+Statement Nodes
+---------------
+
+``ir.Stmt`` is the base class for all statement nodes, including sequences,
+branches, and loops. It defines the following field:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 15 25 60
+
+ * - Field
+ - Value type
+ - Meaning
+ * - ``loc``
+ - ``ir.Location``
+ - Source information for diagnostics. Defaults to ``UnknownLoc`` and does
+ not affect structural comparison.
+
+Subclasses inherit ``loc`` and add fields for operands or nested statements.
+The shared statement classes are:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 25 35 40
+
+ * - Class
+ - Represents
+ - Additional fields
+ * - ``ir.Bind``
+ - A variable definition
+ - ``var``, ``value``
+ * - ``ir.TensorStore``
+ - A tensor element write; the destination type must support stores
+ - ``dest``, ``indices``, ``value``
+ * - ``ir.Evaluate``
+ - Expression evaluation with its result discarded
+ - ``value``
+ * - ``ir.AssertStmt``
+ - A program assertion
+ - ``condition``, ``error_kind``, ``message_parts``
+ * - ``ir.Return``
+ - A function return
+ - ``value``
+ * - ``ir.Break``
+ - Exit the current loop
+ - None
+ * - ``ir.Continue``
+ - Start the next loop iteration
+ - None
+ * - ``ir.SeqStmt``
+ - A sequence of statements in execution order
+ - ``seq``
+ * - ``ir.If``
+ - Select a branch
+ - ``condition``, ``then_case``, ``else_case``
+ * - ``ir.For``
+ - Repeat a body over an integer range
+ - ``var``, ``begin``, ``end``, ``step``, ``kind``, ``body``,
``annotations``
+ * - ``ir.While``
+ - Repeat while a condition is true
+ - ``condition``, ``body``
+ * - ``ir.RegionStmt``
+ - An operator-defined statement with a nested body
+ - ``op``, ``args``, ``attrs``, ``body_params``, ``body``, ``result_vars``
+
+Statement Sequences
+-------------------
+
+A return that uses a bound variable must follow its definition.
+``ir.SeqStmt`` records the execution order of such statements.
+For example, consider a function that adds one to its input:
+
+.. code-block:: python
+
+ def increment(x):
+ result = x + 1
+ return result
+
+The following code constructs that function body as IR objects:
+
+.. literalinclude:: examples/statements_and_functions.py
+ :language: python
+ :start-after: # begin statements
+ :end-before: # end statements
+ :prepend: from tvm_tirx import ir, tirx
+
+The Python variable ``added`` holds the expression for ``x + 1``. ``ir.Var``
+creates the IR variable ``result``, and ``ir.Bind(result, added)`` defines it
+with that expression as its value. ``added.ty`` gives ``result`` the addition's
+``int32`` type. ``ir.Return(result)`` refers to the same variable, and
+``SeqStmt`` places the definition before the return.
+
+Nested sequences are flattened, and a sequence introduces no scope of its
+own. Variables defined by ``Bind`` are available to subsequent statements in
+the enclosing scope.
+
+Control Flow and Scope
+----------------------
+
+Control-flow statement nodes include ``ir.If`` for branching, ``ir.For`` and
+``ir.While`` for loops, ``ir.Break`` and ``ir.Continue`` for loop control, and
+``ir.Return`` for returning from a function.
+
+The statements inside a branch or loop form its body. Branch bodies are
+stored in ``If.then_case`` and ``If.else_case``, while loop bodies are stored
+in ``For.body`` and ``While.body``. Each body is a ``SeqStmt`` and can contain
+further branches and loops.
+
+For example, this branch returns the incremented value only for positive
+inputs. Both bodies are normalized to ``SeqStmt``:
+
+.. literalinclude:: examples/statements_and_functions.py
+ :language: python
+ :start-after: # begin branch
+ :end-before: # end branch
+
+``If`` and ``While`` use boolean conditions. ``If`` requires an explicit
+``else_case`` argument; pass None when there is no else branch.
+
+A loop instead repeats a body. This loop evaluates each index from zero through
+three:
+
+.. literalinclude:: examples/statements_and_functions.py
+ :language: python
+ :start-after: # begin loop
+ :end-before: # end loop
+
+``For`` takes an exclusive ``end`` and explicit ``step`` and ``kind``
arguments.
+Its variable, bounds, and step must have matching scalar integer types after
+conversion.
+
+The ``kind`` argument controls how loop iterations are executed.
+``ForKind.DEFAULT`` executes them in order, one after another.
+See :class:`~tvm_tirx.ir.ForKind` for the available execution kinds.
+
+``Break`` exits a loop, and ``Continue`` starts its next iteration.
+
+Variable scope determines where a definition can be used. The loop defines
``i``
+inside its body; it cannot be used after the loop. Branch-local definitions
+similarly stay inside their branch. Each variable object has one definition,
+and its uses refer to that object.
+
+``ir.RegionStmt`` lets an operator define how a nested body is executed without
+introducing a new statement class. Its ``op`` field refers to the registered
+``ir.Op`` that defines this behavior. Inputs in ``args`` and ``attrs``
+are evaluated in the enclosing scope. ``body_params`` defines variables
+visible within ``body``;
+``result_vars`` defines variables available to subsequent statements in the
+enclosing sequence. See :class:`~tvm_tirx.ir.RegionStmt` for the operator
+contract.
+
+To check these definition and scope rules, the body needs the context of a
+complete function or module. :doc:`functions_and_modules` supplies that
context;
+``tirx.analysis.check_well_formed`` then checks variable and tensor uses and
loop
+control. Invalid uses raise ``RuntimeError`` by default; ``assert_mode=False``
+returns False instead.
+
+The :download:`statements_and_functions.py
<examples/statements_and_functions.py>`
+example completes and runs the increment function. See
+:doc:`/reference/python/ir` for statement constructors and
+:doc:`/reference/python/tirx.analysis` for verification APIs.
diff --git a/docs/ir/types_and_expressions.rst
b/docs/ir/types_and_expressions.rst
new file mode 100644
index 000000000..4ab478ac0
--- /dev/null
+++ b/docs/ir/types_and_expressions.rst
@@ -0,0 +1,274 @@
+.. 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.
+
+.. _expressions-and-types:
+
+Types and Expressions
+=====================
+
+Expressions represent program values and computations, while types describe
+properties such as dtypes, tensor shapes, and function signatures. This shared
+representation can be extended with custom types and operators. Programs using
+these constructs can be built, inspected, and transformed through the existing
+IR interfaces.
+
+.. _user-ir-type:
+.. _user-ir-prim-type:
+
+Types
+-----
+
+``ir.Type`` is the base class for type objects, giving analyses and
+transformations a common interface across different kinds of values.
+It defines the following field:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 15 25 60
+
+ * - Field
+ - Value type
+ - Meaning
+ * - ``loc``
+ - ``ir.Location``
+ - Source information for diagnostics. Defaults to ``UnknownLoc`` and does
+ not affect structural comparison.
+
+Concrete subclasses inherit ``loc`` and add fields describing different kinds
+of values. Common examples include:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 30 35 35
+
+ * - Class
+ - Describes
+ - Additional fields
+ * - ``ir.PrimType``
+ - A scalar or vector dtype
+ - ``dtype``
+ * - ``ir.PtrType``
+ - A pointer's element type and memory scope
+ - ``element_type``, ``scope``
+ * - ``ir.TupleType``
+ - Ordered tuple field types
+ - ``fields``
+ * - ``ir.FuncType``
+ - Function argument and return types
+ - ``arg_types``, ``ret_type``
+ * - ``tirx.TensorType``
+ - A tensor's shape, element type, and storage properties
+ - ``shape``, ``dtype``, ``scope``, ``data_alignment``, ``layout``
+
+Some of these fields refer to other type objects. For example, ``FuncType``
+stores the types of a function's arguments and result:
+
+.. literalinclude:: examples/types_and_expressions.py
+ :language: python
+ :start-after: # begin composite type
+ :end-before: # end composite type
+ :prepend: from tvm_tirx import ir, tirx
+
+``signature`` describes a function that takes and returns int32 values.
+``integer_type.dtype`` holds the int32 dtype descriptor;
+``signature.arg_types[0]`` and ``signature.ret_type`` both refer to the
+``integer_type`` object. Tuple fields and pointer element types similarly
+refer to other type objects.
+
+.. _user-tensor-type:
+
+``tirx.TensorType`` combines an element type with shape and storage
information:
+
+.. literalinclude:: examples/types_and_expressions.py
+ :language: python
+ :start-after: # begin constant tensor type
+ :end-before: # end constant tensor type
+
+``tile_type`` describes an ``(8, 128)`` tensor. Its ``shape`` contains integer
+expressions for ``8`` and ``128``, and its ``dtype`` is
+``ir.PrimType("float32")``. Dimensions can also use variables and arithmetic.
+
+Use ``==`` to compare type structure and ``same_as`` to compare object
identity.
+
+.. _user-ir-var:
+
+Expressions
+-----------
+
+Expression nodes make the inputs to a computation explicit, so analyses can
+inspect its operands and transformations can replace parts of the computation.
+``ir.Expr`` is the base class for these nodes. It defines these fields:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 15 25 60
+
+ * - Field
+ - Value type
+ - Meaning
+ * - ``ty``
+ - ``ir.Type``
+ - Result type. ``ir.MissingType`` marks type information that has not yet
+ been supplied or computed.
+ * - ``loc``
+ - ``ir.Location``
+ - Source information for diagnostics. Defaults to ``UnknownLoc`` and does
+ not affect structural comparison.
+
+Subclasses inherit ``ty`` and ``loc`` and add fields for their stored values or
+operands. The following examples show common expression classes and selected
+fields:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 30 35 35
+
+ * - Classes (examples)
+ - Represents
+ - Additional fields
+ * - ``ir.Var``
+ - A variable
+ - ``name``
+ * - ``ir.IntImm``, ``ir.FloatImm``, ``ir.StringImm``
+ - Literal constants
+ - ``value``
+ * - ``ir.prim.Add``, ``ir.prim.Sub``, ``ir.prim.Mul``
+ - Binary arithmetic
+ - ``a``, ``b``
+ * - ``ir.prim.EQ``, ``ir.prim.LT``
+ - Comparisons
+ - ``a``, ``b``
+ * - ``ir.Op``
+ - A registered operator
+ - ``name``, ``doc``, ``args_info``
+ * - ``ir.Call``
+ - A call to an operator or function
+ - ``op``, ``args``, ``attrs``, ``ty_args``
+ * - ``ir.TensorLoad``
+ - A tensor element read
+ - ``source``, ``indices``
+
+Arithmetic and comparison nodes for primitive values are available under
+``ir.prim``. The binary nodes above store their left and right operand
+expressions in ``a`` and ``b``.
+
+Registered operators are themselves expressions of class ``ir.Op``. Each has
+a name, such as ``prim.exp``, and metadata describing its arguments.
+``ir.Op.get`` retrieves the operator object by name; ``ir.Call`` uses it as
+the callee of an invocation:
+
+.. literalinclude:: examples/types_and_expressions.py
+ :language: python
+ :start-after: # begin operator expressions
+ :end-before: # end operator expressions
+
+``exp_call.op`` refers to ``exp_op``, and ``exp_call.args`` contains the input
+``value``. The same operator object can be used by many calls, each with its
+own arguments. ``ir.prim.exp(value)`` is a helper that constructs this call
+form. :doc:`operators_and_calls` explains operator registration, attributes,
+type inference, and validation.
+
+Operands and call arguments are themselves expressions, allowing these forms
+to be combined into larger computations.
+
+.. _user-ir-const:
+
+To describe ``x + 1``, create a variable with the integer type above and a
+constant of that type. Adding them constructs an ``Add`` expression:
+
+.. literalinclude:: examples/types_and_expressions.py
+ :language: python
+ :start-after: # begin scalar expressions
+ :end-before: # end scalar expressions
+
+``x`` is a ``Var`` and ``added`` is an ``Add``. Both carry ``integer_type`` in
+their ``ty`` field. Reuse the same variable object when a symbol appears in
+several places: another ``ir.Var("x", integer_type)`` is a distinct symbol,
+even with the same name.
+
+A tensor variable follows the same pattern, with a ``TensorType`` in its
+``ty`` field. Its shape can reuse the integer expressions described above:
+
+.. literalinclude:: examples/types_and_expressions.py
+ :language: python
+ :start-after: # begin tensor expressions
+ :end-before: # end tensor expressions
+
+The first dimension stores the expression ``2*n``. The same ``n`` can appear
+in loop bounds and indices. The element load ``tensor[0, 0]`` is another
+expression, with type ``ir.PrimType("float32")``.
+
+The symbol ``n`` has a known int64 type even though its value is not yet known.
+Constructing these types and expressions does not allocate tensor storage.
+
+.. _user-ir-register-type:
+
+Register New Types
+------------------
+
+Custom types let extensions describe their own kinds of values and use them
+in existing IR variables and function signatures. New types can be added from
+Python or Rust without modifying the C++ core.
+Define a subclass of ``ir.Type`` and register it under its own key. Its
+reflected fields let structural comparison and printing handle it in the same
+way as built-in types.
+
+In Python, use ``tvm_ffi.dataclasses.py_class``; in Rust, use
+``tvm_ffi::rs_class`` and ``tvm_ffi::def_object_ref!``.
+
+For example, this type describes a pair of values by the types of its two
+components. In Python, the fields are annotated class attributes, and the
+constructor takes them in order:
+
+.. literalinclude:: examples/types_and_expressions.py
+ :language: python
+ :start-after: # begin custom type
+ :end-before: # end custom type
+ :prepend: from tvm_ffi.dataclasses import py_class
+ from tvm_tirx import ir
+
+In Rust, the object struct starts with the ``TypeObj`` it extends, and the
+generated ``new`` for ``PairTypeObj`` takes that parent object and then the
+fields:
+
+.. code-block:: rust
+
+ use tvm_tirx::ir::{Location, Type, TypeObj};
+ use tvm_tirx::tvm_ffi::{self, def_object_ref, rs_class};
+
+ #[rs_class("my_ext.PairType", final)]
+ pub struct PairTypeObj {
+ base: TypeObj,
+ pub first: Type,
+ pub second: Type,
+ }
+
+ def_object_ref!(pub PairType, Type, PairTypeObj);
+
+ impl PairType {
+ pub fn new(first: impl Into<Type>, second: impl Into<Type>) -> Self {
+ PairTypeObj::new(TypeObj::new(Location::unknown()), first,
second).into()
+ }
+ }
+
+Once registered, the type can be assigned to variables and used inside other
+types, just like a built-in type.
+
+See :class:`~tvm_tirx.tirx.TensorType` for the complete constructor and storage
+options, and :doc:`/reference/python/ir` for expression and type APIs.
+Download :download:`types_and_expressions.py
<examples/types_and_expressions.py>`
+for runnable examples from this chapter.
diff --git a/python/tvm_tirx/ir/expr.py b/python/tvm_tirx/ir/expr.py
index b7467b029..5dbb5ed74 100644
--- a/python/tvm_tirx/ir/expr.py
+++ b/python/tvm_tirx/ir/expr.py
@@ -150,6 +150,10 @@ class GlobalVar(Expr):
GlobalVar is used to refer to the global functions
stored in the IRModule.
+ Each construction creates a distinct reference, even when names match.
+ Use ``module.get_global_var(name)`` to retrieve the stored reference when
+ constructing calls or replacing definitions.
+
Parameters
----------
name_hint: str
@@ -703,6 +707,11 @@ class Call(_CallableExprWithOp):
contract explicitly after inputs are ready. Constructing a ``Call`` never
emits a builder statement, including calls to allocation operators.
+ For an Op callee, a registered fixed return type takes precedence over its
+ type-inference callback. For another callee with a FuncType, inference uses
+ the signature's return type. Without an applicable rule, the result is
+ MissingType. An explicit ``ty``, including MissingType, bypasses inference.
+
Parameters
----------
op : Expr or str
@@ -759,15 +768,26 @@ class Call(_CallableExprWithOp):
return ty, op, args, attrs, ty_args, loc
def validate(self) -> None:
- """Check the registered operator contract without changing this
Call."""
+ """Run the registered operator's validator without changing this Call.
+
+ An Op callee with a validator invokes that validator; other callees and
+ operators without validators perform no checks. Success returns None,
+ and failed checks raise a diagnostic exception. The checks depend on
the
+ registered validator and do not generally establish agreement between
+ the stored result type and the type implied by the inputs.
+ """
_ffi_api.CallValidate(self)
def reinfer_type(call: Call) -> "tvm_tirx.ir.Type":
"""Derive a Call's result type from its current inputs without changing
the Call.
- The operator must register a fixed return type or a context-free inference
- rule. This does not invoke the operator's validator.
+ For an Op callee, a registered ``TFixedReturnType`` takes precedence over
+ ``FInferType``. For another callee whose type is a FuncType, the result is
+ that function type's return type. Without an applicable rule, this returns
+ MissingType. Inference errors propagate.
+
+ This does not invoke the operator's validator or overwrite ``call.ty``.
"""
return _ffi_api.CallReinferType(call)
diff --git a/python/tvm_tirx/ir/module.py b/python/tvm_tirx/ir/module.py
index 1a4405d1e..5e221b6e4 100644
--- a/python/tvm_tirx/ir/module.py
+++ b/python/tvm_tirx/ir/module.py
@@ -69,6 +69,13 @@ class IRModule(Node, Scriptable):
self.__pyfuncs__ = {}
def clone(self) -> "IRModule":
+ """Create a module object with the same functions and metadata.
+
+ The clone initially shares its function objects and other contents with
+ the original. Adding, replacing, or removing function-map entries on
the
+ clone preserves the original module's entries. Contents are not
+ recursively copied.
+ """
return _ffi_api.IRModuleClone(self)
def functions_items(self):
@@ -130,12 +137,18 @@ class IRModule(Node, Scriptable):
return _ffi_api.IRModuleContains(self, var)
def update(self, other):
- """Insert functions in another Module to current one.
+ """Insert function definitions from another module or dictionary in
place.
+
+ Definitions are merged using their GlobalVar keys. A dictionary's
+ string keys create fresh GlobalVar objects; a fresh key with the same
+ name as an existing key raises RuntimeError. To replace an existing
+ definition, reuse its stored GlobalVar or use assignment by name.
Parameters
----------
- other: IRModule
- The module to merge into the current Module.
+ other : IRModule or dict[GlobalVar or str, BaseFunc]
+ Function definitions to merge into this module. Module attributes
+ and global_infos are not merged.
"""
if isinstance(other, dict):
other = IRModule(other)
diff --git a/python/tvm_tirx/ir/stmt.py b/python/tvm_tirx/ir/stmt.py
index cf94199c3..13c57b706 100644
--- a/python/tvm_tirx/ir/stmt.py
+++ b/python/tvm_tirx/ir/stmt.py
@@ -114,9 +114,14 @@ class AssertStmt(Stmt):
class ForKind(IntEnum):
"""Execution semantics of a :class:`For` loop.
- ``DEFAULT`` executes iterations sequentially; ``PARALLEL`` permits parallel
- execution. ``VECTORIZED`` requests a vectorized body and ``UNROLLED``
requests
- unrolling. Transformations must preserve the selected execution semantics.
+ * ``DEFAULT`` executes iterations in order, one after another.
+ * ``PARALLEL`` allows different iterations to run concurrently.
+ * ``VECTORIZED`` asks the compiler to combine iterations into vector
+ operations that process multiple elements at once.
+ * ``UNROLLED`` asks the compiler to expand the loop into repeated copies
+ of its body.
+
+ Transformations must preserve the selected execution semantics.
"""
DEFAULT = 0 # Ordinary loop with sequential iteration semantics.
diff --git a/python/tvm_tirx/tirx/function.py b/python/tvm_tirx/tirx/function.py
index b1da7e8ca..d3c6fae0c 100644
--- a/python/tvm_tirx/tirx/function.py
+++ b/python/tvm_tirx/tirx/function.py
@@ -43,7 +43,9 @@ class Function(BaseFunc, Scriptable):
body : Stmt or Sequence[Stmt] or None
The function body, or None for a declaration without an implementation.
ret_type : tvm_tirx.ir.Type, optional
- The return type; defaults to MissingType until supplied or inferred.
+ The declared return type. Omitting it or supplying MissingType produces
+ an empty TupleType, the type returned by the C++ VoidType() helper.
+ Construction does not infer a return type from the body.
attrs : tvm_tirx.ir.DictAttrs, optional
Function attributes, empty by default.
loc : Location, optional
@@ -90,7 +92,8 @@ class Function(BaseFunc, Scriptable):
The new body, or None for a declaration.
loc : Location, optional
- The source location of the function.
+ The source location of the new function. Defaults to UNKNOWN_LOC;
+ pass ``loc=self.loc`` to preserve the original location.
Returns
-------
