This is an automated email from the ASF dual-hosted git repository.
Seven-Streams 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 783ccba99 [DOCS] Reuse Quick Start examples throughout the IR
walkthrough (#82)
783ccba99 is described below
commit 783ccba996c6529f5047284fe0fe768ebbc878f1
Author: Shushi Hong <[email protected]>
AuthorDate: Mon Oct 26 22:42:54 2026 -0400
[DOCS] Reuse Quick Start examples throughout the IR walkthrough (#82)
Rename the introductory tutorial to Quick Start and reuse its CPU/GPU
add-one programs throughout the IR walkthrough. Add field diagrams to
explain how IR nodes relate, with supplementary vector and sigmoid
examples for concepts beyond the base program.
Clarify Call construction, type inference, validation, statement bodies,
variable scope, and TVMScript region syntax.
Validation:
- Pre-commit checks passed.
- Sphinx HTML build passed with warnings treated as errors.
- Companion examples passed IR/type checks, print/parse round trips, and
CPU numerical checks using LLVM and C.
---------
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 | 18 +-
docs/get_started/examples/add_one_cuda.py | 19 +-
docs/get_started/quick_start.rst | 77 +++++
docs/index.rst | 2 +-
docs/ir/examples/sigmoid.py | 217 -------------
docs/ir/examples/tvmscript_sigmoid.py | 87 ------
docs/ir/functions_and_modules.rst | 236 ++++++++++-----
docs/ir/operators_and_calls.rst | 439 ++++++++++++++++++++++-----
docs/ir/overview.rst | 59 ++--
docs/ir/printing_and_parsing.rst | 333 ++++++++++++++++----
docs/ir/statements.rst | 254 +++++++++++-----
docs/ir/types_and_expressions.rst | 486 ++++++++++++++++++------------
13 files changed, 1395 insertions(+), 891 deletions(-)
diff --git a/docs/get_started/example.rst b/docs/get_started/example.rst
deleted file mode 100644
index 241295607..000000000
--- a/docs/get_started/example.rst
+++ /dev/null
@@ -1,59 +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.
-
-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
index 0a455dfac..d5ead44f4 100644
--- a/docs/get_started/examples/add_one_cpu.py
+++ b/docs/get_started/examples/add_one_cpu.py
@@ -17,24 +17,24 @@
"""Add one to each element on the CPU. Requires LLVM support."""
# begin program
-import numpy as np
+import torch
-from tvm_tirx import runtime, tirx
+from tvm_tirx import tirx
from tvm_tirx.script import tirx as T
n = T.dynamic("n")
@T.function
-def add_one(A: T.Tensor((n,), "float32"), B: T.Tensor((n,), "float32")):
+def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.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())
+add_one_compiled = tirx.compile(add_one, target="llvm")
+a = torch.arange(10, dtype=torch.float32)
+b = torch.empty(a.shape, dtype=a.dtype)
+add_one_compiled(a, b)
+torch.testing.assert_close(b, a + 1)
+print(b)
# end program
diff --git a/docs/get_started/examples/add_one_cuda.py
b/docs/get_started/examples/add_one_cuda.py
index 2db908940..b2bf61ed3 100644
--- a/docs/get_started/examples/add_one_cuda.py
+++ b/docs/get_started/examples/add_one_cuda.py
@@ -17,27 +17,26 @@
"""Add one to each element on an NVIDIA GPU. Requires LLVM and CUDA support."""
# begin program
-import numpy as np
+import torch
-from tvm_tirx import cuda, runtime, tirx
+from tvm_tirx import tirx
from tvm_tirx.script import tirx as T
n = T.dynamic("n")
@T.function
-def add_one(A: T.Tensor((n,), "float32"), B: T.Tensor((n,), "float32")):
+def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.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())
+add_one_compiled = tirx.compile(add_one, target="cuda")
+a = torch.arange(10, dtype=torch.float32, device="cuda:0")
+b = torch.empty(a.shape, dtype=a.dtype, device=a.device)
+add_one_compiled(a, b)
+torch.testing.assert_close(b, a + 1)
+print(b)
# end program
diff --git a/docs/get_started/quick_start.rst b/docs/get_started/quick_start.rst
new file mode 100644
index 000000000..d301bf91e
--- /dev/null
+++ b/docs/get_started/quick_start.rst
@@ -0,0 +1,77 @@
+.. 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.
+
+Quick Start
+===========
+
+In this chapter, we'll write and run a small TIRx program using TVMScript,
+which lets us describe tensor computations with Python syntax. We'll also
+inspect the IR that TVMScript creates from the code. The examples use
+PyTorch tensors; the :doc:`/dev/source_build` guide explains how to set up
+TIRx.
+
+Our example computes ``B[i] = A[i] + 1``, where ``A`` is the input tensor
+and ``B`` is the output. Both tensors have the same dynamic length ``n``,
+so the compiled program can handle different input lengths without
+recompilation.
+
+We'll show two versions of this program: one for the CPU and one for an
+NVIDIA GPU. We'll first compile and run the CPU version using LLVM.
+
+.. literalinclude:: examples/add_one_cpu.py
+ :language: python
+ :start-after: # begin program
+ :end-before: # end program
+
+``@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. The compiled function accepts
+PyTorch tensors directly, and ``torch.testing.assert_close`` checks its
+output.
+
+The same ``add_one`` object lets us inspect the program's IR. We can get
+the type of parameter ``A`` through ``params[0].ty``, then read its symbolic
+length from ``shape[0]``:
+
+.. code-block:: python
+
+ input_type = add_one.params[0].ty
+ length = input_type.shape[0]
+
+The :doc:`IR Walkthrough </ir/overview>` explains these IR objects and how
+to work with them.
+
+We can also run this computation on an NVIDIA GPU. In addition to LLVM,
+it still needs CUDA support, an NVIDIA driver, the CUDA toolkit, and
+``cuda-bindings`` to run.
+
+.. literalinclude:: examples/add_one_cuda.py
+ :language: python
+ :start-after: # begin program
+ :end-before: # end program
+
+``T.device_entry`` marks the GPU computation and configures blocks of 256
+threads. Each thread handles one element, so the block count is rounded
+up to cover ``n`` elements. The ``if i < n`` check keeps extra threads in
+the last block from accessing the tensors. We compile the function for
+CUDA with ``tirx.compile`` and run it with input and output tensors on
+``"cuda:0"``, then check the result with PyTorch as in the CPU example.
+
+The complete examples are available as
+:download:`add_one_cpu.py <examples/add_one_cpu.py>` and
+:download:`add_one_cuda.py <examples/add_one_cuda.py>`. You can run them with
+``python add_one_cpu.py`` and ``python add_one_cuda.py``, respectively.
diff --git a/docs/index.rst b/docs/index.rst
index d9fe8b1ce..cacf5409d 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -26,7 +26,7 @@ Table of Contents
:caption: Get Started
:titlesonly:
- get_started/example
+ get_started/quick_start
.. toctree::
:maxdepth: 1
diff --git a/docs/ir/examples/sigmoid.py b/docs/ir/examples/sigmoid.py
deleted file mode 100644
index cf1157446..000000000
--- a/docs/ir/examples/sigmoid.py
+++ /dev/null
@@ -1,217 +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.
-"""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/tvmscript_sigmoid.py
b/docs/ir/examples/tvmscript_sigmoid.py
deleted file mode 100644
index 49205c30e..000000000
--- a/docs/ir/examples/tvmscript_sigmoid.py
+++ /dev/null
@@ -1,87 +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.
-"""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/functions_and_modules.rst
b/docs/ir/functions_and_modules.rst
index fb48fc27f..2de69dfcd 100644
--- a/docs/ir/functions_and_modules.rst
+++ b/docs/ir/functions_and_modules.rst
@@ -22,34 +22,23 @@ 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.
+implementation and locate definitions through global references.
Functions
---------
-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.
+:class:`ir.BaseFunc <tvm_tirx.ir.BaseFunc>` is the common base class for
+IR function definitions. It provides the ``attrs`` field for named
+function properties. :class:`tirx.Function <tvm_tirx.tirx.Function>` is
+its subclass for tensor programs, adding parameters, a return type, and
+a statement body:
-``ir.BaseFunc`` is the base class for all function definitions and is also
-an expression. Its ``attrs`` field holds named function properties:
+.. code-block:: text
-.. list-table::
- :header-rows: 1
- :widths: 15 25 60
-
- * - Field
- - Value type
- - Meaning
- * - ``attrs``
- - ``ir.DictAttrs``
- - Named function properties.
+ ir.BaseFunc attrs
+ └── tirx.Function params, ret_type, body
-:class:`tirx.Function <tvm_tirx.tirx.Function>` extends ``ir.BaseFunc`` with
-parameters, a result type, and a statement body:
+The main fields of ``tirx.Function`` are:
.. list-table::
:header-rows: 1
@@ -60,82 +49,163 @@ parameters, a result type, and a statement body:
* - ``params``
- Parameter variables in calling order.
* - ``ret_type``
- - The declared result type.
+ - The declared return type.
* - ``body``
- - A ``SeqStmt``, or None for a declaration without an implementation.
-
-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``.
+ - A ``SeqStmt``, or ``None`` for a declaration without an implementation.
+ * - ``attrs``
+ - An ``ir.DictAttrs`` containing named function properties.
-The program defines a function with input, bias, and output tensors.
-Its IR definition combines these three parameters with the batch and
-channel loops:
+In TVMScript, ``@T.function`` produces this IR object from a function
+definition. The CPU example is defined as:
-.. literalinclude:: examples/sigmoid.py
+.. literalinclude:: ../get_started/examples/add_one_cpu.py
:language: python
- :start-after: # begin function construction
- :end-before: # end function construction
- :dedent: 4
+ :start-at: from tvm_tirx.script import tirx as T
+ :end-at: B[i] = A[i] + 1
+
+The name ``add_one`` refers to the resulting ``tirx.Function``. Its
+parameters are the variables ``A`` and ``B``, and its body contains the
+loop. The loop has its own body containing the tensor store:
+
+.. code-block:: text
+
+ add_one: tirx.Function
+ ├── params[0]: ir.Var A
+ ├── params[1]: ir.Var B
+ ├── ret_type: ir.TupleType no return value
+ └── body: ir.SeqStmt
+ └── seq[0]: ir.For for i in range(n)
+ └── body: ir.SeqStmt
+ └── seq[0]: ir.TensorStore B[i] = A[i] + 1
+
+``add_one.ty`` stores an ``ir.FuncType`` describing the function's
+signature: its parameter types and return type. We can inspect these
+types directly:
+
+.. code-block:: python
+
+ signature = add_one.ty
+ parameter_types = signature.arg_types
+ return_type = signature.ret_type
+
+``parameter_types`` contains the tensor types of ``A`` and ``B`` in
+calling order. The entries in ``add_one.params`` are the variables
+themselves. ``return_type`` is the same type stored in
+``add_one.ret_type``: an empty tuple type, ``ir.TupleType([])``, indicating
+no return value. This function writes its results into ``B``.
+
+A function can also return a value explicitly. The ``add_one_scalar``
+helper from :doc:`operators_and_calls` returns a single float32 value:
+
+.. code-block:: python
+
+ @T.function(private=True)
+ def add_one_scalar(value: T.float32) -> T.float32:
+ return value + 1
+
+The return annotation ``-> T.float32`` sets its ``ret_type`` to
+``ir.PrimType("float32")``, and the ``return`` becomes an ``ir.Return``
+statement in its body. The result of ``add_one`` goes into an output
+parameter; the result of ``add_one_scalar`` is a return value.
+
+So far, ``@T.function`` has built a function object for us from source
+code. Compiler code can also build this kind of object directly with
+``tirx.Function``. To focus on the function constructor, we can reuse
+the parameters and loop already stored in ``add_one``:
-``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.
+.. code-block:: python
+
+ from tvm_tirx import ir, tirx
+
+ A, B = add_one.params
+ loop = add_one.body[0]
+
+``A`` and ``B`` are the IR parameter variables, and ``loop`` is the
+``ir.For`` shown in the tree above. We can pass these objects to the
+constructor to create another function with the same computation:
+
+.. code-block:: python
+
+ ir_function = tirx.Function([A, B], loop, ret_type=ir.TupleType([]))
+
+The new function takes ``A`` and ``B`` as parameters and runs ``loop``
+as its body. The constructor wraps this single statement in a ``SeqStmt``.
+``ret_type=ir.TupleType([])`` specifies that the function returns no value;
+its loop writes the result to ``B``. A function that returns a value needs
+its return type supplied through ``ret_type``; the constructor does not
+infer it from the body.
+
+We will put ``ir_function`` in a module and compile it below. To make the
+compiled function available under the name ``"add_one"``, we set its
+``global_symbol`` attribute:
+
+.. code-block:: python
+
+ ir_function = ir_function.with_attr("global_symbol", "add_one")
+
+``with_attr``, inherited from ``ir.BaseFunc``, returns a new function with
+this entry in ``attrs``. For public functions written in TVMScript,
+``@T.function`` sets
+``global_symbol`` to the Python function name automatically.
Modules
-------
-: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``.
+A program can contain several functions that call one another.
+:class:`ir.IRModule <tvm_tirx.ir.IRModule>` collects their definitions so
+compiler passes can find the functions being called and transform the
+program as a whole. We can put the function we just constructed,
+``ir_function``, in a module under the name ``"add_one"``:
-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:
+.. code-block:: python
-.. literalinclude:: examples/sigmoid.py
- :language: python
- :start-after: # begin module construction
- :end-before: # end module construction
+ module = ir.IRModule({"add_one": ir_function})
-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.
+The module's ``functions`` field maps global references to function
+definitions. Each key is an :class:`ir.GlobalVar <tvm_tirx.ir.GlobalVar>`,
+and each value is a function definition whose class inherits
+``ir.BaseFunc``. Here, that value is our ``tirx.Function`` object,
+``ir_function``. The constructor creates a ``GlobalVar`` named
+``"add_one"`` from the string key and associates it with this function:
-Use ``module.get_global_var(name)`` to retrieve the ``ir.GlobalVar`` reference
-for an existing function:
+.. code-block:: text
-.. literalinclude:: examples/sigmoid.py
- :language: python
- :start-after: # begin module lookup
- :end-before: # end module lookup
+ module: ir.IRModule
+ └── functions
+ └── key: ir.GlobalVar ("add_one") → value: tirx.Function (ir_function)
-``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.
+The ``GlobalVar`` serves as a reference; the function body belongs to
+the associated ``tirx.Function``. We can retrieve the reference by name,
+then use it to look up the definition:
-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:
+.. code-block:: python
-.. literalinclude:: examples/sigmoid.py
- :language: python
- :start-after: # begin module update
- :end-before: # end module update
-
-``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.
+ add_one_ref = module.get_global_var("add_one")
+ function = module[add_one_ref]
+
+The first line returns the ``GlobalVar`` key, and the second returns
+the ``tirx.Function`` stored under it. ``module["add_one"]`` retrieves
+the same definition directly by name. A call can use ``add_one_ref`` as
+its ``op`` to identify this function, as in the function-call example in
+:doc:`operators_and_calls`.
+
+The name ``"add_one"`` used in this lookup is the function's name within
+the module. The function's ``attrs["global_symbol"]`` specifies its
+exported name for code generation. Both names are ``"add_one"`` in this
+example, but they serve these separate purposes.
+
+``tirx.compile`` accepts a module as well as a standalone function. We
+can compile ``module`` with the same LLVM target used in Quick Start,
+then create CPU tensors and call its exported ``add_one`` function:
+
+.. code-block:: python
+
+ import torch
+
+ compiled = tirx.compile(module, target="llvm")
+ a = torch.arange(10, dtype=torch.float32)
+ b = torch.empty(a.shape, dtype=a.dtype)
+ compiled["add_one"](a, b)
+
+``compiled["add_one"]`` retrieves the exported function. The call reads
+``a`` and writes ``a + 1`` into ``b``, just as in Quick Start.
diff --git a/docs/ir/operators_and_calls.rst b/docs/ir/operators_and_calls.rst
index b38230080..bfb0ca656 100644
--- a/docs/ir/operators_and_calls.rst
+++ b/docs/ir/operators_and_calls.rst
@@ -18,44 +18,52 @@
Operators and Calls
===================
-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.
+In this chapter, we will introduce operators and calls. A registered
+operator provides the shared argument requirements, type rules, and effect
+information for a computation. A call refers to an operator or function
+and supplies argument expressions for that use. We will examine how
+the IR records these definitions and calls, infers result types, and
+checks arguments.
Operators
---------
-``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.
+An operator identifies a computation, such as sigmoid or a CUDA
+thread-index lookup. :class:`ir.Op <tvm_tirx.ir.Op>` represents a
+registered operator and records its name, argument descriptions, and
+attributes. The operator's registration also provides rules for
+result-type inference and argument validation.
-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.
+The GPU example uses the registered operator ``tirx.cuda.thread_idx`` to
+obtain a thread coordinate. ``ir.Op.get`` retrieves its ``ir.Op`` object,
+and ``op.args_info`` describes its parameters. We can extract their names:
-The program uses the same exponential operator for each input element.
-Its shared argument and effect information can be inspected through the
operator:
+.. code-block:: python
-.. literalinclude:: examples/sigmoid.py
- :language: python
- :start-after: # begin operator lookup
- :end-before: # end operator lookup
+ from tvm_tirx import ir
+ op = ir.Op.get("tirx.cuda.thread_idx")
+ argument_names = [argument.name for argument in op.args_info]
+ print(argument_names)
+
+The output is:
+
+.. code-block:: text
-``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.
+ ['axis']
+
+The list contains one parameter name, ``axis``. In the GPU example,
+``T.cuda.thread_idx("x")`` passes the string ``"x"`` for this parameter.
+This selects the x axis, so the call reads the current thread's x-axis
+coordinate within its block.
Calls
-----
-: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:
+:class:`ir.Call <tvm_tirx.ir.Call>` represents one use of an operator or
+function with specific argument expressions. Its ``op`` field references
+the expression being called, known as the callee, and ``args`` holds the
+arguments for this use. The main fields describe what is called, what
+is passed to it, and what type of value it produces:
.. list-table::
:header-rows: 1
@@ -68,79 +76,346 @@ callee, its arguments, and any attributes for the
invocation:
* - ``args``
- Value arguments in calling order.
* - ``ty_args``
- - Explicit type arguments, empty by default.
+ - Types supplied as arguments to the callee, empty by default.
* - ``attrs``
- - Attributes for this invocation, or None. A dictionary becomes
- ``ir.DictAttrs``.
+ - Operator-specific configuration or annotations for this call,
+ stored as an ``Attrs`` object, or ``None``.
+ * - ``ty``
+ - The result type of this call, inferred or explicitly supplied.
-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:
+The thread-index call refers to the registered ``tirx.cuda.thread_idx``
+operator. We can construct it directly through the IR API:
.. code-block:: python
- x = ir.Var("x", "float32")
- call = ir.prim.sigmoid(x)
+ op = ir.Op.get("tirx.cuda.thread_idx")
+ x_call = ir.Call(op, [ir.StringImm("x")])
-: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.
+``ir.Op.get`` retrieves the shared operator definition. ``ir.Call``
+creates an expression whose ``op`` refers to that definition and whose
+``args`` contains the axis argument. Its ``ty`` records the int32 result:
-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`.
+.. code-block:: text
-.. _types-and-validation:
+ x_call: ir.Call
+ ├── op → shared ir.Op tirx.cuda.thread_idx
+ ├── args[0]: ir.StringImm "x"
+ └── ty: ir.PrimType int32
-Type Inference and Validation
------------------------------
+``op.args_info[0].name`` is the parameter name ``"axis"``;
+``x_call.args[0]`` supplies its argument value ``"x"``. The call describes
+reading the current thread's x-axis coordinate when the compiled GPU
+program runs.
-``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:
+For this registered operator, ``ir.Call`` infers ``ty`` from the
+operator's type rule. An explicit ``ty`` supplies the result type directly
+and skips inference.
+
+The constructor also accepts the registered name directly:
+``ir.Call("tirx.cuda.thread_idx", [ir.StringImm("x")])``.
+
+The program uses the x axis. As an additional example, we can construct
+a call for the y axis with the same operator:
.. 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.
+ y_call = ir.Call(op, [ir.StringImm("y")])
+
+Both ``x_call.op`` and ``y_call.op`` reference the shared ``op``. Each
+call has its own ``args`` array, containing ``"x"`` or ``"y"``, and
+both calls produce an int32 coordinate:
-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 :meth:`call.validate()
<tvm_tirx.ir.Call.validate>`.
-Type inference can also reject inputs while computing the result type.
+.. code-block:: text
-To see how validation catches missing arguments, we can construct a
-``prim.exp`` call with no input and an explicit result type:
+ x_call.op ──┐
+ ├──> op: ir.Op tirx.cuda.thread_idx
+ y_call.op ──┘
-.. literalinclude:: examples/sigmoid.py
+A call's result can be an operand in a larger expression. In the GPU
+example, two calls contribute to the element index ``i``:
+
+.. literalinclude:: ../get_started/examples/add_one_cuda.py
:language: python
- :start-after: # begin validation failure
- :end-before: # end validation failure
- :append: invalid.validate()
+ :start-at: i = T.cuda.block_idx
+ :end-at: i = T.cuda.block_idx
+ :dedent: 4
-The explicit result type bypasses type inference. ``invalid.validate()`` then
-raises ``TypeError`` because ``prim.exp`` requires one input argument.
+``T.cuda.block_idx("x")`` supplies the block coordinate, and
+``T.cuda.thread_idx("x")`` supplies the thread's coordinate within that
+block. Each helper constructs an ``ir.Call``. The multiplication and
+addition use the calls' results as operands:
-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:
+.. code-block:: text
-.. literalinclude:: examples/sigmoid.py
- :language: python
- :start-after: # begin provisional call
- :end-before: # end provisional call
-
-``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.
+ ir.prim.Add
+ ├── a: ir.prim.Mul
+ │ ├── a: ir.Call T.cuda.block_idx("x")
+ │ └── b: ir.IntImm 256
+ └── b: ir.Call T.cuda.thread_idx("x")
+
+Expression Arguments
+~~~~~~~~~~~~~~~~~~~~
+
+A call can also receive a computation as an argument.
+As an extension of the scalar add-one expression, we can apply sigmoid to
+``scalar_value + 1``. The variable ``scalar_value`` has type
+``ir.PrimType("float32")``:
+
+.. code-block:: python
+
+ scalar_value = ir.Var("scalar_value", "float32")
+ scalar_plus_one = scalar_value + 1
+ activated_plus_one = ir.prim.sigmoid(scalar_plus_one)
+
+``ir.prim.sigmoid(...)`` constructs and returns an ``ir.Call`` object.
+``activated_plus_one`` refers to this call, which is also a primitive
+expression because its result type is ``ir.PrimType("float32")``.
+We can construct the same call explicitly:
+
+.. code-block:: python
+
+ sigmoid_op = ir.Op.get("prim.sigmoid")
+ activated_plus_one = ir.Call(sigmoid_op, [scalar_plus_one])
+
+In both forms, the call's ``op`` refers to the registered ``prim.sigmoid``
+operator, and ``args[0]`` refers to the addition expression
``scalar_plus_one``:
+
+.. code-block:: text
+
+ activated_plus_one: ir.Call
+ ├── op → shared ir.Op prim.sigmoid
+ ├── args[0]: ir.prim.Add scalar_value + 1
+ │ ├── a: ir.Var scalar_value
+ │ └── b: ir.FloatImm 1.0
+ └── ty: ir.PrimType float32
+
+When the compiled program runs, sigmoid receives the result of the
+addition. This expression therefore computes
+``1 / (1 + exp(-(scalar_value + 1)))`` and produces a float32 value.
+
+Function Calls
+--------------
+
+The callee can also be a function defined in IR. A function stored in an
+``IRModule`` is referenced through an ``ir.GlobalVar``. For example, a
+small helper can put the scalar add-one computation in a function:
+
+.. code-block:: python
+
+ from tvm_tirx.script import tirx as T
+
+ @T.function(private=True)
+ def add_one_scalar(value: T.float32) -> T.float32:
+ return value + 1
+
+``private=True`` makes this an internal helper. We can place it in a
+module and construct a call that passes ``scalar_value`` as its argument:
+
+.. code-block:: python
+
+ call_module = ir.IRModule({"add_one_scalar": add_one_scalar})
+ callee = call_module.get_global_var("add_one_scalar")
+ function_call = ir.Call(callee, [scalar_value], ty=add_one_scalar.ret_type)
+
+``callee`` is the module's global reference to ``add_one_scalar``.
+``function_call.op`` refers to this ``GlobalVar``, and
+``function_call.args[0]`` supplies the value for the helper's parameter
+``value``. A function-valued callee with a ``FuncType`` supplies a return
+type for inference. Here, ``callee.ty`` is ``MissingType``: the reference
+itself does not carry the function's signature. We therefore supply the
+return type from the definition through ``ty``. The call holds the
+reference and argument; the module holds the function's definition.
+Other function-valued IR expressions can also serve as callees.
+
+Calls in Function Bodies
+------------------------
+
+A call contributes to a function's computation through the expression or
+statement that uses it. We can extend the CPU example to apply sigmoid
+to each add-one result:
+
+.. code-block:: python
+
+ n = T.dynamic("n")
+
+ @T.function
+ def add_one_sigmoid(A: T.Tensor((n,), T.float32), B: T.Tensor((n,),
T.float32)):
+ for i in range(n):
+ B[i] = T.sigmoid(A[i] + 1)
+
+The TVMScript helper ``T.sigmoid`` builds a call to ``prim.sigmoid``.
+The tensor store uses that call as its ``value`` expression:
+
+.. code-block:: python
+
+ activation_store = add_one_sigmoid.body[0].body[0]
+ activation_call = activation_store.value
+
+.. code-block:: text
+
+ activation_store: ir.TensorStore
+ ├── dest: ir.Var B
+ └── value: ir.Call prim.sigmoid
+ └── args[0]: ir.prim.Add A[i] + 1
+
+When the function runs, the addition supplies the sigmoid input and the
+store writes its result to ``B[i]``.
+
+Calls for Their Effects
+~~~~~~~~~~~~~~~~~~~~~~~
+
+Some registered operators perform a memory write themselves. A common
+example is a masked store, which writes only selected vector lanes. A
+vector version of add-one can use masked reads and writes to process
+four elements at a time, including a final group with fewer than four
+elements:
+
+.. code-block:: python
+
+ @T.function
+ def masked_add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,),
T.float32)):
+ for base in range(0, n, 4):
+ indices = T.Ramp(base, 1, 4)
+ values = T.masked_load(A, indices, indices < n)
+ T.masked_store(B, values + 1, indices, indices < n)
+
+``indices`` contains four consecutive positions starting at ``base``.
+The mask ``indices < n`` enables only positions within the tensor for
+both the read and the write. For example, if ``n = 10``, valid indices
+range from 0 to 9. The final group contains indices 8, 9, 10, and 11, so
+the mask enables only the first two lanes.
+
+``T.masked_store`` builds a call to ``tirx.masked_store``, which performs
+the write and returns ``void``. The TVMScript helper places this call in
+an ``ir.Evaluate`` statement inside the loop body:
+
+.. code-block:: python
+
+ memory_loop = masked_add_one.body[0]
+ write_statement = memory_loop.body[2]
+ write_call = write_statement.value
+
+.. code-block:: text
+
+ memory_loop.body: ir.SeqStmt
+ ├── seq[0]: ir.Bind indices
+ ├── seq[1]: ir.Bind values
+ └── seq[2]: ir.Evaluate
+ └── value: ir.Call tirx.masked_store
+ └── ty: ir.PrimType void
+
+Through the IR API, ``ir.Evaluate(write_call)`` constructs the
+corresponding statement. Even with a ``void`` result type, the call is
+an expression; ``Evaluate`` gives it a place in the statement sequence.
+
+The operators' effect attributes explain why these two calls need
+different treatment when their results are unused:
+
+.. code-block:: python
+
+ activation_effect = activation_call.op.get_attr("TCallEffectKind")
+ write_effect = write_call.op.get_attr("TCallEffectKind")
+
+``activation_effect`` is ``ir.CallEffectKind.Pure``: the sigmoid
+operation has no memory side effects. ``write_effect`` is
+``ir.CallEffectKind.UpdateState``. Compiler passes cannot discard this
+write merely because its result is unused; it updates the output tensor.
+
+Type Arguments and Attributes
+-----------------------------
+
+Some operators need a type as an input to describe the value they
+produce. For example, a tensor allocation needs the shape and element
+type of the storage to create. This additional example describes a
+four-element float32 tensor in local storage:
+
+.. code-block:: python
+
+ from tvm_tirx import tirx
+
+ temporary_type = tirx.TensorType((4,), "float32", scope="local")
+ allocation_call = ir.Call(
+ "tirx.alloc_tensor", [], ty_args=[temporary_type], attrs={}
+ )
+
+``ty_args[0]`` supplies the tensor type to the allocation operator. Its
+type rule uses this argument as the result type, so ``allocation_call.ty``
+is ``temporary_type``. A call can therefore produce a tensor as well as a
+scalar or vector value.
+A surrounding ``ir.Bind`` can name the allocated tensor for subsequent
+reads and writes, as described in :doc:`statements`.
+
+The ``attrs`` argument supplies annotations for an individual call. A
+dictionary becomes an ``ir.DictAttrs`` object; the operator determines
+which annotations it supports and how they are used. These annotations
+belong to this call, while the attributes on ``ir.Op`` describe
+properties shared by every call to that operator.
+
+.. _types-and-validation:
+
+.. _type-inference-and-validation:
+
+Call Types and Validation
+-------------------------
+
+The earlier examples inferred a call's result type or supplied it
+explicitly. During IR construction, we may also need a call whose result
+type is not known yet. ``ty=ir.Type.missing()`` records this state and
+defers inference. The same scalar sigmoid input, ``scalar_plus_one``,
+lets us compare all three choices:
+
+.. code-block:: python
+
+ sigmoid_op = ir.Op.get("prim.sigmoid")
+ inferred_call = ir.Call(sigmoid_op, [scalar_plus_one])
+ typed_call = ir.Call(sigmoid_op, [scalar_plus_one], ty="float32")
+ untyped_call = ir.Call(sigmoid_op, [scalar_plus_one], ty=ir.Type.missing())
+
+``inferred_call`` and ``typed_call`` both store a float32 result type.
+``untyped_call`` explicitly supplies ``MissingType``, so its result type
+remains undetermined.
+
+The same operator also accepts a four-lane float32 vector. We can leave
+this call's result type undetermined, then infer it from the vector
+argument:
+
+.. code-block:: python
+
+ vector_value = ir.Var("vector_value", "float32x4")
+ untyped_sigmoid = ir.Call(sigmoid_op, [vector_value], ty=ir.Type.missing())
+
+With the vector argument in place,
+:func:`ir.reinfer_type <tvm_tirx.ir.reinfer_type>` applies the sigmoid
+operator's type rule and returns the result type. We can use it to
+construct a new call with that type:
+
+.. code-block:: python
+
+ inferred = ir.reinfer_type(untyped_sigmoid)
+ rebuilt_call = ir.Call(sigmoid_op, untyped_sigmoid.args, ty=inferred)
+
+``ir.reinfer_type`` returns ``ir.PrimType("float32x4")`` without changing
+``untyped_sigmoid.ty``, which remains ``MissingType``. ``rebuilt_call``
+stores the returned type.
+
+The sigmoid operator also requires one input argument. ``rebuilt_call``
+passes ``vector_value`` as that input. The following call leaves it out:
+
+.. code-block:: python
+
+ invalid_sigmoid = ir.Call(sigmoid_op, [], ty="float32x4")
+
+``ty="float32x4"`` supplies the result type directly, so the constructor
+skips inference. It still creates the incomplete ``ir.Call`` object
+because argument checks do not run automatically. We can use
+:meth:`Call.validate() <tvm_tirx.ir.Call.validate>` to check whether the
+arguments satisfy the operator's requirements:
+
+.. code-block:: python
+
+ invalid_sigmoid.validate()
+
+``invalid_sigmoid.validate()`` raises ``TypeError`` for the missing input.
+``untyped_sigmoid.validate()`` succeeds because the call supplies
+``vector_value``. Its ``ty`` remains ``MissingType`` because validation
+checks the arguments without inferring the result type.
diff --git a/docs/ir/overview.rst b/docs/ir/overview.rst
index b86bd798d..1d371dfbf 100644
--- a/docs/ir/overview.rst
+++ b/docs/ir/overview.rst
@@ -18,10 +18,19 @@
IR Overview
===========
-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.
+An intermediate representation (IR) describes a program in a form that compiler
+components can inspect and transform. Throughout this walkthrough, we'll use
+the CPU and GPU examples from :doc:`/get_started/quick_start` to introduce
+the main IR objects and show how they fit together.
+
+TVMScript constructs IR objects to describe the computation in these examples.
+PyTorch tensors hold the input and output data, and the compiled functions
+perform the computation when called.
+
+TIRx shares its IR object system across C++, Python, and Rust through TVM FFI,
+so compiler components written in different languages can work together. The
+same infrastructure supports custom types and operators and can be reused to
+build analyses, transformation passes, and backends.
The main components of TIRx IR are:
@@ -37,14 +46,18 @@ The main components of TIRx IR are:
- ``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``
+ * - Statements
+ - ``ir.Stmt``
+ - Describe execution order and control flow.
+ - ``B[i] = A[i] + 1``
+ * - Source locations
+ - ``ir.Location``
+ - Record where IR nodes originate in the source.
+ - The filename and line number of ``A[i] + 1``.
* - Functions
- ``ir.BaseFunc``, ``tirx.Function``
- Define parameters and a computation body.
@@ -53,22 +66,8 @@ The main components of TIRx IR are:
- ``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. 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 detailed inheritance hierarchy of the core IR
classes.
+The hierarchy below shows the inheritance relationships between these classes:
.. code-block:: text
@@ -155,16 +154,20 @@ The following tree shows the detailed inheritance
hierarchy of the core IR class
│ └── ir.CallSiteLoc Callee and caller
locations
└── ir.IRModule Global function
definitions
+The hierarchy shows inheritance between classes. The smaller diagrams in
+the following chapters show selected fields of individual objects. Each
+branch names a field, list entry, or mapping entry; for example,
+``seq[0]`` can refer to an ``ir.Bind`` stored inside an ``ir.SeqStmt``.
+
.. _choose-a-chapter:
What to Read Next
-----------------
-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 CPU example provides the tensor types, arithmetic, and loop; the GPU
+example adds a variable binding, a conditional, and calls for thread
+coordinates. Short additional examples introduce scalar and vector
+values through the IR API.
* The :doc:`types_and_expressions` chapter explains how values, computations,
and tensor shapes are represented.
diff --git a/docs/ir/printing_and_parsing.rst b/docs/ir/printing_and_parsing.rst
index dc9573a40..849647fdb 100644
--- a/docs/ir/printing_and_parsing.rst
+++ b/docs/ir/printing_and_parsing.rst
@@ -20,85 +20,308 @@
Printing and Parsing
====================
-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.
+In this chapter, we will introduce printing and parsing with TVMScript,
+the Python-based syntax used to describe tensor programs in Quick Start.
+The printer turns IR objects into readable TVMScript code, and the parser
+constructs IR objects from that code. We will use these interfaces to
+inspect the add-one function, save and load programs, and examine a
+compiler pass's changes.
-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.
+Printing IR
+-----------
-The example program defines ``bias_sigmoid`` inside a Python factory:
+In the CPU example, the TVMScript decorator ``@T.function`` parses the
+function definition into a ``tirx.Function`` IR object:
-.. literalinclude:: examples/tvmscript_sigmoid.py
+.. literalinclude:: ../get_started/examples/add_one_cpu.py
:language: python
- :start-at: from __future__ import annotations
- :end-at: return bias_sigmoid
+ :start-at: from tvm_tirx.script import tirx as T
+ :end-at: B[i] = A[i] + 1
-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.
+The ``script()`` method returns the function's TVMScript source:
-For printing and parsing, we reuse ``module`` from
-:doc:`functions_and_modules`.
+.. code-block:: python
+
+ source = add_one.script()
+ print(source)
-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:
+The output makes the symbolic dimension and constant's type explicit:
.. code-block:: python
- function = module["bias_sigmoid_130"]
- source = function.script()
- print(source)
+ from __future__ import annotations
-The printed loop structure shows the iteration over samples and channels:
+ # from tvm_tirx.script import tirx as T
-.. code-block:: text
+ @T.function
+ def add_one[n](A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
+ for i in range(n):
+ B[i] = A[i] + T.float32(1.0)
+
+The original code declares ``n`` with ``T.dynamic``. The printer places
+that declaration in ``[n]`` after the function name, so the printed
+definition includes the symbol needed by its tensor shapes and loop
+bound. ``T.float32(1.0)`` records the float32 constant used by the addition.
+The printed form represents the same IR as the original definition.
- for i in range(T.int64(0), T.int64(2) * n):
- for j in range(T.int64(0), T.int64(130)):
- ...
+The same method works on modules and smaller IR objects. For example,
+``add_one.params[0].ty.script()`` returns ``T.Tensor((n,), T.float32)``
+for the input tensor type.
-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.
+.. _parsing-source:
-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.
+Parsing TVMScript
+-----------------
-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:
+The :doc:`TVMScript parser </reference/python/script.parser>` provides
+``from_source`` to construct IR from a string containing TVMScript code.
+We can use it with the add-one definition:
.. code-block:: python
from tvm_tirx.script import from_source
+ source_text = """
+ from tvm_tirx.script import tirx as T
+
+ @T.function
+ def add_one[n](A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
+ for i in range(n):
+ B[i] = A[i] + T.float32(1.0)
+ """
+ parsed_function = from_source(source_text)
+ print(type(parsed_function).__name__)
+ print(parsed_function.script())
+
+The output is:
+
+.. code-block:: text
+
+ Function
+ from __future__ import annotations
+
+ # from tvm_tirx.script import tirx as T
+
+ @T.function
+ def add_one[n](A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
+ for i in range(n):
+ B[i] = A[i] + T.float32(1.0)
+
+``parsed_function`` is a ``tirx.Function`` containing the parameter
+variables, loop, and tensor store. It can be inspected or compiled
+through the same interfaces as ``add_one``.
+
+The import inside ``source_text`` tells the parser what ``T`` refers to.
+The earlier ``source`` returned by ``add_one.script()`` has this import
+commented out. In that case, ``extra_vars`` supplies the existing ``T``
+namespace:
+
+.. code-block:: python
+
+ from tvm_tirx import ir
+ from tvm_tirx.script import tirx as T
+
restored_function = from_source(source, extra_vars={"T": T})
- ir.assert_structural_equal(function, restored_function)
+ ir.assert_structural_equal(add_one, restored_function)
+
+The printed ``[n]`` declares the symbolic dimension, so it needs no
+external binding.
+
+Saving and Loading a Module
+---------------------------
+
+An IR module can be saved as TVMScript for later inspection or reuse.
+We can store the CPU function in a module and write its source to a file.
+To let that file supply its own namespaces when parsed, we include the
+imports with ``ir.comment_imports=False``:
+
+.. code-block:: python
+
+ from pathlib import Path
+
+ module = ir.IRModule({"add_one": add_one})
+ module_source = module.script(extra_config={"ir.comment_imports": False})
+ output_path = Path("add_one_module.py")
+ output_path.write_text(module_source)
+
+We can read the saved source and reconstruct the module:
+
+.. code-block:: python
+
+ restored_module = from_source(output_path.read_text())
+ ir.assert_structural_equal(module, restored_module)
+
+The imports in the file supply the namespaces needed by the parser, so
+this call needs no ``extra_vars``. ``restored_module`` is an ``ir.IRModule``
+with the same function definitions as ``module``.
+
+Inspecting IR
+-------------
+
+The printer can also help us locate a selected computation in the source
+or see how a compiler pass changes the program.
+
+.. _highlighting-nodes:
+
+Highlighting IR Objects
+~~~~~~~~~~~~~~~~~~~~~~~
+
+The function's body lets us reach the addition in ``B[i] = A[i] + 1``.
+To find this expression in the printed program, we can ask the printer to
+underline it:
+
+.. code-block:: python
+
+ loop = add_one.body[0]
+ store = loop.body[0]
+ addition = store.value
+ marked_source = add_one.script(obj_to_underline=[addition])
+ print(marked_source)
+
+The marked lines include:
+
+.. code-block:: text
+
+ B[i] = A[i] + T.float32(1.0)
+ ^^^^^^^^^^^^^^^^^^^^^
-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``.
+``obj_to_underline`` selects the IR objects to mark. For a diagnostic,
+``print_line_numbers=True`` adds line numbers, ``obj_to_annotate`` can
+attach messages to selected objects, and ``num_context_lines`` limits the
+surrounding source.
-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:
+For interactive inspection,
+:meth:`add_one.show() <tvm_tirx.tirx.Function.show>` displays highlighted
+source directly and accepts the same options.
-.. literalinclude:: examples/sigmoid.py
+Inspecting a Transformation
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The printer also helps us check what a compiler pass changes. We can
+extend the CPU example with a bounds check inside the loop:
+
+.. code-block:: python
+
+ @T.function
+ def guarded_add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,),
T.float32)):
+ for i in range(n):
+ if i < n:
+ B[i] = A[i] + 1
+
+The loop already guarantees ``i < n``, so this extra branch is redundant.
+The ``StmtSimplify`` pass can remove it. We can print the module before
+and after applying that pass:
+
+.. code-block:: python
+
+ from tvm_tirx import ir, tirx
+
+ before_simplify = ir.IRModule({"guarded_add_one": guarded_add_one})
+ print(before_simplify.script())
+ after_simplify = tirx.transform.StmtSimplify()(before_simplify)
+ print(after_simplify.script())
+
+The second output contains the simplified loop:
+
+.. code-block:: python
+
+ for i in range(n):
+ B[i] = A[i] + T.float32(1.0)
+
+Building Program Variants
+-------------------------
+
+For functions that differ only in a constant, ``from_source`` can parse
+the same template with different ``extra_vars``. This template replaces
+the add-one constant with ``increment``:
+
+.. code-block:: python
+
+ variant_source = """
+ @T.function
+ def add_constant[n](A: T.Tensor((n,), T.float32), B: T.Tensor((n,),
T.float32)):
+ for i in range(n):
+ B[i] = A[i] + increment
+ """
+
+The same text can produce versions that add one or two:
+
+.. code-block:: python
+
+ add_one_variant = from_source(
+ variant_source, extra_vars={"T": T, "increment": 1.0}
+ )
+ add_two_variant = from_source(
+ variant_source, extra_vars={"T": T, "increment": 2.0}
+ )
+ print(add_two_variant.script())
+
+The parser turns the ``increment`` value supplied through ``extra_vars``
+into a constant. The addition in ``add_two_variant`` is therefore
+``A[i] + T.float32(2.0)``; both functions still take only ``A`` and ``B``
+as parameters.
+
+Source Locations
+~~~~~~~~~~~~~~~~
+
+The ``filename`` argument labels generated source in parser diagnostics.
+The following example removes the required type annotation from ``A``
+in the template:
+
+.. code-block:: python
+
+ invalid_source = variant_source.replace("A: T.Tensor((n,), T.float32)", "A")
+ try:
+ from_source(
+ invalid_source,
+ extra_vars={"T": T, "increment": 1.0},
+ filename="generated_add_constant.py",
+ )
+ except SyntaxError as error:
+ print(error)
+
+The output is:
+
+.. code-block:: text
+
+ Parameter 'A' requires an annotation (generated_add_constant.py, line 3)
+
+The program comes from ``invalid_source``; ``filename`` supplies the
+name shown in the message. Successfully parsed
+nodes also record their positions through the
+:ref:`source-location fields <user-ir-source-locations>` introduced earlier.
+
+.. _tvmscript-region-syntax:
+
+Region Syntax
+-------------
+
+The parser also determines which statements belong to a region. In the
+GPU example, TVMScript recognizes ``T.device_entry`` as a region
+construct. Its standalone form starts the device region:
+
+.. literalinclude:: ../get_started/examples/add_one_cuda.py
:language: python
- :pyobject: save_program
+ :start-at: @T.function
+ :end-at: B[i] = A[i] + 1
+
+For this form, TVMScript takes all following statements in the same
+enclosing block as the region's body. The enclosing block here is the
+function body: the region extends from ``T.device_entry(...)`` to the
+end of the function. It contains the assignment to ``i`` and the entire
+``if`` branch, even though those statements have the same indentation
+as ``T.device_entry(...)``.
+
+An explicit ``with`` block shows the same region through indentation:
+
+.. code-block:: python
+
+ @T.function
+ def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
+ with 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
-``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``.
+Both forms produce the same ``ir.RegionStmt``. The :doc:`statements`
+chapter explains the IR structure inside that region.
diff --git a/docs/ir/statements.rst b/docs/ir/statements.rst
index f5d7c2368..21bcf01b9 100644
--- a/docs/ir/statements.rst
+++ b/docs/ir/statements.rst
@@ -20,13 +20,12 @@
Statements
==========
-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.
+In this chapter, we will introduce statements, which describe the actions
+a function performs and their execution order. A statement can bind an
+intermediate value to a variable, write a result to a tensor, or control
+execution through a branch or loop. These statements form the function's
+body, which compiler passes can inspect and rewrite through the IR
+interfaces.
.. _statements-and-sequences:
@@ -83,78 +82,183 @@ writes to sequences, branches, and loops:
Statement Sequences
-------------------
-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.
+A function body often contains several statements that run one after
+another. ``ir.SeqStmt`` stores these statements in execution order in
+its ``seq`` field. The statements inside a loop or branch are also
+stored in a ``SeqStmt``.
-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 CPU example writes one output element in each loop iteration:
+
+.. literalinclude:: ../get_started/examples/add_one_cpu.py
+ :language: python
+ :start-at: for i in range(n):
+ :end-at: B[i] = A[i] + 1
+ :dedent: 4
+
+The indented assignment ``B[i] = A[i] + 1`` becomes an
+``ir.TensorStore`` statement describing a write to a tensor element.
+The loop body is a ``SeqStmt`` containing this statement. We can also
+construct this body through the IR API, using the parameter variables
+from ``add_one``:
+
+.. code-block:: python
+
+ from tvm_tirx import ir
+
+ A, B = add_one.params
+ i = ir.Var("i", "int32")
+ store = ir.TensorStore(B, [i], A[i] + 1)
+ body = ir.SeqStmt([store])
+
+``ir.TensorStore`` takes the destination tensor ``B``, the index list
+``[i]``, and the value expression ``A[i] + 1``. These arguments populate
+its ``dest``, ``indices``, and ``value`` fields. ``ir.SeqStmt([store])``
+holds this store statement as its first entry:
+
+.. code-block:: text
+
+ body: ir.SeqStmt
+ └── seq[0]: ir.TensorStore B[i] = A[i] + 1
+ ├── dest: ir.Var B
+ ├── indices[0]: ir.Var i
+ └── value: ir.prim.Add A[i] + 1
+
+When the compiled function runs, the store computes ``A[i] + 1`` and
+writes the result to element ``i`` of ``B``. We can retrieve this
+statement with ``body[0]``, which is shorthand for ``body.seq[0]``.
+
+A small extension of the CPU loop names the addition's result before
+writing it to the output:
+
+.. code-block:: python
+
+ for i in range(n):
+ value = A[i] + 1
+ B[i] = value
+
+The assignment to ``value`` becomes an ``ir.Bind`` statement: its
+``var`` field holds the variable, and its ``value`` field holds the
+addition expression. The following ``TensorStore`` uses that variable
+as the value to write. We can construct these two statements and put
+them in a sequence:
+
+.. code-block:: python
+
+ value = ir.Var("value", "float32")
+ binding = ir.Bind(value, A[i] + 1)
+ store_from_value = ir.TensorStore(B, [i], value)
+ sequence = ir.SeqStmt([binding, store_from_value])
+
+.. code-block:: text
+
+ sequence: ir.SeqStmt
+ ├── seq[0]: ir.Bind value = A[i] + 1
+ │ ├── var: ir.Var value
+ │ └── value: ir.prim.Add A[i] + 1
+ └── seq[1]: ir.TensorStore B[i] = value
+ └── value → same ir.Var value
+
+When the function runs, the binding computes the addition and defines
+``value``, then the store writes this value to ``B[i]``.
+The next section places the original single-store ``body`` inside the
+CPU loop.
+
+.. _control-flow-and-scope:
+
+Control Flow
+------------
+
+Control-flow statements determine which parts of a program execute and
+how often. A loop repeats a body of statements, while a branch selects
+which body runs.
+
+The CPU example visits every element with a loop:
+
+.. literalinclude:: ../get_started/examples/add_one_cpu.py
+ :language: python
+ :start-at: for i in range(n):
+ :end-at: B[i] = A[i] + 1
+ :dedent: 4
+
+The loop becomes an ``ir.For``. Its ``var`` field identifies the loop
+variable ``i``, and its ``body`` field holds the ``SeqStmt`` containing
+the tensor store. We can construct this loop using ``i`` and ``body``
+from the previous section:
.. code-block:: python
- biased = ir.Var("biased", "float32")
- biased_value = X[i, j] + bias[j]
- result = sigmoid(biased)
-
- body = ir.SeqStmt(
- [
- ir.Bind(biased, biased_value),
- ir.TensorStore(Y, [i, j], result),
- ]
- )
-
-``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 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.
-
-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.
-
-The program processes one channel at a time. The surrounding loops cover
-samples and channels, keeping each tensor access within its bounds:
-
-.. literalinclude:: examples/sigmoid.py
+ loop = ir.For(i, 0, n, 1, ir.ForKind.DEFAULT, body)
+
+The arguments ``0``, ``n``, and ``1`` set the ``begin``, ``end``, and
+``step`` fields. ``ForKind.DEFAULT`` selects execution in iteration
+order. On each iteration, the loop executes the statements in ``body``:
+
+.. code-block:: text
+
+ loop: ir.For
+ ├── var: ir.Var i
+ ├── begin: ir.IntImm 0
+ ├── end: ir.Var n
+ ├── step: ir.IntImm 1
+ ├── kind: ir.ForKind.DEFAULT
+ └── body: ir.SeqStmt
+ └── seq[0]: ir.TensorStore B[i] = A[i] + 1
+
+The end bound is exclusive, so ``i`` runs from 0 through ``n - 1``.
+Each iteration executes the store in ``body`` with that value of ``i``.
+The loop defines ``i`` for use inside ``For.body``, which is its scope.
+The ``For`` establishes this scope, and its ``SeqStmt`` body records
+the order of the statements within it. The
+:class:`~tvm_tirx.ir.ForKind` API reference describes the other execution
+kinds.
+
+The GPU example distributes this work across threads. ``T.device_entry``
+introduces the GPU computation. Each thread computes an index ``i`` and
+checks whether it has an element to process:
+
+.. literalinclude:: ../get_started/examples/add_one_cuda.py
:language: python
- :start-after: # begin loop
- :end-before: # end loop
+ :start-at: T.device_entry(
+ :end-at: B[i] = A[i] + 1
: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.
+The ``if i < n`` becomes an ``ir.If``. Its ``condition`` field holds
+the comparison, an ``ir.prim.LT`` expression. Its ``then_case`` field
+holds a ``SeqStmt`` containing the tensor store:
+
+.. code-block:: text
+
+ ir.If
+ ├── condition: ir.prim.LT i < n
+ ├── then_case: ir.SeqStmt
+ │ └── seq[0]: ir.TensorStore B[i] = A[i] + 1
+ └── else_case: None
+
+The condition produces a boolean. When it is true, the thread executes
+the store in ``then_case``. An else branch has its own ``SeqStmt`` in
+``else_case``. This example has no else branch, so ``else_case`` is
+``None`` and threads whose indices are outside the tensor perform no
+tensor access. These branch bodies can also contain further loops and
+branches.
+
+The standalone ``T.device_entry(...)`` groups the remaining statements
+in this function into a device region. An
+:class:`ir.RegionStmt <tvm_tirx.ir.RegionStmt>` represents this region in
+the IR. Its ``op`` identifies ``tirx.device_entry``, and its ``body`` is
+a ``SeqStmt`` containing the index definition followed by the branch:
+
+.. code-block:: text
+
+ ir.RegionStmt
+ ├── op → ir.Op tirx.device_entry
+ └── body: ir.SeqStmt
+ ├── seq[0]: ir.Bind i = ...
+ └── seq[1]: ir.If if i < n
+ └── then_case: ir.SeqStmt
+ └── seq[0]: ir.TensorStore B[i] = A[i] + 1
+
+The ``Bind`` defines ``i`` before the ``If`` uses it. The variable is
+available to the condition and tensor store within the region. The
+:ref:`Printing and Parsing <tvmscript-region-syntax>` chapter explains
+the grouping rule for this standalone syntax and the equivalent
+``with`` form.
diff --git a/docs/ir/types_and_expressions.rst
b/docs/ir/types_and_expressions.rst
index 2bc991cba..9a5d5d6c8 100644
--- a/docs/ir/types_and_expressions.rst
+++ b/docs/ir/types_and_expressions.rst
@@ -20,11 +20,11 @@
Types and Expressions
=====================
-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.
+In this chapter, we will introduce types and expressions. Expressions
+represent values and computations, and their types describe properties
+such as dtypes, tensor shapes, and function signatures. The add-one
+examples will show how to inspect these objects and construct expressions
+directly with the IR API.
.. _user-ir-type:
@@ -65,88 +65,88 @@ Next, we will look at two common types in tensor programs
in more detail:
PrimType
~~~~~~~~
-``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.
+:class:`ir.PrimType <tvm_tirx.ir.PrimType>` describes integer,
+floating-point, and boolean values, either as scalars or vectors. Its
+``dtype`` records the element kind, bit width, and number of lanes.
-In the sigmoid program, primitive types describe both shape symbols and
-computed values. The dynamic shape symbol ``n`` is declared as:
+In the program's tensor annotations, ``T.float32`` is an
+``ir.PrimType("float32")`` object.
-.. literalinclude:: examples/tvmscript_sigmoid.py
+The program declares its dynamic length as:
+
+.. literalinclude:: ../get_started/examples/add_one_cpu.py
:language: python
- :start-at: n = T.dynamic
- :end-at: n = T.dynamic
- :dedent: 4
+ :start-at: from tvm_tirx.script import tirx as T
+ :end-at: n = T.dynamic
-where ``n`` is a 64-bit integer symbol with type ``ir.PrimType("int64")``.
+where ``n`` is a 32-bit integer symbol with type ``ir.PrimType("int32")``.
+The loop index ``i`` has the same primitive type.
-The sigmoid calculation uses floating-point values:
+We can express the addition ``A[i] + 1`` directly through the IR API,
+with a scalar variable standing for the input element:
-.. literalinclude:: examples/tvmscript_sigmoid.py
- :language: python
- :start-at: biased =
- :end-at: Y[i, j] =
- :dedent: 16
+.. code-block:: python
+
+ from tvm_tirx import ir
+ scalar_value = ir.Var("scalar_value", "float32")
+ scalar_plus_one = scalar_value + 1
-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.
+``scalar_value`` has type ``ir.PrimType("float32")``, so it represents a single
+32-bit floating-point value. The result ``scalar_plus_one`` has the
+same type, just like ``A[i] + 1`` in the program.
-To create a variable of a specific type, you can call the corresponding
-``PrimType`` instance without arguments:
+.. _user-ir-vector-expression:
+
+``PrimType`` also supports vector values. We can extend the same addition
+to four float32 values at once:
.. code-block:: python
- f32 = ir.PrimType("float32")
- x = f32()
+ vector_value = ir.Var("vector_value", "float32x4")
+ vector_plus_one = vector_value + 1
-``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.
+The type of ``vector_value`` is ``ir.PrimType("float32x4")``. The ``x4`` suffix
+specifies four float32 lanes within one vector value. The addition adds
+one to each lane, and ``vector_plus_one`` has the same vector type.
.. _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.
+:class:`tirx.TensorType <tvm_tirx.tirx.TensorType>` is an IR type object
+that describes a whole tensor. It records the tensor's dimensions and
+the type of each element.
-In the sigmoid program, the input ``X`` is declared as:
+The program declares its input and output as:
-.. literalinclude:: examples/tvmscript_sigmoid.py
+.. literalinclude:: ../get_started/examples/add_one_cpu.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.
+ :start-at: def add_one(
+ :end-at: def add_one(
+
+``T.Tensor((n,), T.float32)`` is the TVMScript notation for a tensor type;
+the parser represents it as a ``tirx.TensorType`` object. Its ``shape``
+field contains one dimension, ``n``, and its ``dtype`` field is
+``ir.PrimType("float32")``:
+
+.. code-block:: text
+
+ tirx.TensorType T.Tensor((n,), T.float32)
+ ├── shape[0]: ir.Var n
+ └── dtype: ir.PrimType float32
+
+The dimension ``n`` is symbolic: it is not fixed to a particular number
+when the function is compiled. Each call determines ``n`` from the length
+of ``A``. Since ``B`` uses the same symbol in its type, its length must
+match ``A``. The loop then uses ``n`` as its bound.
Expressions
-----------
-``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:
+``ir.Expr`` is the base class for expressions, which represent program
+values and computations. All expression classes inherit its ``ty`` field,
+which records the type of value they represent:
.. list-table::
:header-rows: 1
@@ -159,8 +159,34 @@ records its result type:
- ``ir.Type``
- The type of value represented by the expression.
-Expressions include variables, constants, arithmetic, tensor reads, and calls.
-These classes store their values or operands in the following fields:
+We can inspect the types of ``scalar_value`` and ``scalar_plus_one``
+through ``ty``. Both are primitive types, so their ``dtype`` fields give
+the dtype names:
+
+.. code-block:: python
+
+ print(scalar_value.ty.dtype)
+ print(scalar_plus_one.ty.dtype)
+
+The output is:
+
+.. code-block:: text
+
+ float32
+ float32
+
+The variable's type was supplied by the ``"float32"`` argument in
+``ir.Var("scalar_value", "float32")``. The addition's constructor
+determines its result type from its operands; this process is called
+type inference.
+
+An expression under construction can also have an ``ir.MissingType``
+in ``ty`` while its type is undetermined. The
+:ref:`Call Types and Validation <types-and-validation>` section shows
+how a missing result type can be inferred later.
+
+Each expression class adds fields for its value or operands, alongside
+the shared ``ty`` field:
.. list-table::
:header-rows: 1
@@ -191,175 +217,265 @@ These classes store their values or operands in the
following fields:
- A tensor element read
- ``source``, ``indices``
-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`):
+.. _user-ir-const:
-.. literalinclude:: examples/tvmscript_sigmoid.py
- :language: python
- :start-at: function = program["bias_sigmoid"]
- :end-at: location = row_extent.loc
- :dedent: 4
+Primitive Expressions
+~~~~~~~~~~~~~~~~~~~~~
+
+A primitive expression is an expression whose ``ty`` is an ``ir.PrimType``.
+This term groups expressions by their result type. Variables, tensor
+reads, arithmetic nodes, and calls can all belong to this group when
+they produce primitive scalar or vector values.
+
+**Scalar expressions.** The CPU example uses the int32 expressions ``n``
+and ``i`` for the loop bound and index. The tensor read ``A[i]`` and the
+addition ``A[i] + 1`` both produce a single float32 value. In the GPU
+example, ``i < n`` produces a boolean that decides whether a thread
+writes its result. These scalar result types are stored as ``PrimType``
+objects in ``ty``.
+
+``A`` itself represents the whole tensor and has a ``TensorType``.
+The expression ``A[i]`` reads element ``i`` of this tensor. TVMScript
+represents the read as an ``ir.TensorLoad`` node. Its ``source`` field
+refers to ``A``, and ``indices`` contains ``i``, selecting the element
+to read:
+
+.. code-block:: text
+
+ ir.TensorLoad A[i]
+ ├── ty: ir.PrimType float32
+ ├── source: ir.Var A
+ │ └── ty: tirx.TensorType (n,) of float32
+ └── indices[0]: ir.Var i
+
+**Vector expressions.** A primitive expression can also produce a vector
+value. The earlier add-one result ``vector_plus_one`` has type
+``ir.PrimType("float32x4")``, representing four float32 lanes within
+one value.
-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.
+A call whose result type is a ``PrimType`` also belongs to this category.
+For example, the GPU expression ``T.cuda.thread_idx("x")`` is an
+``ir.Call`` with result type ``ir.PrimType("int32")``.
.. _user-ir-prim-operations:
Primitive Operations
~~~~~~~~~~~~~~~~~~~~
-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.
+The previous section grouped expressions by result type. The operations
+below describe the computations that produce those values. The
+``ir.prim`` namespace provides node classes and Python helpers for
+arithmetic, comparisons, boolean logic, casts, and other primitive
+computations:
-TIRx overloads Python operators such as ``+`` and ``/`` to construct IR
-expressions. In our program, we use them as follows:
+.. list-table::
+ :header-rows: 1
+ :widths: 55 45
-.. literalinclude:: examples/tvmscript_sigmoid.py
- :language: python
- :start-at: biased =
- :end-at: Y[i, j] =
- :dedent: 16
+ * - Python expression
+ - IR node constructed
+ * - ``scalar_value + 1``
+ - ``ir.prim.Add``
+ * - ``ir.prim.sigmoid(scalar_value)``
+ - ``ir.Call`` to ``prim.sigmoid``
-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.
+Python operators such as ``+``, ``*``, and ``<`` construct arithmetic
+and comparison nodes. The Python helper ``ir.prim.sigmoid`` constructs
+a call to a registered operator; the returned IR node is an ``ir.Call``.
+The :ref:`Call <user-ir-call>` section below examines its fields.
+The :doc:`ir.prim </reference/python/ir.prim>` API reference lists the
+available classes and construction helpers.
-**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:
+In the CPU example, ``A[i] + 1`` becomes an ``ir.prim.Add``. A binary node
+stores its left operand in ``a`` and its right operand in ``b``:
.. code-block:: text
- ir.prim.Add X[i, j] + bias[j]
- ├── a: ir.TensorLoad X[i, j]
- └── b: ir.TensorLoad bias[j]
+ ir.prim.Add A[i] + 1
+ ├── a: ir.TensorLoad A[i]
+ └── b: ir.FloatImm 1.0
-Each tensor read produces one float32 value, and the addition combines
-the input element with its channel's bias.
+The arithmetic constructor derives the result type from its operands.
+Since ``A[i]`` is float32, it converts the source literal ``1`` to the
+float32 constant shown as ``ir.FloatImm``. The result is also float32.
-**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.
+These constructors also support vector operands. In ``vector_value + 1``,
+the scalar constant is broadcast to all four lanes, producing the
+four-lane expression ``vector_plus_one``.
-.. _user-ir-const:
+The GPU example combines multiplication, addition, and comparison to
+compute and check each thread's element index:
-Primitive 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.
-
-**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:
-
-.. code-block:: python
+.. literalinclude:: ../get_started/examples/add_one_cuda.py
+ :language: python
+ :start-at: i = T.cuda.block_idx
+ :end-at: B[i] = A[i] + 1
+ :dedent: 4
- from tvm_tirx import ir
+The GPU launch uses 256 threads per block. The multiplication by 256
+becomes an ``ir.prim.Mul``, and the sum becomes an ``ir.prim.Add``.
+The condition ``i < n`` is an ``ir.prim.LT`` whose ``a`` and ``b``
+operands are the index and tensor length:
- 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")
+.. code-block:: text
-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.
+ ir.prim.LT i < n
+ ├── a: ir.Var i
+ └── b: ir.Var n
-**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.
+The comparison checks whether the thread has a valid element to process.
.. _user-ir-var:
Var
~~~
-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.
+``ir.Var`` represents a named value with a type. Variables can be
+function parameters, loop indices, or intermediate values defined by
+bindings.
-We can declare the variables in our sigmoid example using the IR API:
+The tensor parameters ``A`` and ``B`` in the program are represented by
+``ir.Var`` objects. The PyTorch tensors ``a`` and ``b`` are passed to these
+parameters when the compiled function is called. In the CPU example, we
+can retrieve the parameter variables from ``add_one.params`` and inspect
+their types through ``ty``:
-.. literalinclude:: examples/sigmoid.py
- :language: python
- :start-after: # begin variables
- :end-before: # end variables
- :dedent: 4
+.. code-block:: python
+
+ A, B = add_one.params
+ input_type = A.ty
+ output_type = B.ty
+
+Each parameter is a ``Var`` with its own ``name`` and ``ty`` fields.
+For ``A``, these are:
+
+.. code-block:: text
-``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.
+ A: ir.Var
+ ├── name: "A"
+ └── ty: tirx.TensorType (n,) of float32
+
+TVMScript creates these parameter variables when it parses the function.
+We can also create variables directly with ``ir.Var(name, ty)``. The
+scalar add-one example in PrimType uses this constructor:
+
+.. code-block:: python
+
+ from tvm_tirx import ir
+ scalar_value = ir.Var("scalar_value", "float32")
+ scalar_plus_one = scalar_value + 1
+
+``ir.Var("scalar_value", "float32")`` creates a variable whose ``name``
+is ``"scalar_value"`` and whose ``ty`` is ``ir.PrimType("float32")``.
.. _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``.
+:class:`ir.Call <tvm_tirx.ir.Call>` represents a call to an operator or
+function. Its ``op`` field refers to the operator or function expression
+being called, and ``args`` holds the argument expressions. A registered
+operator is represented by ``ir.Op``.
-In sigmoid, ``1 / (1 + exp(-x))``, the exponential is a call inside a larger
-arithmetic expression:
+In the GPU example, the index expression contains calls to
+``T.cuda.block_idx("x")`` and ``T.cuda.thread_idx("x")``.
+These TVMScript helpers construct calls to registered CUDA operators.
+We can construct the thread-index call directly through the IR API:
.. code-block:: python
- from tvm_tirx import ir, tirx
- from tvm_tirx.script import from_source
- from tvm_tirx.script import tirx as T
+ from tvm_tirx import ir
+
+ op = ir.Op.get("tirx.cuda.thread_idx")
+ x_call = ir.Call(op, [ir.StringImm("x")])
- exp_op = ir.Op.get("prim.exp")
+``op`` is the ``ir.Op`` object for ``tirx.cuda.thread_idx``. The call's
+``op`` field refers to this object, and its ``args`` array contains the
+string expression ``"x"``. Its ``ty`` is ``ir.PrimType("int32")``:
- def sigmoid(x):
- one = ir.const(1, x.ty)
- exp_neg_x = ir.Call(exp_op, [-x])
- return one / (one + exp_neg_x)
+.. code-block:: text
-``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.
+ x_call: ir.Call
+ ├── op → ir.Op tirx.cuda.thread_idx
+ ├── args[0]: ir.StringImm "x"
+ └── ty: ir.PrimType int32
-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:
+The ``"x"`` argument selects the x axis. When the compiled GPU function
+runs, this call reads the current thread's x-axis coordinate within its
+block. The index arithmetic uses the coordinates returned by the
+thread-index and block-index calls to compute ``i``.
+
+A common extension of add-one is an activation applied to its result.
+We can pass the earlier addition expression ``scalar_plus_one`` to the
+built-in ``prim.sigmoid`` operator:
.. code-block:: python
- biased_value = X[i, j] + bias[j]
- result = sigmoid(biased)
+ sigmoid_op = ir.Op.get("prim.sigmoid")
+ activated_plus_one = ir.Call(sigmoid_op, [scalar_plus_one])
+
+The call's argument is now the addition expression:
-``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]``.
+.. code-block:: text
+
+ activated_plus_one: ir.Call
+ ├── op → ir.Op prim.sigmoid
+ ├── args[0]: ir.prim.Add scalar_value + 1
+ └── ty: ir.PrimType float32
+
+``activated_plus_one.args[0]`` refers to ``scalar_plus_one``, so sigmoid
+receives the addition's result. The expression computes
+``1 / (1 + exp(-(scalar_value + 1)))``. The operator's type rule gives
+the call the same type as its input, ``ir.PrimType("float32")``. This
+call is therefore also a primitive expression. The Python helper
+``ir.prim.sigmoid(scalar_plus_one)`` constructs the same call.
+
+The :doc:`operators_and_calls` chapter also covers calls to IR functions
+and explains how registered operators provide result-type inference
+and argument validation.
+
+.. _user-ir-source-locations:
+
+Source Locations
+----------------
+
+IR nodes can carry source locations to help diagnostics point back to
+the original code. The ``loc`` field stores this information, and
+:class:`ir.SourceLoc <tvm_tirx.ir.SourceLoc>` records the source file
+and range.
+
+The parameter ``A`` is declared on this line of the CPU program:
+
+.. literalinclude:: ../get_started/examples/add_one_cpu.py
+ :language: python
+ :start-at: def add_one(
+ :end-at: def add_one(
+ :lineno-match:
+
+The parameter's ``loc`` points back to this declaration. We can read
+``start_line`` to report its line number:
+
+.. code-block:: python
+
+ parameter = add_one.params[0]
+ location = parameter.loc
+ line = location.start_line
+ print(f"Parameter {parameter.name} is declared on line {line} of the CPU
program.")
+
+The output is:
+
+.. code-block:: text
+
+ Parameter A is declared on line 29 of the CPU program.
+
+``source_name.name`` also records the path to the CPU program's source
+file:
+
+.. code-block:: text
-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.
+ parameter: ir.Var A
+ └── loc: ir.SourceLoc
+ ├── source_name: ir.SourceName
+ │ └── name path to the CPU program's source file
+ └── start_line line of A's declaration
