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 7c0d37abf [DOCS] Rework the IR Walkthrough around a shared TVMScript
program (#62)
7c0d37abf is described below
commit 7c0d37abf99670583750aba1dcdc19ed5ecc3242
Author: Shushi Hong <[email protected]>
AuthorDate: Sun Oct 25 19:21:28 2026 -0400
[DOCS] Rework the IR Walkthrough around a shared TVMScript program (#62)
This pr adds a runnable bias + sigmoid program to Get Started and use
its scalar, vector, and symbolic-shape computation throughout the IR
Walkthrough. Clarify the core IR concepts, consolidate the examples, and
remove custom registration sections and repetitive closing notes.
---------
Signed-off-by: tlopex <[email protected]>
Signed-off-by: Yuchuan <[email protected]>
Co-authored-by: Yuchuan <[email protected]>
---
docs/get_started/example.rst | 59 ++++
docs/get_started/examples/add_one_cpu.py | 40 +++
docs/get_started/examples/add_one_cuda.py | 43 +++
docs/index.rst | 7 +
docs/ir/examples/operators_and_calls.py | 56 ----
docs/ir/examples/printing_and_parsing.py | 132 ---------
docs/ir/examples/sigmoid.py | 217 +++++++++++++++
docs/ir/examples/statements_and_functions.py | 78 ------
docs/ir/examples/tvmscript_sigmoid.py | 87 ++++++
docs/ir/examples/types_and_expressions.py | 100 -------
docs/ir/functions_and_modules.rst | 129 ++++-----
docs/ir/operators_and_calls.rst | 183 +++++--------
docs/ir/overview.rst | 118 ++++----
docs/ir/printing_and_parsing.rst | 159 +++++------
docs/ir/statements.rst | 164 +++++-------
docs/ir/types_and_expressions.rst | 387 +++++++++++++++++----------
16 files changed, 1021 insertions(+), 938 deletions(-)
diff --git a/docs/get_started/example.rst b/docs/get_started/example.rst
new file mode 100644
index 000000000..241295607
--- /dev/null
+++ b/docs/get_started/example.rst
@@ -0,0 +1,59 @@
+.. 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.
+
+Your First TIRx Program
+=======================
+
+In this chapter, we'll write and run a small TIRx program using TVMScript,
+TIRx's Python-based notation for IR, then look at the IR it builds. Our example
+adds one to each element of a tensor: ``B[i] = A[i] + 1``.
+
+``A`` is the input tensor, and ``B`` stores the result. Both tensors have the
+same dynamic length ``n``. We can call the compiled program with different
+lengths without recompiling.
+
+With an LLVM-enabled TIRx build, you can download
+:download:`add_one_cpu.py <examples/add_one_cpu.py>` and run ``python
add_one_cpu.py``.
+Here's the complete program:
+
+.. literalinclude:: examples/add_one_cpu.py
+ :language: python
+ :start-after: # begin program
+ :end-before: # end program
+
+Here, ``@T.function`` turns the ``add_one`` definition into IR. The code
+in ``add_one`` loops over the input and writes each result to ``B``.
+``tirx.compile`` compiles it for the CPU, and the NumPy assertion checks
+its output. We'll explore the IR objects behind these constructs in the
+:doc:`IR Walkthrough </ir/overview>`.
+
+We can also run this computation on an NVIDIA GPU. In addition to LLVM,
+this example needs CUDA support, an NVIDIA driver, the CUDA toolkit, and
+``cuda-bindings``. You can download
+:download:`add_one_cuda.py <examples/add_one_cuda.py>` and run ``python
add_one_cuda.py``.
+Here's the complete GPU program:
+
+.. literalinclude:: examples/add_one_cuda.py
+ :language: python
+ :start-after: # begin program
+ :end-before: # end program
+
+Here, ``T.device_entry`` marks the GPU computation. Each thread handles one
+element, and the launch uses enough blocks of 256 threads to cover ``n``
+elements. The ``if i < n`` check skips any extra threads. ``tirx.compile``
+compiles it for CUDA, and ``cuda(0)`` selects the GPU for the input and output
+tensors. The NumPy assertion checks the result in the same way as the CPU
example.
diff --git a/docs/get_started/examples/add_one_cpu.py
b/docs/get_started/examples/add_one_cpu.py
new file mode 100644
index 000000000..0a455dfac
--- /dev/null
+++ b/docs/get_started/examples/add_one_cpu.py
@@ -0,0 +1,40 @@
+# 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.
+"""Add one to each element on the CPU. Requires LLVM support."""
+
+# begin program
+import numpy as np
+
+from tvm_tirx import runtime, tirx
+from tvm_tirx.script import tirx as T
+
+n = T.dynamic("n")
+
+
[email protected]
+def add_one(A: T.Tensor((n,), "float32"), B: T.Tensor((n,), "float32")):
+ for i in range(n):
+ B[i] = A[i] + 1
+
+
+compiled = tirx.compile(add_one, target="llvm")
+a = runtime.tensor(np.arange(10, dtype="float32"))
+b = runtime.empty(a.shape, a.dtype)
+compiled(a, b)
+np.testing.assert_allclose(b.numpy(), a.numpy() + 1)
+print(b.numpy())
+# end program
diff --git a/docs/get_started/examples/add_one_cuda.py
b/docs/get_started/examples/add_one_cuda.py
new file mode 100644
index 000000000..2db908940
--- /dev/null
+++ b/docs/get_started/examples/add_one_cuda.py
@@ -0,0 +1,43 @@
+# 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.
+"""Add one to each element on an NVIDIA GPU. Requires LLVM and CUDA support."""
+
+# begin program
+import numpy as np
+
+from tvm_tirx import cuda, runtime, tirx
+from tvm_tirx.script import tirx as T
+
+n = T.dynamic("n")
+
+
[email protected]
+def add_one(A: T.Tensor((n,), "float32"), B: T.Tensor((n,), "float32")):
+ T.device_entry(launch=T.cuda.LaunchConfig(grid=((n + 255) // 256,),
block=(256,)))
+ i = T.cuda.block_idx("x") * 256 + T.cuda.thread_idx("x")
+ if i < n:
+ B[i] = A[i] + 1
+
+
+compiled = tirx.compile(add_one, target="cuda")
+dev = cuda(0)
+a = runtime.tensor(np.arange(10, dtype="float32"), device=dev)
+b = runtime.empty(a.shape, a.dtype, device=dev)
+compiled(a, b)
+np.testing.assert_allclose(b.numpy(), a.numpy() + 1)
+print(b.numpy())
+# end program
diff --git a/docs/index.rst b/docs/index.rst
index 5f5d9967d..d9fe8b1ce 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -21,6 +21,13 @@ tvm-tirx
Table of Contents
-----------------
+.. toctree::
+ :maxdepth: 1
+ :caption: Get Started
+ :titlesonly:
+
+ get_started/example
+
.. toctree::
:maxdepth: 1
:caption: IR Walkthrough
diff --git a/docs/ir/examples/operators_and_calls.py
b/docs/ir/examples/operators_and_calls.py
deleted file mode 100644
index 385c4ce53..000000000
--- a/docs/ir/examples/operators_and_calls.py
+++ /dev/null
@@ -1,56 +0,0 @@
-# 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
deleted file mode 100644
index 1c5ed9738..000000000
--- a/docs/ir/examples/printing_and_parsing.py
+++ /dev/null
@@ -1,132 +0,0 @@
-# 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/sigmoid.py b/docs/ir/examples/sigmoid.py
new file mode 100644
index 000000000..cf1157446
--- /dev/null
+++ b/docs/ir/examples/sigmoid.py
@@ -0,0 +1,217 @@
+# 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.
+"""This example builds and executes model-specific bias + sigmoid programs.
+
+TIRx needs LLVM support. The example runs from the repository root with:
+ uv run --no-sync python docs/ir/examples/sigmoid.py
+
+X and Y have shape (2*n, channels); bias has one value per channel.
+The code built by make_bias_sigmoid processes one element at a time.
+It writes Y[i, j] = sigmoid(X[i, j] + bias[j]) without an intermediate tensor.
+The symbol n is derived from tensor shapes at runtime.
+Channels selects a model specialization.
+"""
+
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+# begin sigmoid expression
+from tvm_tirx import ir, tirx
+from tvm_tirx.script import from_source
+from tvm_tirx.script import tirx as T
+
+exp_op = ir.Op.get("prim.exp")
+
+
+def sigmoid(x):
+ one = ir.const(1, x.ty)
+ exp_neg_x = ir.Call(exp_op, [-x])
+ return one / (one + exp_neg_x)
+
+
+# end sigmoid expression
+
+# begin primitive types
+index_type = ir.PrimType("int64")
+element_type = ir.PrimType("float32")
+# end primitive types
+
+
+def make_bias_sigmoid(channels=128, activation=sigmoid):
+ """The factory builds scalar computations inside batch and channel
loops."""
+ # begin tensor types
+ n = ir.Var("n", "int64")
+ matrix_type = tirx.TensorType((2 * n, channels), "float32")
+ bias_type = tirx.TensorType((channels,), "float32")
+ # end tensor types
+
+ # begin variables
+ X = ir.Var("X", matrix_type)
+ bias = ir.Var("bias", bias_type)
+ Y = ir.Var("Y", matrix_type)
+ i, j = [ir.Var(name, "int64") for name in ("i", "j")]
+ biased = ir.Var("biased", "float32")
+ # end variables
+
+ zero, step = [ir.const(value, index_type) for value in (0, 1)]
+ end = ir.const(channels, index_type)
+
+ # begin biased activation
+ biased_value = X[i, j] + bias[j]
+ result = activation(biased)
+ # end biased activation
+
+ # begin statement sequence
+ body = ir.SeqStmt(
+ [
+ ir.Bind(biased, biased_value),
+ ir.TensorStore(Y, [i, j], result),
+ ]
+ )
+ # end statement sequence
+
+ # begin loop
+ channel_loop = ir.For(j, zero, end, step, ir.ForKind.DEFAULT, body)
+ batch_loop = ir.For(i, zero, 2 * n, step, ir.ForKind.DEFAULT, channel_loop)
+ # end loop
+
+ # begin function construction
+ function = tirx.Function([X, bias, Y], batch_loop,
ret_type=ir.TupleType([]))
+ function = function.with_attr("global_symbol", "bias_sigmoid")
+ # end function construction
+ return function
+
+
+# begin module construction
+channels_128 = make_bias_sigmoid(128).with_attr("global_symbol",
"bias_sigmoid_128")
+channels_130 = make_bias_sigmoid(130).with_attr("global_symbol",
"bias_sigmoid_130")
+module = ir.IRModule({"bias_sigmoid_128": channels_128, "bias_sigmoid_130":
channels_130})
+# end module construction
+
+# begin module lookup
+entry = module.get_global_var("bias_sigmoid_130")
+function = module[entry]
+# end module lookup
+
+batch_loop = function.body[0]
+channel_loop = batch_loop.body[0]
+binding, store = channel_loop.body.seq
+activation = store.value
+exp_call = activation.b.b
+
+# begin operator lookup
+op = ir.Op.get("prim.exp")
+argument_names = [argument.name for argument in op.args_info]
+effect = op.get_attr("TCallEffectKind")
+# end operator lookup
+
+# begin validation failure
+invalid = ir.Call("prim.exp", [], ty="float32")
+# end validation failure
+try:
+ invalid.validate()
+except TypeError:
+ pass
+else:
+ raise AssertionError("The exponential operator requires its input
argument")
+
+# begin provisional call
+provisional = ir.Call(op, exp_call.args, ty=ir.Type.missing())
+inferred = ir.reinfer_type(provisional)
+rebuilt_call = ir.Call(op, provisional.args, ty=inferred)
+# end provisional call
+exp_call.validate()
+rebuilt_call.validate()
+
+# begin module update
+variant = module.clone()
+replacement = make_bias_sigmoid(130, activation=ir.prim.sigmoid)
+variant[entry] = replacement.with_attr("global_symbol", "bias_sigmoid_130")
+# end module update
+
+# begin module serialization
+restored_module = ir.load_json(ir.save_json(module))
+ir.assert_structural_equal(module, restored_module)
+# end module serialization
+
+# begin print program
+source = function.script()
+print(source)
+# end print program
+
+# begin printing and parsing
+restored_function = from_source(source, extra_vars={"T": T})
+ir.assert_structural_equal(function, restored_function)
+# end printing and parsing
+
+
+# begin save program
+def save_program(program, output_path):
+ source = program.script(extra_config={"ir.comment_imports": False})
+ output_path.write_text(source)
+ restored = from_source(output_path.read_text())
+ ir.assert_structural_equal(program, restored)
+ return restored
+
+
+# end save program
+
+
+def run():
+ """Both specializations and the replacement run at two batch sizes."""
+ import numpy as np
+
+ from tvm_tirx import runtime
+
+ assert tirx.analysis.check_well_formed(module)
+ assert tirx.analysis.check_well_formed(variant)
+ assert inferred == element_type
+ ir.assert_structural_equal(exp_call, rebuilt_call)
+ assert module[entry].same_as(channels_130)
+
+ with TemporaryDirectory() as directory:
+ cached = save_program(variant, Path(directory) / "activation.py")
+ # begin compilation
+ compiled = tirx.compile(cached, target="llvm")
+ # end compilation
+ baseline = tirx.compile(module, target="llvm")
+
+ rng = np.random.default_rng(0)
+ for channels in (128, 130):
+ for batch in (2, 16):
+ # begin execution
+ inputs = rng.normal(size=(batch, channels)).astype("float32")
+ offsets = rng.normal(size=channels).astype("float32")
+ x = runtime.tensor(inputs)
+ bias = runtime.tensor(offsets)
+ output = runtime.tensor(np.full_like(inputs, np.nan))
+ name = f"bias_sigmoid_{channels}"
+ compiled[name](x, bias, output)
+
+ expected = 1.0 / (1.0 + np.exp(-(inputs + offsets)))
+ np.testing.assert_allclose(output.numpy(), expected,
rtol=1e-6, atol=1e-7)
+ # end execution
+ baseline_output = runtime.tensor(np.full_like(inputs, np.nan))
+ baseline[name](x, bias, baseline_output)
+ np.testing.assert_allclose(baseline_output.numpy(), expected,
rtol=1e-6, atol=1e-7)
+ print(
+ f"bias + sigmoid: shape {(batch, channels)}, both
implementations match NumPy"
+ )
+
+
+if __name__ == "__main__":
+ run()
diff --git a/docs/ir/examples/statements_and_functions.py
b/docs/ir/examples/statements_and_functions.py
deleted file mode 100644
index 1126b04aa..000000000
--- a/docs/ir/examples/statements_and_functions.py
+++ /dev/null
@@ -1,78 +0,0 @@
-# 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/tvmscript_sigmoid.py
b/docs/ir/examples/tvmscript_sigmoid.py
new file mode 100644
index 000000000..49205c30e
--- /dev/null
+++ b/docs/ir/examples/tvmscript_sigmoid.py
@@ -0,0 +1,87 @@
+# 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.
+"""This example defines, compiles, and runs a bias + sigmoid program in
TVMScript.
+
+TIRx needs LLVM support. The example runs from the repository root with:
+ uv run --no-sync python docs/ir/examples/tvmscript_sigmoid.py
+
+X and Y have shape (2*n, channels); bias has one value per channel.
+The code in bias_sigmoid processes one element at a time. run() compiles a
+130-channel specialization and compares the output against NumPy.
+"""
+
+# begin program
+from __future__ import annotations
+
+from tvm_tirx import ir, tirx
+from tvm_tirx.script import tirx as T
+
+
+def make_bias_sigmoid(channels=130):
+ n = T.dynamic("n", "int64")
+
+ @T.function
+ def bias_sigmoid(
+ X: T.Tensor((2 * n, channels), T.float32),
+ bias: T.Tensor((channels,), T.float32),
+ Y: T.Tensor((2 * n, channels), T.float32),
+ ):
+ for i in range(2 * n):
+ for j in range(T.int64(channels)):
+ biased = X[i, j] + bias[j]
+ Y[i, j] = T.float32(1) / (T.float32(1) + T.exp(-biased))
+
+ return bias_sigmoid
+
+
+program = ir.IRModule({"bias_sigmoid": make_bias_sigmoid(130)})
+
+
+def run():
+ import numpy as np
+
+ from tvm_tirx import runtime
+
+ function = program["bias_sigmoid"]
+ input_var = function.params[0] # ir.Var named X.
+ tensor_type = input_var.ty # tirx.TensorType: shape=(2 * n, 130),
dtype=float32.
+ row_extent = tensor_type.shape[0] # ir.prim.Mul: a=2, b=n,
ty=ir.PrimType("int64").
+ location = row_extent.loc # Source location of 2 * n in X's type
annotation.
+ print("Input type:", tensor_type.script())
+ print("Rows:", row_extent.script())
+ print("Rows source:", location.source_name.name)
+ print("Rows start:", location.start_line, location.start_column)
+
+ print(program.script())
+ compiled = tirx.compile(program, target="llvm")
+ batch, channels = 16, 130
+ rng = np.random.default_rng(0)
+ inputs = rng.normal(size=(batch, channels)).astype("float32")
+ offsets = rng.normal(size=channels).astype("float32")
+ x = runtime.tensor(inputs)
+ bias = runtime.tensor(offsets)
+ output = runtime.tensor(np.full_like(inputs, np.nan))
+ compiled["bias_sigmoid"](x, bias, output)
+
+ expected = 1.0 / (1.0 + np.exp(-(inputs + offsets)))
+ np.testing.assert_allclose(output.numpy(), expected, rtol=1e-6, atol=1e-7)
+ print("bias + sigmoid matches NumPy.")
+
+
+if __name__ == "__main__":
+ run()
+# end program
diff --git a/docs/ir/examples/types_and_expressions.py
b/docs/ir/examples/types_and_expressions.py
deleted file mode 100644
index 4f8d7b766..000000000
--- a/docs/ir/examples/types_and_expressions.py
+++ /dev/null
@@ -1,100 +0,0 @@
-# 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
index 36d6d5dc1..fb48fc27f 100644
--- a/docs/ir/functions_and_modules.rst
+++ b/docs/ir/functions_and_modules.rst
@@ -18,17 +18,24 @@
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.
+In this chapter, we will introduce functions and modules as the structures
+that organize IR programs. A function combines parameters, a result type,
+and a statement body; a module collects named function definitions. Compiler
+passes use these structures to inspect a function's interface and
+implementation, locate definitions through global references, and replace
+an implementation while preserving references to it.
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:
+For tensor programs, use ``@T.function`` to define functions in TVMScript,
+as in our sigmoid example. The decorator produces a
+:class:`tirx.Function <tvm_tirx.tirx.Function>` IR object. To help users better
+understand the IR, we also show how to construct the same representation
+directly through the IR API.
+
+``ir.BaseFunc`` is the base class for all function definitions and is also
+an expression. Its ``attrs`` field holds named function properties:
.. list-table::
:header-rows: 1
@@ -41,7 +48,8 @@ shared by function definitions:
- ``ir.DictAttrs``
- Named function properties.
-``tirx.Function`` adds these fields:
+:class:`tirx.Function <tvm_tirx.tirx.Function>` extends ``ir.BaseFunc`` with
+parameters, a result type, and a statement body:
.. list-table::
:header-rows: 1
@@ -52,83 +60,82 @@ shared by function definitions:
* - ``params``
- Parameter variables in calling order.
* - ``ret_type``
- - The declared result type; an omitted type becomes an empty
``TupleType``.
+ - The declared result type.
* - ``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:
+The parameter and result types form the function's signature, available as
+an ``ir.FuncType`` through ``ty``. The ``global_symbol`` entry in ``attrs``
+specifies the function's exported name for code generation. For public
+functions, ``@T.function`` sets it to the Python function name by default.
+The IR API example below sets it explicitly with ``with_attr``.
+
+The program defines a function with input, bias, and output tensors.
+Its IR definition combines these three parameters with the batch and
+channel loops:
-.. literalinclude:: examples/statements_and_functions.py
+.. literalinclude:: examples/sigmoid.py
:language: python
:start-after: # begin function construction
:end-before: # end function construction
- :prepend: from tvm_tirx import ir, tirx
+ :dedent: 4
-``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.
+``batch_loop`` contains a channel loop that processes one element at a time.
+The function writes activated values into ``Y``; ``ir.TupleType([])`` declares
+that it returns no value.
+The input and output types have shape ``(2*n, channels)``. The symbol ``n``,
+whose value the compiled function derives from the tensor shapes at runtime,
+is used in these types and the loop bound.
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``:
+:class:`ir.IRModule <tvm_tirx.ir.IRModule>` organizes named functions into a
+program and lets us look up their definitions. An
+:class:`ir.GlobalVar <tvm_tirx.ir.GlobalVar>` is an expression that refers to a
+function stored in the module. This differs from the :ref:`ir.Var
<user-ir-var>`
+objects used for parameters and local values in our example. The module
+associates each function reference with its ``ir.BaseFunc`` definition, such
+as a ``tirx.Function``.
-.. literalinclude:: examples/statements_and_functions.py
+In the bias + sigmoid example, one module keeps two model specializations.
+The same factory builds kernels for 128-channel and 130-channel outputs,
+each with its own exported name:
+
+.. literalinclude:: examples/sigmoid.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.
+Both accept tensors with a runtime row count of ``2*n``. Each specialization
+uses a scalar loop over its channels. Their tensor types record the different
+widths.
-Use the stored reference to construct a call representing ``main(3)``:
+Use ``module.get_global_var(name)`` to retrieve the ``ir.GlobalVar`` reference
+for an existing function:
-.. literalinclude:: examples/statements_and_functions.py
+.. literalinclude:: examples/sigmoid.py
:language: python
- :start-after: # begin global call
- :end-before: # end global call
+ :start-after: # begin module lookup
+ :end-before: # end module lookup
-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.
+``entry`` is the function's ``ir.GlobalVar`` reference; ``module[entry]`` and
+``module["bias_sigmoid_130"]`` retrieve its current definition. A compiler pass
+can replace that definition while keeping ``entry``. If the replacement
+preserves the function's calling interface, calls using this reference can
+remain unchanged.
-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:
+For example, a compiler can keep the original expanded sigmoid and build
+another module that uses the named ``prim.sigmoid`` operator. The factory's
+``activation`` argument selects this representation:
-.. literalinclude:: examples/statements_and_functions.py
+.. literalinclude:: examples/sigmoid.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.
+``variant`` keeps the 128-channel definition and replaces the 130-channel
+definition with one that calls ``prim.sigmoid``. The replacement keeps the
+same tensor interface and exported name, so ``variant[entry]`` now retrieves
+the new implementation through the existing global reference. The original
+``module`` still contains the expanded sigmoid expression.
diff --git a/docs/ir/operators_and_calls.rst b/docs/ir/operators_and_calls.rst
index db33168b8..b38230080 100644
--- a/docs/ir/operators_and_calls.rst
+++ b/docs/ir/operators_and_calls.rst
@@ -18,46 +18,44 @@
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.
+In this chapter, we will introduce operators and calls, from an operator's
+shared definition to the arguments and result type of an individual call.
+We will examine argument requirements, type inference, validation, and
+effects. Built-in and custom operators share the same IR representation,
+so compiler passes can inspect and transform both through the existing
+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.
+``ir.Op`` is an expression that identifies a registered computation. We can
+retrieve it with ``ir.Op.get(name)`` and inspect ``args_info`` for the expected
+inputs.
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:
+The program uses the same exponential operator for each input element.
+Its shared argument and effect information can be inspected through the
operator:
-.. literalinclude:: examples/operators_and_calls.py
+.. literalinclude:: examples/sigmoid.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.
+``argument_names`` is ``["x"]``, and ``effect`` is
+``ir.CallEffectKind.Pure``. This means the exponential call computes its
+result from its input without reading or changing external state. The tensor
+loads supplying that input are separate expressions. The
:class:`~tvm_tirx.ir.Op` API reference
+describes its argument information and attributes.
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:
+:class:`ir.Call <tvm_tirx.ir.Call>` represents an invocation of an operator or
+function. The expression being called is its *callee*. The call records that
+callee, its arguments, and any attributes for the invocation:
.. list-table::
:header-rows: 1
@@ -75,121 +73,74 @@ being called is its *callee*. A call inherits ``ty`` and
``loc`` from
- 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.
+In the bias + sigmoid example, the activation can be an expanded expression,
+``1 / (1 + exp(-x))``, or a single call to ``prim.sigmoid``. A named
+activation operator lets a later compiler pass choose its implementation:
-For example, this call applies the registered exponential operator to a float32
-constant:
+.. code-block:: python
-.. 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.
+ x = ir.Var("x", "float32")
+ call = ir.prim.sigmoid(x)
-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
+:func:`ir.prim.sigmoid <tvm_tirx.ir.prim.sigmoid>` is a helper that constructs
+an ``ir.Call`` to the registered ``prim.sigmoid`` operator. In this example,
+it is equivalent to ``ir.Call("prim.sigmoid", [x])``, which resolves the
+operator name through ``ir.Op.get``. The call's result has the same type
+as ``x``: float32 in this example.
-``function_call`` represents ``f(1.0)``. The variable supplies a signature;
-the example does not provide an implementation of ``f``.
+In the example above, constructing ``call`` records the sigmoid computation
+in the IR; the computation runs when the compiled program executes.
+The same ``ir.Call`` representation can invoke a function as well as a
+registered operator. For example, its ``op`` field can hold a global reference
+to a function in a module, as we'll see in :doc:`functions_and_modules`.
.. _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.
+``ir.Call`` can infer its result type from the operator's registered rule or
+the callee's function signature. The constructor infers the type when ``ty``
+is omitted and stores the supplied type when ``ty`` is explicit.
+``MissingType`` marks a result whose type has not yet been determined.
+For example:
+
+.. code-block:: python
+
+ x = ir.Var("x", "float32")
+ ir.Call("prim.exp", [x]) # Infers float32.
+ ir.Call("prim.exp", [x], ty="float32") # Uses the supplied float32.
+ ir.Call("prim.exp", [x], ty=ir.Type.missing()) # Keeps MissingType; no
inference.
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.
+checking its arguments with :meth:`call.validate()
<tvm_tirx.ir.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:
+To see how validation catches missing arguments, we can construct a
+``prim.exp`` call with no input and an explicit result type:
-.. literalinclude:: examples/operators_and_calls.py
+.. literalinclude:: examples/sigmoid.py
:language: python
:start-after: # begin validation failure
:end-before: # end validation failure
- :append: invalid.validate() # raises TypeError: expected one argument
+ :append: invalid.validate()
+
+The explicit result type bypasses type inference. ``invalid.validate()`` then
+raises ``TypeError`` because ``prim.exp`` requires one input 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:
+When a pass rebuilds a call, its result type must match the arguments.
+Here, ``exp_call`` represents ``T.exp(-biased)``, whose input is a float32
+value. The following code reconstructs the call with a missing type, then uses
+:func:`ir.reinfer_type <tvm_tirx.ir.reinfer_type>` to infer its result type:
-.. literalinclude:: examples/operators_and_calls.py
+.. literalinclude:: examples/sigmoid.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.
+``inferred`` is ``ir.PrimType("float32")``, matching the input's type.
+The inference function returns this type without changing ``provisional``;
+``rebuilt_call`` stores it in its ``ty`` field. The rebuilt call can then be
+checked with ``rebuilt_call.validate()`` before a pass uses it in place of
+the original.
diff --git a/docs/ir/overview.rst b/docs/ir/overview.rst
index db392288b..b86bd798d 100644
--- a/docs/ir/overview.rst
+++ b/docs/ir/overview.rst
@@ -18,20 +18,57 @@
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.
+In :doc:`/get_started/example`, we wrote and ran our first TIRx program. To
+improve efficiency, we need to convert the program into a common intermediate
+representation (IR) and apply compiler passes to it. In this section, we'll
+explore the IR system in tvm-tirx in detail.
+
+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``
+ * - Statements
+ - ``ir.Stmt``
+ - Describe execution order and control flow.
+ - ``B[i] = A[i] + 1``
+ * - Types
+ - ``ir.Type``
+ - Describe value properties, such as dtypes and tensor shapes.
+ - ``T.float32``
+ * - 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})``
+ * - Source locations
+ - ``ir.Location``
+ - Record where IR nodes originate in the source.
+ - The filename and line number of ``A[i] + 1``.
+
+The IR 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.
+together. Built on shared C++ IR infrastructure, 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.
+The following tree shows the detailed inheritance hierarchy of the core IR
classes.
.. code-block:: text
@@ -85,13 +122,7 @@ is available as ``ir.prim.Add`` in Python.
│ ├── 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
+ │ └── prim.Shuffle Vector lane selection
├── ir.Type Types
│ ├── ir.MissingType Unresolved type
information
│ ├── ir.NoneType The type of None
@@ -118,49 +149,30 @@ is available as ``ir.prim.Add`` in Python.
│ ├── 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
+ └── ir.IRModule Global function
definitions
.. _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`.
+After the add-one example in :doc:`/get_started/example`, this walkthrough
+uses a bias + sigmoid program to introduce symbolic shapes and scalar
computations.
+It computes ``Y[i, j] = 1 / (1 + exp(-(X[i, j] + bias[j])))``.
+Download :download:`tvmscript_sigmoid.py <examples/tvmscript_sigmoid.py>`
+and run ``python tvmscript_sigmoid.py`` with an LLVM-enabled TIRx build.
+
+* 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.
diff --git a/docs/ir/printing_and_parsing.rst b/docs/ir/printing_and_parsing.rst
index df2cb8566..dc9573a40 100644
--- a/docs/ir/printing_and_parsing.rst
+++ b/docs/ir/printing_and_parsing.rst
@@ -20,120 +20,85 @@
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.
+In this chapter, we will introduce TVMScript and the interfaces for printing
+and parsing IR. TVMScript provides a Python-based notation for program
+definitions, while the printer and parser convert between IR objects and
+source text. These interfaces allow generated programs to be inspected,
+saved, and reconstructed. Compiler extensions can reuse them with custom
+builder namespaces and 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:
+The ``@T.function`` decorator parses a Python function definition into a
+``tirx.Function``. The parser reads the definition's source from its Python
+file; the numerical computation runs when the compiled function is called.
-.. 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
+The example program defines ``bias_sigmoid`` inside a Python factory:
-``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.
+.. literalinclude:: examples/tvmscript_sigmoid.py
+ :language: python
+ :start-at: from __future__ import annotations
+ :end-at: return bias_sigmoid
-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:
+The factory creates ``n`` with ``T.dynamic`` before the function definition
+and reuses it in the tensor shapes and loop bound. The decorated definition
+records the loops, arithmetic, and indexing as IR.
-.. literalinclude:: examples/printing_and_parsing.py
- :language: python
- :start-after: # begin program factory
- :end-before: # end program factory
+For printing and parsing, we reuse ``module`` from
+:doc:`functions_and_modules`.
-``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.
+The ``script()`` method returns TVMScript source for an IR object, so
+generated loops and expressions can be inspected before lowering. The
+130-channel specialization can be printed with:
-Printing lets you inspect what this Python code generated:
+.. code-block:: python
-.. literalinclude:: examples/printing_and_parsing.py
- :language: python
- :start-after: # begin print generated function
- :end-before: # end print generated function
+ function = module["bias_sigmoid_130"]
+ source = function.script()
+ print(source)
-``script()`` returns source text. The printed ``wide`` function contains these
-loops, using the selected width and element type:
+The printed loop structure shows the iteration over samples and channels:
.. 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
+ for j in range(T.int64(0), T.int64(130)):
+ ...
-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 full source includes the bias addition, sigmoid expression, and output
+store. It keeps ``n`` symbolic so the tensor shapes and outer loop share
+the same symbol when parsed again.
-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:
+Smaller objects can be inspected too. ``function.params[0].ty.script()``
+prints the input tensor type. ``function.show()`` displays highlighted source.
+The :meth:`~tvm_tirx.tirx.Function.show` reference describes line numbers and
+object annotations.
+These interfaces also work on functions produced by a compiler pass.
-.. literalinclude:: examples/printing_and_parsing.py
- :language: python
- :start-after: # begin module round trip
- :end-before: # end module round trip
+The :doc:`TVMScript parser </reference/python/script.parser>` converts source
+text back into IR through ``from_source``. A printed specialization can
+therefore be reconstructed and checked before compilation:
+
+.. code-block:: python
+
+ from tvm_tirx.script import from_source
+
+ restored_function = from_source(source, extra_vars={"T": T})
+ ir.assert_structural_equal(function, restored_function)
+
+The default source comments out its imports, so ``extra_vars`` supplies the
+``T`` namespace from ``tvm_tirx.script.tirx``. The result is a
``tirx.Function``
+that can be inspected or passed to ``tirx.compile``.
-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:
+A self-contained source file needs its imports. With
+``ir.comment_imports=False``, the :doc:`printer
</reference/python/script.printer>`
+includes them in the generated source. The helper below saves and reloads a
+whole module with those imports:
-.. literalinclude:: examples/printing_and_parsing.py
+.. literalinclude:: examples/sigmoid.py
:language: python
- :start-after: # begin standalone source
- :end-before: # end standalone source
+ :pyobject: save_program
-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.
+``output_path`` is a ``pathlib.Path``. The saved file contains the module's
+TVMScript definition and its imports, so ``from_source`` can reconstruct the
+module from that file. The structural check verifies that the reconstructed
+module has the same IR structure as the original. The returned module can
+then be inspected, transformed, or compiled with ``tirx.compile``.
diff --git a/docs/ir/statements.rst b/docs/ir/statements.rst
index 091c6c1d9..f5d7c2368 100644
--- a/docs/ir/statements.rst
+++ b/docs/ir/statements.rst
@@ -20,33 +20,21 @@
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.
+In this chapter, we will introduce the statements that organize a
+function's execution. Statements can be used to describe value bindings, tensor
+writes, and control flow. Statements take several forms, including sequences,
+branches, and loops. Sequences record execution order, while branches and
+loops determine which statements execute and how often. Compiler passes can
+use these structures to construct and transform programs, including bodies
+defined by custom operators.
.. _statements-and-sequences:
-Statement Nodes
----------------
+Statement Classes
+-----------------
-``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:
+``ir.Stmt`` is the base class for all statements, from bindings and tensor
+writes to sequences, branches, and loops:
.. list-table::
:header-rows: 1
@@ -54,7 +42,7 @@ The shared statement classes are:
* - Class
- Represents
- - Additional fields
+ - Fields
* - ``ir.Bind``
- A variable definition
- ``var``, ``value``
@@ -95,96 +83,78 @@ The shared statement classes are:
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
+One of the most common statement types is ``ir.SeqStmt``, which groups
+statements in execution order in its ``seq`` field.
+A sequence can serve as a function body or as the body of a branch or loop.
- def increment(x):
- result = x + 1
- return result
+In our sigmoid example, each iteration binds the biased input before storing
+the sigmoid result. Reusing the tensor variables and indices from
+:ref:`user-ir-var` and the helper from :ref:`user-ir-call`, we can build this
+sequence:
-The following code constructs that function body as IR objects:
+.. code-block:: python
-.. literalinclude:: examples/statements_and_functions.py
- :language: python
- :start-after: # begin statements
- :end-before: # end statements
- :prepend: from tvm_tirx import ir, tirx
+ biased = ir.Var("biased", "float32")
+ biased_value = X[i, j] + bias[j]
+ result = sigmoid(biased)
-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.
+ body = ir.SeqStmt(
+ [
+ ir.Bind(biased, biased_value),
+ ir.TensorStore(Y, [i, j], result),
+ ]
+ )
-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.
+``ir.Bind`` defines ``biased`` with the value of ``X[i, j] + bias[j]``.
+The following ``ir.TensorStore`` evaluates the sigmoid expression ``result``
+using that binding and writes the value to ``Y[i, j]``.
+The sequence forms the body of the channel loop, which we examine next.
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.
+Control-flow statements determine which parts of a program execute, how
+many times they execute, and when execution leaves a loop or function.
+TIRx provides ``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.
+An ``ir.For`` specifies a loop variable, begin and end bounds, a step, and an
+execution kind. The end bound is exclusive. ``ForKind.DEFAULT`` executes
+iterations in order; the :class:`~tvm_tirx.ir.ForKind` API reference describes
+the other execution kinds.
-A loop instead repeats a body. This loop evaluates each index from zero through
-three:
+The program processes one channel at a time. The surrounding loops cover
+samples and channels, keeping each tensor access within its bounds:
-.. literalinclude:: examples/statements_and_functions.py
+.. literalinclude:: examples/sigmoid.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.
+ :dedent: 4
+
+``batch_loop`` advances ``i`` by one sample. ``channel_loop`` advances ``j``
+by one channel, executing ``body`` at each step. Both loops use int64
+variables, bounds, and steps. The channel loop stops before ``end``, the
+channel count, so ``j`` ranges from 0 to ``channels - 1``.
+
+Scope determines where a definition can be used. ``ir.SeqStmt`` groups
+statements without introducing an additional scope. In the example, the
+sequence ensures that ``biased`` is bound before the ``TensorStore`` uses it,
+while the enclosing channel loop limits its visibility to that loop body.
+The loop variables ``i`` and ``j`` are likewise visible only within their
+loop bodies.
+
+Some operators introduce their own statement bodies.
+:class:`ir.RegionStmt <tvm_tirx.ir.RegionStmt>` represents these using a
+registered operator, whose definition specifies the body's parameters.
+
+These scope rules also apply to programs built through the IR API. After
+construction, :func:`tirx.analysis.check_well_formed
<tvm_tirx.tirx.analysis.check_well_formed>`
+can check a function or module for invalid variable and tensor uses and
+misplaced loop-control statements.
diff --git a/docs/ir/types_and_expressions.rst
b/docs/ir/types_and_expressions.rst
index 4ab478ac0..2bc991cba 100644
--- a/docs/ir/types_and_expressions.rst
+++ b/docs/ir/types_and_expressions.rst
@@ -20,36 +20,19 @@
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.
+In this chapter, we will introduce types and expressions and explain how
+they work together in the IR. Types describe properties such as dtypes,
+tensor shapes, and function signatures, while expressions represent values
+and computations. Compiler passes can work with these objects directly to
+inspect or rewrite a program without executing it.
.. _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:
+``ir.Type`` is the base class for all IR types, including primitive types
+for scalar and vector values, tensor types, and function types:
.. list-table::
:header-rows: 1
@@ -57,7 +40,7 @@ of values. Common examples include:
* - Class
- Describes
- - Additional fields
+ - Fields
* - ``ir.PrimType``
- A scalar or vector dtype
- ``dtype``
@@ -74,44 +57,96 @@ of values. Common examples include:
- 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:
+Next, we will look at two common types in tensor programs in more detail:
+``ir.PrimType`` for scalar and vector values, and ``tirx.TensorType`` for
tensors.
-.. 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
+.. _user-ir-prim-type:
-``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.
+PrimType
+~~~~~~~~
-.. _user-tensor-type:
+``ir.PrimType`` is a widely used type in TIRx IR, which describes integer,
+floating-point, and boolean values, either as scalars or vectors. Its ``dtype``
+records the element kind, bit width, and number of lanes.
-``tirx.TensorType`` combines an element type with shape and storage
information:
+In the sigmoid program, primitive types describe both shape symbols and
+computed values. The dynamic shape symbol ``n`` is declared as:
-.. literalinclude:: examples/types_and_expressions.py
+.. literalinclude:: examples/tvmscript_sigmoid.py
:language: python
- :start-after: # begin constant tensor type
- :end-before: # end constant tensor type
+ :start-at: n = T.dynamic
+ :end-at: n = T.dynamic
+ :dedent: 4
-``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.
+where ``n`` is a 64-bit integer symbol with type ``ir.PrimType("int64")``.
-Use ``==`` to compare type structure and ``same_as`` to compare object
identity.
+The sigmoid calculation uses floating-point values:
-.. _user-ir-var:
+.. literalinclude:: examples/tvmscript_sigmoid.py
+ :language: python
+ :start-at: biased =
+ :end-at: Y[i, j] =
+ :dedent: 16
+
+In this example, ``X`` and ``bias`` are declared as tensors with float32
+elements, so ``X[i, j]`` and ``bias[j]`` have that type. ``T.float32(1)``
supplies
+float32 constants for the formula. The arithmetic operations and ``T.exp``
+preserve this type, so ``biased`` and the sigmoid result are automatically
+inferred as float32 during expression construction. In the IR, this type is
+represented by ``ir.PrimType("float32")``.
+You can also pass an expression to a ``PrimType`` instance to cast it to that
type.
+
+To create a variable of a specific type, you can call the corresponding
+``PrimType`` instance without arguments:
+
+.. code-block:: python
+
+ f32 = ir.PrimType("float32")
+ x = f32()
+
+``PrimType`` also provides methods to inspect a type, like ``is_scalar()``,
+``matches_code(...)``, and ``matches_element_type(...)``. See the
+:class:`~tvm_tirx.ir.PrimType` API reference for more details.
+
+.. _user-tensor-type:
+
+TensorType
+~~~~~~~~~~
+
+Most computations in TIRx involve tensors. To define a tensor type, we need
+two important properties: its element type and shape.
+:class:`tirx.TensorType <tvm_tirx.tirx.TensorType>` stores these in its
``dtype``
+and ``shape`` fields, respectively. The element type is a ``PrimType``, and
+each dimension is an integer expression.
+
+In the sigmoid program, the input ``X`` is declared as:
+
+.. literalinclude:: examples/tvmscript_sigmoid.py
+ :language: python
+ :start-at: X: T.Tensor
+ :end-at: X: T.Tensor
+ :dedent: 8
+
+where ``X`` is the parameter name and the part after ``:`` is its *type
+annotation*, which tells TVMScript what kind of value the parameter accepts.
+This annotation becomes a ``tirx.TensorType`` in the IR: ``shape`` describes
+``2*n`` rows and ``channels`` columns, and ``dtype`` is
+``ir.PrimType("float32")``.
+
+**Symbolic shapes.** A dimension such as ``2*n`` can vary between calls.
+In this example, ``channels`` is fixed at 130 when the program is constructed,
+while ``n`` is determined from the input's row count at runtime. An input
+with shape ``(16, 130)`` gives ``n=8``. A later call can use a different
+even row count without recompiling the function. Both fixed shapes such as
+``(16, 130)`` and symbolic shapes such as ``(2*n, 130)`` use the same
+``TensorType`` representation.
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:
+``ir.Expr`` is the base class for all expressions, which represent program
+values and the operations that produce them. Each expression's ``ty`` field
+records its result type:
.. list-table::
:header-rows: 1
@@ -122,16 +157,10 @@ inspect its operands and transformations can replace
parts of the computation.
- 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.
+ - The type of value represented by the expression.
-Subclasses inherit ``ty`` and ``loc`` and add fields for their stored values or
-operands. The following examples show common expression classes and selected
-fields:
+Expressions include variables, constants, arithmetic, tensor reads, and calls.
+These classes store their values or operands in the following fields:
.. list-table::
:header-rows: 1
@@ -162,113 +191,175 @@ fields:
- 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``.
+In our sigmoid example, the input parameter ``X`` is represented by an
+``ir.Var`` expression. We can inspect its type through the ``program`` module
+(introduced later in :doc:`functions_and_modules`):
-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/tvmscript_sigmoid.py
+ :language: python
+ :start-at: function = program["bias_sigmoid"]
+ :end-at: location = row_extent.loc
+ :dedent: 4
+
+The ``ir.prim.Mul`` expression for ``2 * n`` is one example of a primitive
operation.
+Next, we'll look at how primitive operations represent the arithmetic in the
+sigmoid formula.
+
+.. _user-ir-prim-operations:
+
+Primitive Operations
+~~~~~~~~~~~~~~~~~~~~
-.. literalinclude:: examples/types_and_expressions.py
+Primitive operations give the compiler a common way to represent and reason
+about basic computations. The ``ir.prim`` namespace provides expressions such
+as ``Add`` for arithmetic and ``LE`` for comparisons. These expressions have
+defined semantics, explicit operands, and result types, allowing compiler
+passes to analyze computations and perform optimizations such as constant
+folding and algebraic simplification. See the
+:doc:`ir.prim </reference/python/ir.prim>` API reference for the available
+expression classes and construction helpers.
+
+TIRx overloads Python operators such as ``+`` and ``/`` to construct IR
+expressions. In our program, we use them as follows:
+
+.. literalinclude:: examples/tvmscript_sigmoid.py
:language: python
- :start-after: # begin operator expressions
- :end-before: # end operator expressions
+ :start-at: biased =
+ :end-at: Y[i, j] =
+ :dedent: 16
+
+In this code, the ``+`` and ``/`` operations become
+``ir.prim.Add`` and ``ir.prim.Div`` expressions, respectively.
+Each binary expression has two operand fields: ``a`` for the expression on its
+left and ``b`` for the expression on its right.
-``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.
+**Bias addition.** The ``+`` in the definition of ``biased`` becomes an
+``ir.prim.Add``. Its ``a`` field points to the ``X[i, j]`` tensor read, and
+its ``b`` field points to the ``bias[j]`` tensor read. The addition has this
+structure:
-Operands and call arguments are themselves expressions, allowing these forms
-to be combined into larger computations.
+.. code-block:: text
+
+ ir.prim.Add X[i, j] + bias[j]
+ ├── a: ir.TensorLoad X[i, j]
+ └── b: ir.TensorLoad bias[j]
+
+Each tensor read produces one float32 value, and the addition combines
+the input element with its channel's bias.
+
+**Sigmoid.** In the final line, the ``+`` in
+``T.float32(1) + T.exp(-biased)`` becomes an ``ir.prim.Add``. The ``/`` becomes
+an ``ir.prim.Div`` whose ``a`` operand supplies the numerator, ``1``, and
+whose ``b`` operand is the addition that computes the denominator.
+The exponential ``T.exp(-biased)`` is represented by an ``ir.Call``, which
+we will discuss in the :ref:`user-ir-call` section.
.. _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:
+Primitive Expressions
+~~~~~~~~~~~~~~~~~~~~~
-.. literalinclude:: examples/types_and_expressions.py
- :language: python
- :start-after: # begin scalar expressions
- :end-before: # end scalar expressions
+A primitive expression is an expression whose ``ty`` is an ``ir.PrimType``.
+The expression may be a variable, a tensor element read, an arithmetic
+calculation, or a call. Its result type tells us whether it produces an
+integer, a floating-point value, a boolean, or a vector of these values.
+The float32 tensor reads and computations in our sigmoid example are primitive
+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.
+**Type inference.** Result types are usually determined when expressions are
+constructed and stored in their ``ty`` fields. Overloaded arithmetic operators
+such as ``+`` match operand types, promoting compatible types and inserting
+casts where needed. Comparisons produce boolean results:
-A tensor variable follows the same pattern, with a ``TensorType`` in its
-``ty`` field. Its shape can reuse the integer expressions described above:
+.. code-block:: python
-.. literalinclude:: examples/types_and_expressions.py
- :language: python
- :start-after: # begin tensor expressions
- :end-before: # end tensor expressions
+ from tvm_tirx import ir
-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")``.
+ x = ir.Var("x", "float32")
+ (x + 1).ty # ir.PrimType("float32")
+ (x > 0).ty # ir.PrimType("bool")
+ ir.Call("prim.exp", [x]).ty # 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.
+For ``ir.Call``, omitting ``ty`` uses the operator's registered result-type
+rule or the callee's function signature. ``ir.reinfer_type(call)`` recomputes
+the result type without modifying ``call.ty`` or validating the call.
+See :ref:`types-and-validation` for inference rules and explicit validation.
-.. _user-ir-register-type:
+**Index expressions.** The loop ``for i in range(2 * n)`` uses ``2*n`` to
+determine how many rows to process. The multiplication's ``ty`` is
+``ir.PrimType("int64")``. The loop indices ``i`` and ``j`` also have this
+type and select one row and one channel, respectively.
-Register New Types
-------------------
+.. _user-ir-var:
-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.
+Var
+~~~
-In Python, use ``tvm_ffi.dataclasses.py_class``; in Rust, use
-``tvm_ffi::rs_class`` and ``tvm_ffi::def_object_ref!``.
+Variables can be function parameters, loop indices, or intermediate values
+defined by bindings. ``ir.Var`` can be used to represent these values together
+with their names and types.
-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:
+We can declare the variables in our sigmoid example using the IR API:
-.. literalinclude:: examples/types_and_expressions.py
+.. literalinclude:: examples/sigmoid.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.
+ :start-after: # begin variables
+ :end-before: # end variables
+ :dedent: 4
+
+``X``, ``bias``, and ``Y`` have tensor types. ``i`` selects a sample, and ``j``
+selects a channel. These indices have int64 types. ``biased`` holds one
+float32 value after adding the channel's bias to the input element.
+
+.. _user-ir-call:
+
+Call
+~~~~
+
+After defining a function (introduced later in :doc:`functions_and_modules`),
+we can call it with arguments. :doc:`ir.Call <operators_and_calls>` represents
+such a call in the IR. Its ``op`` field (an ``ir.Op`` for a registered
operator)
+identifies the function or operator being called, and ``args`` holds the
+argument expressions. Using an unregistered operator name in ``ir.Op.get``
+or ``ir.Call`` raises ``AttributeError``.
+
+In sigmoid, ``1 / (1 + exp(-x))``, the exponential is a call inside a larger
+arithmetic expression:
+
+.. code-block:: python
+
+ from tvm_tirx import ir, tirx
+ from tvm_tirx.script import from_source
+ from tvm_tirx.script import tirx as T
+
+ exp_op = ir.Op.get("prim.exp")
+
+ def sigmoid(x):
+ one = ir.const(1, x.ty)
+ exp_neg_x = ir.Call(exp_op, [-x])
+ return one / (one + exp_neg_x)
+
+``exp_op`` is the registered ``ir.Op`` for the exponential operation.
+``ir.Call(exp_op, [-x])`` constructs a call with the negated input as its
+argument. The sum ``one + exp_neg_x`` forms the denominator; the division
+produces sigmoid's result.
+The helper preserves the input's primitive type, float32 in this example.
+
+We can construct IR by calling a Python helper directly. Here,
+``sigmoid(biased)`` runs the helper above, which builds an ``ir.Call`` for the
+exponential and returns the complete sigmoid expression:
+
+.. code-block:: python
+
+ biased_value = X[i, j] + bias[j]
+ result = sigmoid(biased)
+
+``biased_value`` is an ``ir.prim.Add`` containing two tensor reads. The program
+binds this value to ``biased`` before using it in the sigmoid calculation.
+``result`` is an ``ir.prim.Div`` whose denominator contains another addition
+and the exponential call. The program stores this result in ``Y[i, j]``.
+
+When the compiled program runs, each loop iteration reads the input and bias
+values, computes the activation, and writes it to ``Y``. The biased values stay
+local to that iteration, so bias addition and sigmoid share one traversal
+without an intermediate tensor.
