This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 667058e89999029b409f5712649c0d30b1ab8134 Author: Tianqi Chen <[email protected]> AuthorDate: Mon Oct 19 17:59:32 2026 -0400 [DOCS] Clarify TIRx assignment and allocation receipts (#20800) Clarify parsed assignment and explicit Bind examples after the binding cleanup, and document allocation receipt consumption in programmatic builders. Explain raw Call purity and the scalar allocation receipt distinction; correct the For begin/end summary. This changes only Python docstrings. Validation:43 executable API examples,95 focused builder/SSA/layout tests, Ruff, and strict Sphinx with local-link checks. --- tirx/python/tvm_tirx/ir/expr.py | 3 +- tirx/python/tvm_tirx/ir/stmt.py | 8 +++-- tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py | 40 +++++++++++++++++++++-- tirx/python/tvm_tirx/tirx/transform/transform.py | 11 ++++--- 4 files changed, 50 insertions(+), 12 deletions(-) diff --git a/tirx/python/tvm_tirx/ir/expr.py b/tirx/python/tvm_tirx/ir/expr.py index e646e6da8a..d1f99d6d9d 100644 --- a/tirx/python/tvm_tirx/ir/expr.py +++ b/tirx/python/tvm_tirx/ir/expr.py @@ -700,7 +700,8 @@ class Call(_CallableExprWithOp): or a missing type when no deduction is available. Explicit types, including ``Type.missing()``, are preserved exactly. Inference errors propagate. Construction permits provisional IR; :meth:`validate` checks the operator - contract explicitly after inputs are ready. + contract explicitly after inputs are ready. Constructing a ``Call`` never + emits a builder statement, including calls to allocation operators. Parameters ---------- diff --git a/tirx/python/tvm_tirx/ir/stmt.py b/tirx/python/tvm_tirx/ir/stmt.py index d3cd1cba81..bc7f484bdc 100644 --- a/tirx/python/tvm_tirx/ir/stmt.py +++ b/tirx/python/tvm_tirx/ir/stmt.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Shared statement nodes and sequence helpers.""" +"""Shared statement nodes.""" from collections.abc import Mapping, Sequence from enum import IntEnum @@ -38,7 +38,9 @@ class Bind(Stmt): """Bind node. Bind a variable to a value in the enclosing scope. - Bind has no body field. + Parsed TVMScript normally creates bindings with ordinary assignments. Use + explicit ``T.Bind(var, value)`` when constructing a binding for an existing + variable. Bind has no body field. The bound variable is visible in all subsequent statements within the same enclosing scope (SeqStmt, ForNode.body, etc.). @@ -125,7 +127,7 @@ class ForKind(IntEnum): @tvm_ffi.register_object("ir.For") class For(Stmt): - """Iterate from ``min`` while below ``min + extent``, incrementing by ``step``. + """Iterate from ``begin`` while below ``end``, incrementing by ``step``. Parameters ---------- diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py index b2784a083a..46ed3cc65b 100644 --- a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py +++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py @@ -210,9 +210,43 @@ def alloc_tensor( belong to ``attrs``; ``annotations`` accepts a Python dictionary whose scalar values are normalized to typed IR constants. Under an active builder, emit one binding and return its tensor variable in an ``AlreadyEmitted`` - receipt. Outside a builder return - the allocation Call. Construct ``ir.Call("tirx.alloc_tensor", ...)`` directly - when a pure expression is needed inside a builder. + receipt. Parsed assignments consume this receipt automatically; programmatic + code uses ``T.bind_(receipt)`` or ``receipt.value`` to obtain the variable + without emitting another binding. Outside a builder return the allocation + Call. Construct ``ir.Call("tirx.alloc_tensor", ...)`` directly when a pure + expression is needed inside a builder. + + Tensor allocation conveniences such as ``T.alloc_local`` and + ``T.alloc_shared`` use the same receipt convention. Scalar allocation and + declaration helpers instead return ``AlreadyEmitted[TensorLoad]`` in an + active builder, consumed through the same binding protocol. + + Examples + -------- + Parsed TVMScript uses an ordinary assignment. The allocation creates one + binding, and subsequent tensor accesses use its variable. + + .. code-block:: python + + from tvm_tirx import ir + from tvm_tirx.script import tirx as T + from tvm_tirx.script.ir_builder import IRBuilder + + @T.function(private=True) + def example(): + tensor = T.alloc_tensor(ty_args=[T.Tensor((4,), "int32", scope="local")]) + tensor[0] = 1 + + tensor_type = T.Tensor((4,), "int32", scope="local") + with IRBuilder() as builder: + raw = ir.Call("tirx.alloc_tensor", [], ty_args=[tensor_type]) + receipt = T.alloc_tensor(ty_args=[tensor_type]) + tensor = T.bind_(receipt, name="tensor") + assert tensor.same_as(receipt.value) + binding = builder.get() + assert isinstance(binding, ir.Bind) + assert binding.var.same_as(tensor) + assert not binding.value.same_as(raw) """ if attrs is not None and annotations is not None: raise TypeError("alloc_tensor cannot mix attrs and annotations") diff --git a/tirx/python/tvm_tirx/tirx/transform/transform.py b/tirx/python/tvm_tirx/tirx/transform/transform.py index 2d1135f0a7..59d18a0ab9 100644 --- a/tirx/python/tvm_tirx/tirx/transform/transform.py +++ b/tirx/python/tvm_tirx/tirx/transform/transform.py @@ -732,7 +732,8 @@ def ConvertSSA(): Examples -------- - The input deliberately reuses a variable and disables the SSA check. + The input uses explicit ``T.Bind`` statements to reuse a variable and disables + the SSA check. Ordinary parsed assignments create distinct variable bindings. .. code-block:: python @@ -746,18 +747,18 @@ def ConvertSSA(): class Before: @T.function(private=True) def main(): - T.bind(16, var=v) + T.Bind(v, 16) T.evaluate(v) - T.bind(32, var=v) + T.Bind(v, 32) T.evaluate(v) @I.ir_module class Expected: @T.function(private=True) def main(): - v: T.let[T.int32] = 16 + v = 16 T.evaluate(v) - v_1: T.let[T.int32] = 32 + v_1 = 32 T.evaluate(v_1) def run_pass():
