This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new a5886f7d4a [DOCS] Clarify TIRx public APIs and executable pass
examples (#20797)
a5886f7d4a is described below
commit a5886f7d4aaec9d06d9ab2086c595708b172f07a
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 19 14:19:00 2026 -0400
[DOCS] Clarify TIRx public APIs and executable pass examples (#20797)
Clarify TIRx public IR nodes, constructor fields, and transformation
prerequisites. Add executable Before/Expected TVMScript examples for
individual passes and module-pass decorators, concise pipeline parameter
references, and layout construction examples.
Document static-library loading and type-alias initialization
parameters, with API descriptions matching the current layout-check and
loop-bound contracts.
---
tirx/python/tvm_tirx/backend/cuda/op.py | 83 +-
.../tvm_tirx/backend/cuda/transform/__init__.py | 135 ++-
tirx/python/tvm_tirx/backend/trn/pipeline.py | 16 +-
.../tvm_tirx/backend/trn/transform/__init__.py | 44 +-
.../backend/trn/transform/naive_allocator.py | 30 +
.../backend/trn/transform/private_tensor_alloc.py | 47 +-
tirx/python/tvm_tirx/codegen/executable.py | 18 +-
tirx/python/tvm_tirx/ir/attrs.py | 14 +-
tirx/python/tvm_tirx/ir/base.py | 5 +-
tirx/python/tvm_tirx/ir/expr.py | 97 +-
tirx/python/tvm_tirx/ir/function.py | 13 +-
tirx/python/tvm_tirx/ir/global_info.py | 6 +-
tirx/python/tvm_tirx/ir/location.py | 10 +-
tirx/python/tvm_tirx/ir/module.py | 31 +-
tirx/python/tvm_tirx/ir/op.py | 7 +-
tirx/python/tvm_tirx/ir/prim/expr.py | 66 +-
tirx/python/tvm_tirx/ir/prim/vector_expr.py | 14 +-
tirx/python/tvm_tirx/ir/stmt.py | 48 +-
tirx/python/tvm_tirx/ir/transform/core.py | 167 ++-
tirx/python/tvm_tirx/ir/transform/instrument.py | 64 +-
tirx/python/tvm_tirx/ir/type.py | 22 +-
tirx/python/tvm_tirx/tirx/compilation_pipeline.py | 119 ++-
tirx/python/tvm_tirx/tirx/function.py | 31 +-
tirx/python/tvm_tirx/tirx/index_map.py | 22 +-
tirx/python/tvm_tirx/tirx/layout.py | 94 +-
tirx/python/tvm_tirx/tirx/tensor.py | 32 +-
tirx/python/tvm_tirx/tirx/thread_scope.py | 6 +-
tirx/python/tvm_tirx/tirx/tile/dispatch.py | 23 +-
tirx/python/tvm_tirx/tirx/transform/__init__.py | 7 +-
.../tvm_tirx/tirx/transform/function_pass.py | 74 +-
tirx/python/tvm_tirx/tirx/transform/transform.py | 1112 +++++++++++++++++++-
tirx/python/tvm_tirx/tirx/type.py | 10 +-
32 files changed, 2063 insertions(+), 404 deletions(-)
diff --git a/tirx/python/tvm_tirx/backend/cuda/op.py
b/tirx/python/tvm_tirx/backend/cuda/op.py
index b88aeb16da..dbebcd77eb 100644
--- a/tirx/python/tvm_tirx/backend/cuda/op.py
+++ b/tirx/python/tvm_tirx/backend/cuda/op.py
@@ -547,12 +547,74 @@ _TCGEN05_CTA_GROUP = (1, 2)
@tvm_ffi.register_object("tirx.cuda.TCGen05InstrDescriptorAttrs")
class TCGen05InstrDescriptorAttrs(Attrs):
- """Static options for the dense tcgen05 instruction descriptor."""
+ """Static options carried by :func:`tcgen05_encode_instr_descriptor`.
+
+ Parameters
+ ----------
+ d_dtype : str
+ Element type of the output matrix D.
+ a_dtype : str
+ Element type of the input matrix A.
+ b_dtype : str
+ Element type of the input matrix B.
+ M : int
+ Number of rows in the output tile.
+ N : int
+ Number of columns in the output tile.
+ K : int
+ Reduction dimension of the input tiles.
+ trans_a : bool
+ Whether to transpose A.
+ trans_b : bool
+ Whether to transpose B.
+ n_cta_groups : int
+ Number of CTA groups, 1 (default) or 2.
+ neg_a : bool
+ Whether to negate A; defaults to False.
+ neg_b : bool
+ Whether to negate B; defaults to False.
+ sat_d : bool
+ Whether to saturate D; defaults to False.
+ is_sparse : bool
+ Whether to select sparse MMA; defaults to False.
+ """
@tvm_ffi.register_object("tirx.cuda.TCGen05InstrDescriptorBlockScaledAttrs")
class TCGen05InstrDescriptorBlockScaledAttrs(Attrs):
- """Static options for the block-scaled tcgen05 instruction descriptor."""
+ """Static options carried by
:func:`tcgen05_encode_instr_descriptor_block_scaled`.
+
+ Parameters
+ ----------
+ d_dtype : str
+ Element type of the output matrix D.
+ a_dtype : str
+ Element type of the input matrix A.
+ b_dtype : str
+ Element type of the input matrix B.
+ sfa_dtype : str
+ Element type of the scale factors for A.
+ sfb_dtype : str
+ Element type of the scale factors for B.
+ M : int
+ Number of rows in the output tile.
+ N : int
+ Number of columns in the output tile.
+ K : int
+ Reduction dimension of the input tiles.
+ trans_a : bool
+ Whether to transpose A.
+ trans_b : bool
+ Whether to transpose B.
+ n_cta_groups : int
+ Number of CTA groups, 1 (default) or 2.
+ neg_a : bool
+ Whether to negate A; defaults to False.
+ neg_b : bool
+ Whether to negate B; defaults to False.
+ is_sparse : bool
+ Whether to select sparse MMA; defaults to False.
+ """
_encode_instr_descriptor = _make_op_api(
@@ -1473,7 +1535,22 @@ def nvshmem_barrier_all(*, ty=None, loc: Location =
UNKNOWN_LOC):
@tvm_ffi.register_object("tirx.cuda.TensorMapEncodeTiledAttr")
class TensorMapEncodeTiledAttr(Attrs):
- """Descriptor dtype and fixed options for tiled tensor-map encoding."""
+ """Fixed options for :func:`tensormap_encode_tiled`.
+
+ Parameters
+ ----------
+ descriptor_dtype : str or tvm_ffi.dtype
+ Tensor element dtype encoded in the descriptor.
+ rank : int
+ Number of tensor dimensions. Dimensions, strides and pointers are
+ separate operands of the encoding call, not fields of this object.
+ interleave, swizzle, l2_promotion, oob_fill : int
+ CUDA tensor-map enum values for the corresponding encoding modes.
+ Each defaults to zero.
+ force_cu_dtype : int
+ Explicit CUDA tensor-map dtype enum, or -1 to derive it from
+ ``descriptor_dtype``.
+ """
def __init__(
self,
diff --git a/tirx/python/tvm_tirx/backend/cuda/transform/__init__.py
b/tirx/python/tvm_tirx/backend/cuda/transform/__init__.py
index dd946a328b..ba076230aa 100644
--- a/tirx/python/tvm_tirx/backend/cuda/transform/__init__.py
+++ b/tirx/python/tvm_tirx/backend/cuda/transform/__init__.py
@@ -24,32 +24,141 @@ from tvm_ffi import get_global_func
def LowerIket():
"""Lower frontend-only NVIDIA IKET annotations.
- This pass must run after ``tvm_tirx.tirx.transform.SplitHostDevice`` and
before
- ``tvm_tirx.tirx.transform.MakePackedAPI``. It strips annotations for
regular
- compilation and emits NVIDIA IKET metadata and NativeDump placeholders
- when the IRModule is explicitly IKET-enabled. Trace collection and
- postprocessing are owned by the external ``run-iket`` process.
-
- Returns
- -------
- fpass : tvm_tirx.ir.transform.Pass
- The result pass.
+ Run after ``SplitHostDevice`` and before ``MakePackedAPI``. Strip
+ annotations for regular compilation; emit metadata and NativeDump
+ placeholders for IKET-enabled modules. The external ``run-iket`` process
+ owns trace collection and postprocessing.
+
+ Examples
+ --------
+ .. code:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import tirx as T
+ from tvm_tirx.backend.cuda.transform import LowerIket
+
+ @T.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((1,), "int32")):
+ T.cuda.iket_mark("store")
+ A[0] = 1
+
+ @T.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((1,), "int32")):
+ A[0] = 1
+
+ after = LowerIket()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
"""
return get_global_func("tirx.backend.cuda.transform.LowerIket")()
def BindBackendConfig():
- """Resolve per-entry compilation settings and targets before tile
dispatch."""
+ """Resolve CUDA function and device-entry compilation settings.
+
+ Run after ``BindTarget`` and before tile dispatch. Function defaults and
+ entry overrides become ``target`` and JSON ``backend_config`` attributes;
+ non-CUDA functions are unchanged.
+
+ Examples
+ --------
+ .. code:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import tirx as T
+ from tvm_tirx.backend.config import backend_config_json
+ from tvm_tirx.backend.cuda.backend_config import resolve_backend_config
+ from tvm_tirx.backend.cuda.transform import BindBackendConfig
+
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
+ config = backend_config_json({"cuda":
resolve_backend_config(target=target)})
+
+ @T.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((1,), "int32")):
+ T.func_attr({"target": target})
+ with T.device_entry():
+ A[0] = 1
+
+ @T.ir_module
+ class Expected:
+ @T.function
+ def main(A: T.Tensor((1,), "int32")):
+ T.func_attr({"target": target, "backend_config": config})
+ with T.device_entry(attrs={"target": target, "backend_config":
config}):
+ A[0] = 1
+
+ after = BindBackendConfig()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+ """
return get_global_func("tirx.backend.cuda.transform.BindBackendConfig")("")
def SpecializeDeviceHelpers():
- """Clone private device helpers into each independently compiled CUDA
group."""
+ """Specialize private device helpers for each CUDA compilation group.
+
+ Run after ``SplitHostDevice``. Kernel target and backend configuration
define
+ each group. Called helpers receive those attributes and are cloned when
+ their settings differ; helpers already in the same group can be reused.
+ Unreferenced device-only helpers are omitted from the output module.
+ """
return
get_global_func("tirx.backend.cuda.transform.SpecializeDeviceHelpers")()
def SpecializeEntryHelpers():
- """Specialize shared helper calls before target-sensitive tile dispatch."""
+ """Specialize helpers for each CUDA device entry's resolved settings.
+
+ Run after ``BindBackendConfig`` and before ``TileDispatch`` and
+ ``SplitHostDevice``. Clone helpers when entry settings differ, propagating
+ them through the call graph; remove originals that become unreachable.
+
+ Examples
+ --------
+ .. code:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import tirx as T
+ from tvm_tirx.backend.cuda.transform import SpecializeEntryHelpers
+
+ def settings(arch):
+ return {"target": tvm.target.Target({"kind": "cuda", "arch":
arch}),
+ "backend_config": '{"cuda":{"arch":"' + arch + '"}}'}
+
+ host_attrs, entry_attrs = settings("sm_100a"), settings("sm_90a")
+
+ @T.ir_module
+ class Before:
+ @T.function(private=True)
+ def helper(x: T.int32) -> T.int32:
+ T.func_attr(host_attrs)
+ return x + 1
+
+ @T.function
+ def main(A: T.Tensor((1,), "int32")):
+ T.func_attr(host_attrs)
+ with T.device_entry(attrs=entry_attrs):
+ A[0] = Before.helper(A[0])
+
+ @T.ir_module
+ class Expected:
+ @T.function(private=True)
+ def helper_cuda(x: T.int32) -> T.int32:
+ T.func_attr(entry_attrs)
+ return x + 1
+
+ @T.function
+ def main(A: T.Tensor((1,), "int32")):
+ T.func_attr(host_attrs)
+ with T.device_entry(attrs=entry_attrs):
+ A[0] = Expected.helper_cuda(A[0])
+
+ after = SpecializeEntryHelpers()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+ """
return
get_global_func("tirx.backend.cuda.transform.SpecializeEntryHelpers")()
diff --git a/tirx/python/tvm_tirx/backend/trn/pipeline.py
b/tirx/python/tvm_tirx/backend/trn/pipeline.py
index 7b3025189e..afccc76832 100644
--- a/tirx/python/tvm_tirx/backend/trn/pipeline.py
+++ b/tirx/python/tvm_tirx/backend/trn/pipeline.py
@@ -24,7 +24,19 @@ from . import transform as trn_transform
def trn_pipeline(*, target=None):
- """The Trainium pipeline used in tvm_tirx.tirx.compile."""
+ """Construct the Trainium lowering and host/device finalization stages.
+
+ Parameters
+ ----------
+ target : Target or None
+ Target to bind before lowering, or None to retain existing targets.
+
+ Returns
+ -------
+ stages : tuple
+ The lowering ModulePass, host-finalization factory, and
+ device-finalization factory consumed by ``tvm_tirx.tirx.compile``.
+ """
@tvm_tirx.ir.transform.module_pass(opt_level=0)
def _pipeline(
@@ -52,7 +64,7 @@ def trn_pipeline(*, target=None):
def finalize_device_passes_trn(): # pylint: disable=unused-argument
- """The finalization passes for the Trainium backend."""
+ """Construct the statement-simplification stage for Trainium device IR."""
return tvm_tirx.ir.transform.Sequential([tirx.transform.StmtSimplify()])
diff --git a/tirx/python/tvm_tirx/backend/trn/transform/__init__.py
b/tirx/python/tvm_tirx/backend/trn/transform/__init__.py
index 963270889b..3f311cc17f 100644
--- a/tirx/python/tvm_tirx/backend/trn/transform/__init__.py
+++ b/tirx/python/tvm_tirx/backend/trn/transform/__init__.py
@@ -29,12 +29,52 @@ _LAZY_TRANSFORMS = {
def LowerTrainiumLayout():
- """Lower Trainium layouts to backend physical tensor shapes and indices."""
+ """Lower logical tensor accesses to Trainium physical shapes and indices.
+
+ SBUF uses partition/free dimensions; PSUM adds a bank dimension. Global
+ memory is flattened. Parameter ABI shapes stay logical. Run after tile
+ dispatch.
+
+ Examples
+ --------
+ .. code:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import tirx as T
+ from tvm_tirx.backend.trn.axis import P, F
+ from tvm_tirx.backend.trn.transform import LowerTrainiumLayout
+
+ @T.ir_module
+ class Before:
+ @T.function
+ def main():
+ T.device_entry()
+ A = T.alloc_tensor(ty_args=[T.Tensor(
+ (128, 2, 2), "float32", scope="trn.sbuf",
+ layout=T.TileLayout(T.S[(128, 4):(1 @ P, 1 @ F)]))])
+ A[0, 1, 0] = T.float32(1)
+
+ @T.ir_module
+ class Expected:
+ @T.function
+ def main():
+ T.device_entry()
+ A = T.alloc_tensor(ty_args=[T.Tensor((128, 4), "float32",
scope="trn.sbuf")])
+ A[0, 2] = T.float32(1)
+
+ after = LowerTrainiumLayout()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+ """
return get_global_func("tirx.backend.trn.transform.LowerTrainiumLayout")()
def LowerTIRx():
- """Lower TIRx tile op calls for the Trainium backend."""
+ """Run ``TileDispatch`` followed by ``LowerTrainiumLayout``.
+
+ Run ``TrnPrivateTensorAlloc`` and ``TrnNaiveAllocator`` first when tiles
+ need workspaces or allocation addresses. Tile dispatch requires a
+ Trainium target on the function or in the current target scope.
+ """
return tvm_tirx.ir.transform.Sequential(
[tirx.transform.TileDispatch(), LowerTrainiumLayout()],
name="tirx.backend.trn.LowerTIRx",
diff --git a/tirx/python/tvm_tirx/backend/trn/transform/naive_allocator.py
b/tirx/python/tvm_tirx/backend/trn/transform/naive_allocator.py
index 26b5d8eafa..bb06a72538 100644
--- a/tirx/python/tvm_tirx/backend/trn/transform/naive_allocator.py
+++ b/tirx/python/tvm_tirx/backend/trn/transform/naive_allocator.py
@@ -123,6 +123,36 @@ def _allocate_missing_tensors(stmt, alloc_pool_start: int):
@function_pass(opt_level=0, name="TrnNaiveAllocator")
class TrnNaiveAllocator:
+ """Assign consecutive byte offsets to unaddressed tensor allocations.
+
+ Handles constant-sized ``trn.sbuf``, ``shared``, and ``shared.dyn``
tensors;
+ SBUF sizes count bytes per partition. New offsets follow existing
+ allocations, without lifetime-based reuse.
+
+ Examples
+ --------
+ .. code:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import tirx as T
+ from tvm_tirx.backend.trn.transform import TrnNaiveAllocator
+
+ @T.ir_module
+ class Before:
+ @T.function
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((128, 4), "float32",
scope="trn.sbuf")])
+
+ @T.ir_module
+ class Expected:
+ @T.function
+ def main():
+ A = T.alloc_tensor((0,), ty_args=[T.Tensor((128, 4),
"float32", scope="trn.sbuf")])
+
+ after = TrnNaiveAllocator()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+ """
+
def transform_function(self, func, mod, ctx):
alloc_pool_start = _get_alloc_pool_start(func.body)
new_body = _allocate_missing_tensors(func.body, alloc_pool_start)
diff --git a/tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py
b/tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py
index 81ef1ac0ce..d8de6f0131 100644
--- a/tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py
+++ b/tirx/python/tvm_tirx/backend/trn/transform/private_tensor_alloc.py
@@ -139,7 +139,52 @@ def private_alloc(stmt: Stmt, target: Target) -> Stmt:
@function_pass(opt_level=0, name="TrnPrivateTensorAlloc")
class TrnPrivateTensorAlloc:
- """Generate private tensor allocations for each tensor instruction Call"""
+ """Allocate and initialize instruction-private Trainium workspaces.
+
+ Insert workspaces at the device-entry root and attach them to tile calls.
+ Run before tile dispatch with a Trainium target. Calls needing no
+ workspace are unchanged.
+
+ Examples
+ --------
+ .. code:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import tirx as T
+ from tvm_tirx.backend.trn.axis import P, F
+ from tvm_tirx.backend.trn.transform import TrnPrivateTensorAlloc
+
+ tile = T.Tensor((128, 512), "float32", scope="trn.sbuf",
+ layout=T.TileLayout(T.S[(128, 512):(1 @ P, 1 @ F)]))
+
+ @T.ir_module
+ class Before:
+ @T.function
+ def main():
+ T.device_entry()
+ A = T.alloc_tensor(ty_args=[tile])
+ C = T.alloc_tensor(ty_args=[tile])
+ T.trn.tile.activation(C, A, bias=T.float32(1), opcode="exp")
+
+ @T.ir_module
+ class Expected:
+ @T.function
+ def main():
+ T.device_entry()
+ const_bias = T.alloc_tensor(
+ ty_args=[T.Tensor((128, 512), "float32",
scope="trn.sbuf")])
+ with T.nki.tensorized_instruction():
+ for p in T.serial(128, annotations={"nki_dim": "P"}):
+ for f in T.serial(512, annotations={"nki_dim": "F"}):
+ T.nki.memset(const_bias[p, f], T.float32(1))
+ A = T.alloc_tensor(ty_args=[tile])
+ C = T.alloc_tensor(ty_args=[tile])
+ T.trn.tile.activation(C, A, bias=T.float32(1),
const_bias=const_bias, opcode="exp")
+
+ with tvm.target.Target("aws/trn1/trn1.2xlarge"):
+ after = TrnPrivateTensorAlloc()(Before)
+ tvm.ir.assert_structural_equal(after, Expected)
+ """
def transform_function(self, func, mod, ctx):
target = func.attrs.get("target", None)
diff --git a/tirx/python/tvm_tirx/codegen/executable.py
b/tirx/python/tvm_tirx/codegen/executable.py
index 98e10bd0ff..0bc4093afc 100644
--- a/tirx/python/tvm_tirx/codegen/executable.py
+++ b/tirx/python/tvm_tirx/codegen/executable.py
@@ -328,7 +328,19 @@ class Executable:
def load_static_library(path, func_names):
- """Load the .o library at path which implements functions with func_names.
- Unlike the generic load_module the result will remain as a static_library
- and will not be relinked on-the-fly into a .so library."""
+ """Load an object file as a static library module without relinking it.
+
+ Parameters
+ ----------
+ path : str
+ Path to the object file (``.o``) to load.
+ func_names : Sequence[str]
+ Names of the functions implemented by the object file.
+
+ Returns
+ -------
+ module : tvm_ffi.Module
+ Static library module that can be imported into another module for
+ linking during export. It is not converted to a shared library here.
+ """
return _ffi_api.ModuleLoadStaticLibrary(path, func_names)
diff --git a/tirx/python/tvm_tirx/ir/attrs.py b/tirx/python/tvm_tirx/ir/attrs.py
index 02fdecd563..86c7c42a53 100644
--- a/tirx/python/tvm_tirx/ir/attrs.py
+++ b/tirx/python/tvm_tirx/ir/attrs.py
@@ -22,10 +22,11 @@ import tvm_ffi._ffi_api as _tvm_ffi_api
@tvm_ffi.register_object("ir.Attrs")
class Attrs(tvm_ffi.Object):
- """Attribute node, which is mainly use for defining attributes of
operators.
+ """Base class for structured operator and function attributes.
- Used by function registered in python side, such as compute, schedule and
alter_layout.
- Attrs is passed as the first argument to these functions.
+ Concrete subclasses declare named fields exposed through attribute access
or
+ ``attrs["field"]``. Use :func:`make_node` to construct a registered
attribute
+ type by its type key and fields.
"""
def __getitem__(self, item):
@@ -34,7 +35,12 @@ class Attrs(tvm_ffi.Object):
@tvm_ffi.register_object("ir.DictAttrs")
class DictAttrs(Attrs):
- """Dictionary attributes."""
+ """A string-keyed attribute map for functions, modules, and calls.
+
+ Construct with ``DictAttrs(dict={"key": value})`` or
+ ``make_node("ir.DictAttrs", key=value)``. Entries are available through
+ indexing, :meth:`get`, and attribute access; ``dict`` holds the underlying
map.
+ """
def keys(self):
"""Get list of names in the attribute.
diff --git a/tirx/python/tvm_tirx/ir/base.py b/tirx/python/tvm_tirx/ir/base.py
index ed7f76fdab..61ffe54627 100644
--- a/tirx/python/tvm_tirx/ir/base.py
+++ b/tirx/python/tvm_tirx/ir/base.py
@@ -48,9 +48,10 @@ class Node(tvm_ffi.Object):
@register_object("ir.EnvFunc")
class EnvFunc(tvm_ffi.Object):
- """Environment function.
+ """A registered global function referenced and serialized by name.
- This is a global function object that can be serialized by its name.
+ Use ``EnvFunc.get(name)`` to look up a function and call the returned
object
+ with its arguments. Deserialization resolves the same name in the registry.
"""
def __call__(self, *args):
diff --git a/tirx/python/tvm_tirx/ir/expr.py b/tirx/python/tvm_tirx/ir/expr.py
index 8c847f9fd9..e646e6da8a 100644
--- a/tirx/python/tvm_tirx/ir/expr.py
+++ b/tirx/python/tvm_tirx/ir/expr.py
@@ -52,7 +52,12 @@ def _convert_subscript_index(index):
@tvm_ffi.register_object("ir.Expr")
class Expr(Node):
- """Base class of all the expressions."""
+ """Base class of typed IR expressions.
+
+ ``ty`` holds the expression's type, which may be :class:`MissingType`
during
+ construction. ``loc`` records its source location. Concrete subclasses
define
+ operands and may expose operations through their type.
+ """
loc: Location
ty: "tvm_tirx.ir.Type"
@@ -498,10 +503,19 @@ class TupleGetItem(_CallableExprWithOp):
@tvm_ffi.register_object("ir.TensorLoad")
class TensorLoad(_CallableExprWithOp):
- """An indexed load from an expression source.
+ """Load an element or vector from an indexed expression source.
- The source type's registered inference hook validates the inputs and
- derives the primitive result type.
+ The source type's registered inference hook validates the inputs and
derives
+ the primitive result type.
+
+ Parameters
+ ----------
+ source : Expr
+ The tensor or other expression supporting indexed loads.
+ indices : list[Expr]
+ Indices interpreted by the source type, in dimension order.
+ loc : Location, optional
+ The source location of the load.
"""
source: Expr
@@ -523,7 +537,7 @@ class FloatImm(Constant):
Parameters
----------
- ty : str, DataType or PrimType
+ ty : str, tvm_ffi.dtype or PrimType
The data type
value : float
@@ -557,7 +571,7 @@ class IntImm(Constant):
Parameters
----------
- ty : str, DataType or PrimType
+ ty : str, tvm_ffi.dtype or PrimType
The data type
value : int
@@ -602,7 +616,17 @@ class IntImm(Constant):
@tvm_ffi.register_object("ir.GenericConst")
class GenericConst(_ExprCallable, Constant):
- """A literal payload with an explicit expression type."""
+ """A literal payload with an explicit expression type.
+
+ Parameters
+ ----------
+ ty : Type
+ The type assigned to the literal.
+ value : object
+ The payload, convertible to an FFI value.
+ loc : Location, optional
+ The source location of the literal.
+ """
def __init__(self, ty: "tvm_tirx.ir.Type", value, loc: Location =
UNKNOWN_LOC) -> None:
self.__init_handle_by_constructor__(_ffi_api.GenericConst, ty, value,
loc)
@@ -631,7 +655,15 @@ class DataTypeImm(Constant):
@tvm_ffi.register_object("ir.StringImm")
class StringImm(Constant):
- """A string literal with StringType."""
+ """A string literal with :class:`StringType`.
+
+ Parameters
+ ----------
+ value : str
+ The literal text.
+ loc : Location, optional
+ The source location of the literal.
+ """
value: str
@@ -662,13 +694,28 @@ class CallEffectKind(IntEnum):
@tvm_ffi.register_object("ir.Call")
class Call(_CallableExprWithOp):
- """Core function call node.
+ """Call a registered operator or another callable IR expression.
Omitted or ``None`` keyword-only ``ty`` uses available result inference,
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.
+
+ Parameters
+ ----------
+ op : Expr or str
+ The callee expression, or a registered operator name resolved by
``Op.get``.
+ args : sequence[Expr]
+ Positional value arguments.
+ attrs : Attrs or dict, optional
+ Operator-specific attributes. A dictionary becomes ``DictAttrs``.
+ ty_args : sequence[Type], optional
+ Explicit type arguments, empty by default.
+ loc : Location, optional
+ The source location of the call.
+ ty : Type or str, optional
+ Keyword-only result type; a primitive dtype string is accepted.
"""
op: Expr
@@ -762,15 +809,18 @@ class Var(_CallableExprWithOp):
Parameters
----------
- name : str
- The name of the variable.
+ name : str, optional
+ The display name. Variables with the same name remain distinct
bindings.
ty : Optional[Type or str]
- The exact type of the variable. A string denotes a primitive dtype.
+ The variable type, or MissingType when omitted. A string denotes a
+ primitive dtype; ``"handle"`` denotes a pointer to void.
loc : Location, optional
Location that points to the original source code.
+ name_hint : str, optional
+ Keyword-only alternative to ``name``; supply exactly one name spelling.
"""
name: str
@@ -846,7 +896,7 @@ class LambdaExpr(StagingExpr, Scriptable):
return types.
As a StagingExpr, LambdaExpr is eliminated during compilation and does not
- remain in executable IR.
+ remain in executable IR. In TVMScript, use ``T.Lambda`` to construct it.
Parameters
----------
@@ -883,27 +933,18 @@ class LambdaExpr(StagingExpr, Scriptable):
@tvm_ffi.register_object("ir.Range")
class Range(Node, Scriptable):
- """Represent a range in TVM.
+ """A half-open interval stored as ``min`` and ``extent``.
- You do not need to create a Range explicitly.
- Python lists and tuples will be converted automatically to a Range in API
functions.
+ Use :meth:`from_min_extent` when the extent is already known.
Parameters
----------
begin : Expr
- The begin value of the range when end is not None.
- Otherwise it is the extent of the range ``[0, begin)``.
-
- end : Optional[Expr]
- The end value of the range.
-
+ The lower bound when ``end`` is provided; otherwise the extent of
``[0, begin)``.
+ end : Expr, optional
+ The exclusive upper bound. With two arguments, ``extent`` is ``end -
begin``.
loc : Location, optional
- The location of this node in the source code.
-
- Note
- ----
- The constructor creates the range `[begin, end)`
- if the end argument is not None. Otherwise, it creates `[0, begin)`.
+ The source location of the range.
"""
begin: Expr
diff --git a/tirx/python/tvm_tirx/ir/function.py
b/tirx/python/tvm_tirx/ir/function.py
index eb16149882..ef1e4ac6fd 100644
--- a/tirx/python/tvm_tirx/ir/function.py
+++ b/tirx/python/tvm_tirx/ir/function.py
@@ -38,7 +38,12 @@ class CallingConv(IntEnum):
@tvm_ffi.register_object("ir.BaseFunc")
class BaseFunc(_CallableExprWithOp):
- """Base class of all functions."""
+ """Base class for callable IR function definitions.
+
+ ``attrs`` stores function metadata as :class:`DictAttrs`. Attribute update
+ methods return a new function, preserving the original. Concrete function
+ classes define the parameters, body, and function type in ``ty``.
+ """
def __bool__(self) -> bool:
return True
@@ -80,7 +85,7 @@ class BaseFunc(_CallableExprWithOp):
)
def with_attrs(self, attr_map: DictAttrs | dict[str, tvm_ffi.Object]) ->
"BaseFunc":
- """Copy the IRModule and add the given attribute map to it.
+ """Copy the function and add the given attribute map to it.
Parameters
----------
attr_map: Union[DictAttrs, Dict[str, tvm_ffi.Object]]
@@ -96,12 +101,12 @@ class BaseFunc(_CallableExprWithOp):
return _ffi_api.BaseFuncWithAttrs(self, attr_map)
def without_attr(self, attr_key: str) -> "BaseFunc":
- """Create a new copy of the function with an attribute without
provided key.
+ """Copy the function and remove the given attribute key.
Parameters
----------
attr_key : str
- The attribute key to delete from the attrubte pairs.
+ The attribute key to delete from the attribute pairs.
Returns
diff --git a/tirx/python/tvm_tirx/ir/global_info.py
b/tirx/python/tvm_tirx/ir/global_info.py
index c6991cdbf3..0351222c0e 100644
--- a/tirx/python/tvm_tirx/ir/global_info.py
+++ b/tirx/python/tvm_tirx/ir/global_info.py
@@ -21,7 +21,11 @@ import tvm_ffi
@tvm_ffi.register_object("ir.GlobalInfo")
class GlobalInfo(tvm_ffi.Object):
- """Base node for all global info that can appear in the IR"""
+ """Base class for module-level information referenced by the IR.
+
+ Concrete subclasses are stored in ``IRModule.global_infos``, a map from
names
+ to arrays of global information objects.
+ """
def __eq__(self, other):
"""Compare two global info objects for structural equivalence."""
diff --git a/tirx/python/tvm_tirx/ir/location.py
b/tirx/python/tvm_tirx/ir/location.py
index 5f3d7f962d..16d40a7e87 100644
--- a/tirx/python/tvm_tirx/ir/location.py
+++ b/tirx/python/tvm_tirx/ir/location.py
@@ -83,7 +83,15 @@ class SourceLoc(Location):
@register_object("ir.CallSiteLoc")
class CallSiteLoc(Location):
- """A callee's location together with the location of its caller."""
+ """A callee's source location together with its caller's location.
+
+ Parameters
+ ----------
+ callee : Location
+ The location of the called or inlined code.
+ caller : Location
+ The calling location, which may itself be a CallSiteLoc.
+ """
def __init__(self, callee: Location, caller: Location):
self.__init_handle_by_constructor__(_ffi_api.CallSiteLoc, callee,
caller)
diff --git a/tirx/python/tvm_tirx/ir/module.py
b/tirx/python/tvm_tirx/ir/module.py
index 5857721360..1a4405d1ec 100644
--- a/tirx/python/tvm_tirx/ir/module.py
+++ b/tirx/python/tvm_tirx/ir/module.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""IRModule that holds the functions and type definitions."""
+"""Modules containing global functions and shared metadata."""
import tvm_ffi
@@ -27,14 +27,19 @@ from .function import BaseFunc
@tvm_ffi.register_object("ir.IRModule")
class IRModule(Node, Scriptable):
- """IRModule that holds functions and type definitions.
+ """A collection of global functions and module-level metadata.
- IRModule is the basic unit for all IR transformations across the stack.
+ Modules are the input and output of compiler passes. Functions can be
looked
+ up by their :class:`GlobalVar` or name.
Parameters
----------
- functions: Optional[dict].
- Map of global var to BaseFunc
+ functions : dict[GlobalVar or str, BaseFunc], optional
+ Global function definitions. String keys are converted to GlobalVar
objects.
+ attrs : dict, optional
+ Module attributes, stored as DictAttrs.
+ global_infos : dict[str, list[GlobalInfo]], optional
+ Named groups of global information referenced by the IR.
"""
def __init__(self, functions=None, attrs=None, global_infos=None):
@@ -83,10 +88,10 @@ class IRModule(Node, Scriptable):
Parameters
---------
- var: GlobalVar
- The global variable.
+ var: str or GlobalVar
+ The global variable or its name.
- val: Union[Function, Type]
+ val: BaseFunc
The value.
"""
return self._add(var, val, True)
@@ -105,13 +110,13 @@ class IRModule(Node, Scriptable):
Parameters
----------
- var: Union[String, GlobalVar, GlobalTypeVar]
+ var: str or GlobalVar
The name or global variable.
Returns
-------
- val: Union[Function, Type]
- The definition referenced by :code:`var` (either a function or
type).
+ val: BaseFunc
+ The function referenced by :code:`var`.
"""
if isinstance(var, str):
return _ffi_api.IRModuleLookupStr(self, var)
@@ -143,8 +148,8 @@ class IRModule(Node, Scriptable):
Parameters
----------
- var: GlobalVar
- The global variable.
+ var: str or GlobalVar
+ The global variable or its name.
func: tvm_tirx.ir.BaseFunc
The function to be inserted.
diff --git a/tirx/python/tvm_tirx/ir/op.py b/tirx/python/tvm_tirx/ir/op.py
index 6686fe7550..b3e8b55213 100644
--- a/tirx/python/tvm_tirx/ir/op.py
+++ b/tirx/python/tvm_tirx/ir/op.py
@@ -32,7 +32,12 @@ from .expr import Expr
@tvm_ffi.register_object("ir.Op")
class Op(Expr):
- """Primitive operator in the IR."""
+ """A registered primitive operator used as a :class:`Call` callee.
+
+ Retrieve the shared operator with ``Op.get(name)``. Its registration
supplies
+ argument signatures, attributes, validation, and result-type inference;
+ constructing a Call does not itself execute the operator.
+ """
def __init__(self):
raise RuntimeError("Cannot create op, use get instead")
diff --git a/tirx/python/tvm_tirx/ir/prim/expr.py
b/tirx/python/tvm_tirx/ir/prim/expr.py
index 8113d87fa6..acf83bf3fe 100644
--- a/tirx/python/tvm_tirx/ir/prim/expr.py
+++ b/tirx/python/tvm_tirx/ir/prim/expr.py
@@ -27,22 +27,34 @@ from . import _ffi_api as _prim_ffi_api
class BinaryOpExpr(ExprWithOp):
+ """Python base class for binary arithmetic and bitwise expressions.
+
+ Concrete nodes store left and right primitive operands in ``a`` and ``b``.
+ """
+
a: Expr
b: Expr
class CmpExpr(ExprWithOp):
+ """Python base class for comparisons of primitive expressions.
+
+ Concrete nodes compare ``a`` and ``b`` and produce boolean scalar or
vector values.
+ """
+
a: Expr
b: Expr
class LogicalExpr(ExprWithOp):
+ """Python base class for boolean conjunction and disjunction
expressions."""
+
pass
@tvm_ffi.register_object("prim.Cast")
class Cast(ExprWithOp):
- """Cast expression.
+ """Convert a primitive value to another dtype while preserving its lane
count.
Parameters
----------
@@ -66,7 +78,7 @@ class Cast(ExprWithOp):
@tvm_ffi.register_object("prim.LShift")
class LShift(BinaryOpExpr):
- """Left shift node.
+ """Shift the bits of ``a`` left by ``b`` positions.
Parameters
----------
@@ -86,7 +98,7 @@ class LShift(BinaryOpExpr):
@tvm_ffi.register_object("prim.RShift")
class RShift(BinaryOpExpr):
- """Right shift node.
+ """Shift the bits of ``a`` right by ``b`` positions.
Parameters
----------
@@ -106,7 +118,7 @@ class RShift(BinaryOpExpr):
@tvm_ffi.register_object("prim.BitwiseAnd")
class BitwiseAnd(BinaryOpExpr):
- """Bitwise and node.
+ """Compute the bitwise intersection of ``a`` and ``b``.
Parameters
----------
@@ -126,7 +138,7 @@ class BitwiseAnd(BinaryOpExpr):
@tvm_ffi.register_object("prim.BitwiseOr")
class BitwiseOr(BinaryOpExpr):
- """Bitwise or node.
+ """Compute the bitwise union of ``a`` and ``b``.
Parameters
----------
@@ -146,7 +158,7 @@ class BitwiseOr(BinaryOpExpr):
@tvm_ffi.register_object("prim.BitwiseXor")
class BitwiseXor(BinaryOpExpr):
- """Bitwise xor node.
+ """Compute the bitwise exclusive-or of ``a`` and ``b``.
Parameters
----------
@@ -166,7 +178,7 @@ class BitwiseXor(BinaryOpExpr):
@tvm_ffi.register_object("prim.BitwiseNot")
class BitwiseNot(ExprWithOp):
- """Bitwise not node.
+ """Invert every bit of the input value.
Parameters
----------
@@ -185,7 +197,7 @@ class BitwiseNot(ExprWithOp):
@tvm_ffi.register_object("prim.Add")
class Add(BinaryOpExpr):
- """Add node.
+ """Add two primitive values, elementwise for vectors.
Parameters
----------
@@ -205,7 +217,7 @@ class Add(BinaryOpExpr):
@tvm_ffi.register_object("prim.Sub")
class Sub(BinaryOpExpr):
- """Sub node.
+ """Subtract ``b`` from ``a``, elementwise for vectors.
Parameters
----------
@@ -225,7 +237,7 @@ class Sub(BinaryOpExpr):
@tvm_ffi.register_object("prim.Mul")
class Mul(BinaryOpExpr):
- """Mul node.
+ """Multiply two primitive values, elementwise for vectors.
Parameters
----------
@@ -245,7 +257,7 @@ class Mul(BinaryOpExpr):
@tvm_ffi.register_object("prim.Div")
class Div(BinaryOpExpr):
- """Div node.
+ """Divide ``a`` by ``b``; integer division truncates toward zero.
Parameters
----------
@@ -265,7 +277,7 @@ class Div(BinaryOpExpr):
@tvm_ffi.register_object("prim.Mod")
class Mod(BinaryOpExpr):
- """Mod node.
+ """Compute the remainder of integer division truncated toward zero.
Parameters
----------
@@ -285,7 +297,7 @@ class Mod(BinaryOpExpr):
@tvm_ffi.register_object("prim.FloorDiv")
class FloorDiv(BinaryOpExpr):
- """FloorDiv node.
+ """Divide ``a`` by ``b``, rounding the quotient down toward negative
infinity.
Parameters
----------
@@ -305,7 +317,7 @@ class FloorDiv(BinaryOpExpr):
@tvm_ffi.register_object("prim.FloorMod")
class FloorMod(BinaryOpExpr):
- """FloorMod node.
+ """Compute ``a - FloorDiv(a, b) * b``, the remainder of floor division.
Parameters
----------
@@ -325,7 +337,7 @@ class FloorMod(BinaryOpExpr):
@tvm_ffi.register_object("prim.Min")
class Min(BinaryOpExpr):
- """Min node.
+ """Choose the smaller of two primitive values, elementwise for vectors.
Parameters
----------
@@ -345,7 +357,7 @@ class Min(BinaryOpExpr):
@tvm_ffi.register_object("prim.Max")
class Max(BinaryOpExpr):
- """Max node.
+ """Choose the larger of two primitive values, elementwise for vectors.
Parameters
----------
@@ -365,7 +377,7 @@ class Max(BinaryOpExpr):
@tvm_ffi.register_object("prim.EQ")
class EQ(CmpExpr):
- """EQ node.
+ """Compare whether ``a`` equals ``b``, elementwise for vectors.
Parameters
----------
@@ -385,7 +397,7 @@ class EQ(CmpExpr):
@tvm_ffi.register_object("prim.NE")
class NE(CmpExpr):
- """NE node.
+ """Compare whether ``a`` differs from ``b``, elementwise for vectors.
Parameters
----------
@@ -405,7 +417,7 @@ class NE(CmpExpr):
@tvm_ffi.register_object("prim.LT")
class LT(CmpExpr):
- """LT node.
+ """Compare whether ``a`` is less than ``b``, elementwise for vectors.
Parameters
----------
@@ -425,7 +437,7 @@ class LT(CmpExpr):
@tvm_ffi.register_object("prim.LE")
class LE(CmpExpr):
- """LE node.
+ """Compare whether ``a`` is at most ``b``, elementwise for vectors.
Parameters
----------
@@ -445,7 +457,7 @@ class LE(CmpExpr):
@tvm_ffi.register_object("prim.GT")
class GT(CmpExpr):
- """GT node.
+ """Compare whether ``a`` is greater than ``b``, elementwise for vectors.
Parameters
----------
@@ -465,7 +477,7 @@ class GT(CmpExpr):
@tvm_ffi.register_object("prim.GE")
class GE(CmpExpr):
- """GE node.
+ """Compare whether ``a`` is at least ``b``, elementwise for vectors.
Parameters
----------
@@ -485,7 +497,7 @@ class GE(CmpExpr):
@tvm_ffi.register_object("prim.And")
class And(LogicalExpr):
- """And node.
+ """Compute the conjunction of two boolean values.
Parameters
----------
@@ -505,7 +517,7 @@ class And(LogicalExpr):
@tvm_ffi.register_object("prim.Or")
class Or(LogicalExpr):
- """Or node.
+ """Compute the disjunction of two boolean values.
Parameters
----------
@@ -528,7 +540,7 @@ class Or(LogicalExpr):
@tvm_ffi.register_object("prim.Not")
class Not(LogicalExpr):
- """Not node.
+ """Negate a boolean value, elementwise for vectors.
Parameters
----------
@@ -547,7 +559,7 @@ class Not(LogicalExpr):
@tvm_ffi.register_object("prim.Select")
class Select(ExprWithOp):
- """Select node.
+ """Select between two values according to a boolean condition.
Note
----
@@ -595,7 +607,7 @@ class Select(ExprWithOp):
@tvm_ffi.register_object("prim.Let")
class Let(ExprWithOp):
- """Let node.
+ """Bind ``var`` to ``value`` within an expression body and return the body
result.
Parameters
----------
diff --git a/tirx/python/tvm_tirx/ir/prim/vector_expr.py
b/tirx/python/tvm_tirx/ir/prim/vector_expr.py
index 6712d0e45e..827a0887fe 100644
--- a/tirx/python/tvm_tirx/ir/prim/vector_expr.py
+++ b/tirx/python/tvm_tirx/ir/prim/vector_expr.py
@@ -25,7 +25,7 @@ from . import _ffi_api as _prim_ffi_api
@tvm_ffi.register_object("prim.Ramp")
class Ramp(ExprWithOp):
- """Ramp node.
+ """Construct a vector whose lane ``i`` equals ``base + i * stride``.
Parameters
----------
@@ -36,7 +36,7 @@ class Ramp(ExprWithOp):
The stride of the ramp.
lanes : Expr
- The lanes of the expression.
+ The number of vector lanes.
loc : Location, optional
The location of this expression in the source code.
@@ -58,7 +58,7 @@ class Ramp(ExprWithOp):
@tvm_ffi.register_object("prim.Broadcast")
class Broadcast(ExprWithOp):
- """Broadcast node.
+ """Repeat a scalar value across all lanes of a vector.
Parameters
----------
@@ -66,7 +66,7 @@ class Broadcast(ExprWithOp):
The value of the expression.
lanes : Expr
- The lanes of the expression.
+ The number of vector lanes.
loc : Location, optional
The location of this expression in the source code.
@@ -81,15 +81,15 @@ class Broadcast(ExprWithOp):
@tvm_ffi.register_object("prim.Shuffle")
class Shuffle(ExprWithOp):
- """Shuffle node.
+ """Select output lanes from the concatenation of the input vectors.
Parameters
----------
vectors : List[Expr]
- The vectors
+ Input vectors concatenated in the given order.
indices : List[Expr]
- The indices
+ Lane indices into the concatenated input; one index per output lane.
loc : Location, optional
The location of this expression in the source code.
diff --git a/tirx/python/tvm_tirx/ir/stmt.py b/tirx/python/tvm_tirx/ir/stmt.py
index 4b18b32698..b7bdddd877 100644
--- a/tirx/python/tvm_tirx/ir/stmt.py
+++ b/tirx/python/tvm_tirx/ir/stmt.py
@@ -69,7 +69,7 @@ class Bind(Stmt):
@tvm_ffi.register_object("ir.AssertStmt")
class AssertStmt(Stmt):
- """AssertStmt node.
+ """Check a condition and raise the specified error when it is false.
Parameters
----------
@@ -110,12 +110,11 @@ class AssertStmt(Stmt):
class ForKind(IntEnum):
- """The kind of the for loop.
+ """Execution semantics of a :class:`For` loop.
- note
- ----
- ForKind can change the control flow semantics
- of the loop and need to be considered in all TIR passes.
+ ``DEFAULT`` executes iterations sequentially; ``PARALLEL`` permits parallel
+ execution. ``VECTORIZED`` requests a vectorized body and ``UNROLLED``
requests
+ unrolling. Transformations must preserve the selected execution semantics.
"""
DEFAULT = 0 # Ordinary loop with sequential iteration semantics.
@@ -126,7 +125,7 @@ class ForKind(IntEnum):
@tvm_ffi.register_object("ir.For")
class For(Stmt):
- """For node.
+ """Iterate from ``min`` while below ``min + extent``, incrementing by
``step``.
Parameters
----------
@@ -143,7 +142,7 @@ class For(Stmt):
The required loop step, including a typed one for unit stride.
kind : ForKind
- The type of the for.
+ Sequential, parallel, vectorized, or unrolled execution semantics.
body : Stmt | Sequence[Stmt]
The body statement.
@@ -190,7 +189,7 @@ class For(Stmt):
@tvm_ffi.register_object("ir.While")
class While(Stmt):
- """While node.
+ """Repeat the body while the condition evaluates to true.
Parameters
----------
@@ -267,6 +266,23 @@ class RegionStmt(Stmt):
evaluated outside the body-parameter scope. Direct construction and JSON
serialization support result variables; structured script syntax currently
supports only result-free regions.
+
+ Parameters
+ ----------
+ op : Op or str
+ Registered region operator or its name.
+ args : sequence[Expr]
+ Operands evaluated in the enclosing scope.
+ body_params : sequence[Var]
+ Variables bound for the body according to the operator's contract.
+ attrs : DictAttrs or Mapping or None
+ Operator-specific attributes; None gives an empty map.
+ body : Stmt or Sequence[Stmt]
+ Region body, normalized to SeqStmt.
+ result_vars : sequence[Var], optional
+ Variables defined for subsequent statements, empty by default.
+ loc : Location, optional
+ The source location of the region.
"""
op: Op
@@ -332,7 +348,7 @@ class SeqStmt(Stmt):
@tvm_ffi.register_object("ir.If")
class If(Stmt):
- """If node.
+ """Execute one branch according to a boolean condition.
Parameters
----------
@@ -371,7 +387,7 @@ class If(Stmt):
@tvm_ffi.register_object("ir.Evaluate")
class Evaluate(Stmt):
- """Evaluate node.
+ """Evaluate an expression for its effects and discard its result.
Parameters
----------
@@ -391,10 +407,12 @@ class Evaluate(Stmt):
@tvm_ffi.register_object("ir.Break")
class Break(Stmt):
- """Break node.
+ """Exit the innermost enclosing loop.
Parameters
----------
+ loc : Location, optional
+ The source location of the statement.
"""
def __init__(self, loc: Location = UNKNOWN_LOC) -> None:
@@ -403,7 +421,7 @@ class Break(Stmt):
@tvm_ffi.register_object("ir.Return")
class Return(Stmt):
- """Return node.
+ """Exit the enclosing function with a return value.
Parameters
----------
@@ -423,10 +441,12 @@ class Return(Stmt):
@tvm_ffi.register_object("ir.Continue")
class Continue(Stmt):
- """Continue node.
+ """Skip to the next iteration of the innermost enclosing loop.
Parameters
----------
+ loc : Location, optional
+ The source location of the statement.
"""
def __init__(self, loc: Location = UNKNOWN_LOC) -> None:
diff --git a/tirx/python/tvm_tirx/ir/transform/core.py
b/tirx/python/tvm_tirx/ir/transform/core.py
index a013f7d0fe..58f003195f 100644
--- a/tirx/python/tvm_tirx/ir/transform/core.py
+++ b/tirx/python/tvm_tirx/ir/transform/core.py
@@ -194,44 +194,8 @@ def pass_instrument(pi_cls=None):
Parameters
----------
- pi_class : class
+ pi_cls : class
Instrument class. See example below.
-
- Examples
- --------
-
- .. code-block:: python
-
- @tvm_tirx.ir.transform.pass_instrument
- class SkipPass:
- def __init__(self, skip_pass_name):
- self.skip_pass_name = skip_pass_name
-
- # Uncomment to customize
- # def enter_pass_ctx(self):
- # pass
-
- # Uncomment to customize
- # def exit_pass_ctx(self):
- # pass
-
- # If pass name contains keyword, skip it by return False. (return
True: not skip)
- def should_run(self, mod, pass_info)
- if self.skip_pass_name in pass_info.name:
- return False
- return True
-
- # Uncomment to customize
- # def run_before_pass(self, mod, pass_info):
- # pass
-
- # Uncomment to customize
- # def run_after_pass(self, mod, pass_info):
- # pass
-
- skip_annotate = SkipPass("AnnotateSpans")
- with tvm_tirx.ir.transform.PassContext(instruments=[skip_annotate]):
- tvm.compile(mod, "llvm")
"""
def create_pass_instrument(pi_cls):
@@ -247,19 +211,18 @@ def pass_instrument(pi_cls=None):
@tvm_ffi.register_object("transform.PassContext")
class PassContext(tvm_ffi.Object):
- """The basis where a TVM optimization/analysis runs on.
- Each pass context contains a number of auxiliary information that is used
- to help an optimization pass. Such information includes the error reporter
- to record the errors of during the optimization, etc.
+ """Scope optimization levels, pass selection, instruments, and pass
options.
+ Parameters
+ ----------
opt_level : Optional[int]
- The optimization level of this pass.
+ Maximum optimization level enabled in sequential pipelines (default 2).
required_pass : Optional[Union[List[str], Set[str], Tuple[str]]]
- The list of passes that are required by a certain pass.
+ Pass names to enable regardless of their optimization level.
disabled_pass : Optional[Union[List[str], Set[str], Tuple[str]]]
- The list of passes that are disabled.
+ Pass names to skip in sequential pipelines, even if required.
instruments : Optional[Sequence[PassInstrument]]
The list of pass instrument implementations.
@@ -335,9 +298,12 @@ class PassContext(tvm_ffi.Object):
@tvm_ffi.register_object("transform.Pass")
class Pass(tvm_ffi.Object):
- """The base class of all passes. All methods here are just simple wrappers
- that are implemented in the backend. They are defined for users to
- conveniently interact with the base class.
+ """Base interface for callable IRModule transformations and analyses.
+
+ Obtain a concrete pass from a pass factory or :func:`module_pass`, then
+ call ``after = pass_object(before)``. The :attr:`info` property exposes
+ its name and optimization level. The base class defines no transformation
+ of its own.
"""
__slots__ = ("__dict__",)
@@ -365,12 +331,13 @@ class Pass(tvm_ffi.Object):
@tvm_ffi.register_object("transform.ModulePass")
class ModulePass(Pass):
- """A pass that works on tvm_tirx.IRModule. Users don't need to interact
with
- this class directly. Instead, a module pass should be created through
- `module_pass`, because the design of the `module_pass` API is flexible
- enough to handle the creation of a module pass in different manners. In
- addition, all members of a module pass can be accessed from the base class.
- The same rule applies to FunctionPass as well.
+ """A pass whose callback receives and returns a complete IRModule.
+
+ Create instances with :func:`module_pass`, which accepts a function or a
+ class defining ``transform_module``. Call ``after = pass_object(before)``
+ to execute the callback under the current PassContext. Its callback,
+ rather than this wrapper class, defines how the IR changes; see
+ :func:`module_pass` for a TVMScript before/after example.
"""
@@ -390,8 +357,7 @@ class Sequential(Pass):
opt_level : Optional[int]
The optimization level of this sequential pass.
The opt_level of a default sequential pass is set to 0.
- Note that some of the passes within the Sequantial may still not be
executed
- if their opt_level is higher than the provided opt_level.
+ Child passes use the active PassContext's optimization level.
name : Optional[str]
The name of the sequential pass.
@@ -468,43 +434,47 @@ def module_pass(pass_func=None, opt_level=None,
name=None):
Examples
--------
- The following code block decorates a module pass class.
+ Decorate a class for a configurable pass, or a function for a fixed pass.
+ Both forms below simplify a TIRx module.
- .. code-block:: python
+ .. code:: python
- @tvm_tirx.ir.transform.module_pass
- class CustomPipeline:
- def __init__(self, enable_fold):
- self.enable_fold = enable_fold
- self.const_fold = relax.transform.FoldConstant()
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
- def transform_module(self, mod, ctx):
- if self.enable_fold:
- mod = self.const_fold(mod, ctx)
- return mod
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), "int32")):
+ for i in T.serial(4):
+ A[i] = i // 4
- # create an instance of customized pipeline
- pipeline = CustomPipeline(enable_fold=False)
- assert isinstance(pipeline, transform.ModulePass)
- # run the pipeline.
- output_module = pipeline(input_module)
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), "int32")):
+ for i in T.serial(4):
+ A[i] = 0
- The following code creates a module pass by decorating
- a user defined transform function.
+ @tvm.ir.transform.module_pass(opt_level=0)
+ class CustomPipeline:
+ def __init__(self, enable_simplify=True):
+ self.enable_simplify = enable_simplify
- .. code-block:: python
+ def transform_module(self, mod, ctx):
+ if self.enable_simplify:
+ mod = tvm.tirx.transform.StmtSimplify()(mod)
+ return mod
- @tvm_tirx.ir.transform.module_pass(opt_level=2)
- def transform(mod, ctx):
- return relax.transform.FoldConstant(mod)
+ pipeline = CustomPipeline()
+ tvm.ir.assert_structural_equal(pipeline(Before), Expected)
+ tvm.ir.assert_structural_equal(CustomPipeline(False)(Before), Before)
- module_pass = transform
- assert isinstance(module_pass, transform.ModulePass)
- assert module_pass.info.opt_level == 2
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def simplify(mod, ctx):
+ return tvm.tirx.transform.StmtSimplify()(mod)
- # Given a module m, the optimization could be invoked as the follwoing:
- updated_mod = module_pass(m)
- # Now a function abs should be added to the module m.
+ tvm.ir.assert_structural_equal(simplify(Before), Expected)
"""
if opt_level is None:
raise ValueError("Please provide opt_level for the module pass.")
@@ -525,7 +495,7 @@ def module_pass(pass_func=None, opt_level=None, name=None):
def PrintIR(header=""):
- """A special trace pass that prints the header and IR.
+ """Log the header and module IR, returning the module unchanged.
Parameters
----------
@@ -533,7 +503,34 @@ def PrintIR(header=""):
The header to be displayed along with the dump.
Returns
+ -------
+ fpass : Pass
+ A module pass with optimization level zero.
+
+ Examples
--------
- The pass
+ The observation pass returns the module unchanged.
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.evaluate(0)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ T.evaluate(0)
+
+ def run_pass():
+ return tvm.ir.transform.PrintIR("before lowering")(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_transform_api.PrintIR(header)
diff --git a/tirx/python/tvm_tirx/ir/transform/instrument.py
b/tirx/python/tvm_tirx/ir/transform/instrument.py
index 16acf8b041..f271dbd9e0 100644
--- a/tirx/python/tvm_tirx/ir/transform/instrument.py
+++ b/tirx/python/tvm_tirx/ir/transform/instrument.py
@@ -25,14 +25,19 @@ from .core import PassInstrument, pass_instrument
class PassTimingInstrument(PassInstrument):
- """A wrapper to create a passes time instrument that implemented in C++"""
+ """Measure time spent in passes without changing their IR.
+
+ Install in ``PassContext(instruments=[PassTimingInstrument()])`` and call
+ :meth:`render` inside that context to read the accumulated profile.
+ """
def __init__(self):
self.__init_handle_by_constructor__(_ffi_instrument_api.MakePassTimingInstrument)
@staticmethod
def render():
- """Retrieve rendered time profile result
+ """Retrieve the rendered pass timing profile.
+
Returns
-------
string : string
@@ -45,7 +50,7 @@ class PassTimingInstrument(PassInstrument):
timing_inst = PassTimingInstrument()
with tvm_tirx.ir.transform.PassContext(instruments=[timing_inst]):
- relax_mod = relax.transform.FuseOps()(relax_mod)
+ mod = tvm_tirx.tirx.transform.StmtSimplify()(mod)
# before exiting the context, get profile results.
profiles = timing_inst.render()
"""
@@ -54,8 +59,23 @@ class PassTimingInstrument(PassInstrument):
@pass_instrument
class PassPrintingInstrument:
- """A pass instrument to print if before or
- print ir after each element of a named pass."""
+ """Print IR before and/or after passes selected by exact pass name.
+
+ Parameters
+ ----------
+ print_before_pass_names : Sequence[str]
+ Pass names whose input IR should be printed.
+ print_after_pass_names : Sequence[str]
+ Pass names whose output IR should be printed.
+
+ Examples
+ --------
+ This observer leaves the IR unchanged::
+
+ observer = PassPrintingInstrument([], ["tirx.StmtSimplify"])
+ with tvm_tirx.ir.transform.PassContext(instruments=[observer]):
+ after = tvm_tirx.tirx.transform.StmtSimplify()(mod)
+ """
def __init__(self, print_before_pass_names, print_after_pass_names):
self.print_before_pass_names = print_before_pass_names
@@ -72,7 +92,11 @@ class PassPrintingInstrument:
@pass_instrument
class PrintAfterAll:
- """Print the name of the pass, the IR, only after passes execute."""
+ """Print pass metadata and output IR after every executed pass.
+
+ Use ``PassContext(instruments=[PrintAfterAll()])`` to observe a pipeline.
+ This instrument does not transform the module.
+ """
def run_after_pass(self, mod, info):
print(f"After Running Pass: {info}")
@@ -81,7 +105,11 @@ class PrintAfterAll:
@pass_instrument
class PrintBeforeAll:
- """Print the name of the pass, the IR, only before passes execute."""
+ """Print pass metadata and input IR before every executed pass.
+
+ Use ``PassContext(instruments=[PrintBeforeAll()])`` to observe a pipeline.
+ This instrument does not transform the module.
+ """
def run_before_pass(self, mod, info):
print(f"Before Running Pass: {info}")
@@ -90,7 +118,27 @@ class PrintBeforeAll:
@pass_instrument
class DumpIR:
- """Dump the IR after the pass runs."""
+ """Write each pass's output TVMScript to a numbered file.
+
+ Files are named ``000_<pass_name>.py``, ``001_<pass_name>.py``, and so on,
+ with unsafe filename characters replaced by underscores. This instrument
+ observes the module without changing it.
+
+ Parameters
+ ----------
+ dump_dir : pathlib.Path or str
+ Output directory, created when the first pass finishes.
+ refresh : bool
+ If true, remove an existing directory only when all of its entries
+ match the numbered dump-file pattern; otherwise preserve it.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ with
tvm_tirx.ir.transform.PassContext(instruments=[DumpIR("pass_ir")]):
+ after = tvm_tirx.tirx.transform.StmtSimplify()(mod)
+ """
def __init__(self, dump_dir: Path | str, refresh: bool = False):
if isinstance(dump_dir, Path):
diff --git a/tirx/python/tvm_tirx/ir/type.py b/tirx/python/tvm_tirx/ir/type.py
index ffba9b5b94..f81555e48a 100644
--- a/tirx/python/tvm_tirx/ir/type.py
+++ b/tirx/python/tvm_tirx/ir/type.py
@@ -110,12 +110,12 @@ class OpaqueType(Type):
@tvm_ffi.register_object("ir.PrimType")
class PrimType(Type):
- """Primitive data type in the low level IR
+ """A primitive scalar or vector data type, including the ``void`` sentinel.
Parameters
----------
- dtype : str
- The runtime data type relates to the primtype.
+ dtype : str or tvm_ffi.dtype
+ Element kind, bit width, and optional vector lane count, such as
``float32x4``.
"""
def __init__(self, dtype):
@@ -186,7 +186,7 @@ class StringType(Type):
@tvm_ffi.register_object("ir.PtrType")
class PtrType(Type):
- """PtrType used in the low-level TIR.
+ """A pointer to an element type in a specified memory scope.
Parameters
----------
@@ -228,10 +228,7 @@ class TupleType(Type):
class FuncType(Type):
"""Function type.
- A function type consists of a list of type parameters to enable
- the definition of generic functions,
- a set of type constraints which we omit for the time being,
- a sequence of argument types, and a return type.
+ Records the argument types in order and a single return type.
Parameters
----------
@@ -274,6 +271,15 @@ def _init_type_api(namespace: str, target_module: str |
ModuleType) -> None:
instances, or generates constructors. Call it again after loading
additional
bindings to fill missing names or replace generated root aliases with
dialect
aliases. Module-local provenance distinguishes generated aliases from
exports.
+
+ Parameters
+ ----------
+ namespace : str
+ Registered type-key namespace to expose, such as ``"tirx"``. Its type
+ aliases take precedence over root ``"ir"`` aliases with the same name.
+ target_module : str or types.ModuleType
+ Module receiving the aliases, or its name in ``sys.modules``. The
+ module must already be loaded when passed by name.
"""
module = sys.modules[target_module] if isinstance(target_module, str) else
target_module
exports = vars(module)
diff --git a/tirx/python/tvm_tirx/tirx/compilation_pipeline.py
b/tirx/python/tvm_tirx/tirx/compilation_pipeline.py
index 0308c9bb0d..3d12253943 100644
--- a/tirx/python/tvm_tirx/tirx/compilation_pipeline.py
+++ b/tirx/python/tvm_tirx/tirx/compilation_pipeline.py
@@ -27,12 +27,39 @@ from tvm_tirx.backend.cuda import transform as
cuda_transform
def default_tir_pipeline(*, target=None):
- """The default pipeline for TIRx primitive functions."""
+ """Return :func:`tirx_pipeline` with its standard lowering stages.
+
+ Parameters
+ ----------
+ target : tvm_tirx.target.Target or None, optional
+ Target to bind before lowering; None preserves existing function
targets.
+
+ Returns
+ -------
+ stages : tuple
+ Lowering pass, host-finalization factory, and device-finalization
factory.
+ """
return tirx_pipeline(target=target)
def tirx_pipeline(*, prepare_only=False, target=None):
- """The TIRX pipeline used in tvm_tirx.tirx.compile"""
+ """Prepare TIRx functions and optionally lower their module-wide ABI.
+
+ Dispatch tile instructions, lower layouts, flatten tensors, simplify
indices
+ and loops, and verify memory accesses.
+
+ Parameters
+ ----------
+ prepare_only : bool, optional
+ Stop before :func:`finalize_tir_pipeline` when true. Defaults to false.
+ target : tvm_tirx.target.Target or None, optional
+ Target to bind before lowering; None preserves existing function
targets.
+
+ Returns
+ -------
+ stages : tuple
+ Lowering pass, host-finalization factory, and device-finalization
factory.
+ """
@tvm_tirx.ir.transform.module_pass(opt_level=0)
def _pipeline(
@@ -124,7 +151,16 @@ def _with_target_binding(pipeline, target, *,
original_mod=None, prepared_mod=No
def finalize_tir_pipeline():
- """Lower module-wide calling conventions after all dialects are
prepared."""
+ """Split host/device functions, legalize low-precision types and lower the
ABI.
+
+ Run after dialect-specific preparation and target binding. Annotates the
entry
+ function, lowers CUDA helpers and creates packed entry wrappers where
required.
+
+ Returns
+ -------
+ passes : tvm_tirx.ir.transform.Sequential
+ Module-wide finalization passes.
+ """
return tvm_tirx.ir.transform.Sequential(
[
tirx.transform.AnnotateEntryFunc(),
@@ -141,7 +177,16 @@ def finalize_tir_pipeline():
def finalize_host_passes(): # pylint: disable=unused-argument
- """The default finalization passes for TIR backend."""
+ """Lower runtime builtins and target intrinsics in host functions.
+
+ Run after ABI lowering and host/device splitting, with a host target on
each
+ function.
+
+ Returns
+ -------
+ passes : tvm_tirx.ir.transform.Sequential
+ Host finalization passes.
+ """
host_pass_list = [
tirx.transform.LowerTVMBuiltin(),
tirx.transform.LowerIntrin(),
@@ -150,7 +195,16 @@ def finalize_host_passes(): # pylint:
disable=unused-argument
def finalize_device_passes(): # pylint: disable=unused-argument
- """The default finalization passes for TIR backend."""
+ """Simplify device statements and lower target intrinsics.
+
+ Run after host/device splitting, with a target on each device function.
+ Intrinsic expansion uses the target's registered lowering rules.
+
+ Returns
+ -------
+ passes : tvm_tirx.ir.transform.Sequential
+ Device finalization passes.
+ """
device_pass_list = [
tirx.transform.StmtSimplify(),
tirx.transform.LowerIntrin(),
@@ -163,27 +217,40 @@ PIPELINE_MAP = {"default": default_tir_pipeline, "tirx":
tirx_pipeline}
def register_tir_pipeline(name: str, pipeline_factory) -> None:
- """Register a factory returning lowering, host, and device finalization
passes.
+ """Register a factory returning lowering and host/device finalization
stages.
Factories accepting an explicit ``target`` keyword own target binding. They
must lower S-TIR thread bindings before binding the complete module and
specializing helpers. Other factories receive inputs with targets already
bound; they cannot classify calls inside unlowered S-TIR thread bindings.
+
+ Parameters
+ ----------
+ name : str
+ Name used to select the pipeline; replaces any existing registration.
+ pipeline_factory : Callable
+ Factory returning the lowering pass and host/device finalization
factories.
"""
PIPELINE_MAP[name] = pipeline_factory
def get_tir_pipeline(name: str | None = None, *, target=None, **kwargs):
- """Get pre-build pipeline by name
+ """Construct a registered pipeline by name.
Parameters
----------
- name : Optional[str]
- Name of the pipeline
-
- target : Optional[tvm_tirx.target.Target]
+ name : str or None, optional
+ Registered pipeline name. None selects ``"default"``.
+ target : tvm_tirx.target.Target or None, optional
Target to bind during lowering, after thread launches become explicit.
+ **kwargs
+ Additional keyword arguments forwarded to the registered factory.
+
+ Returns
+ -------
+ stages : tuple
+ Lowering pass, host-finalization factory, and device-finalization
factory.
"""
if name is None:
name = "default"
@@ -204,11 +271,15 @@ _LEGACY_TARGET_PREPARERS = []
def register_default_tir_pipeline_selector(selector) -> None:
- """Register a callback returning a pipeline tuple or None for (module,
target).
+ """Register a callback selecting a pipeline for a module and target.
- Selectors accepting ``target_to_bind`` pass it to the selected pipeline
- factory. Existing two-argument selectors inspect the bound module, which is
- reused by the selected pipeline without binding it a second time.
+ Parameters
+ ----------
+ selector : Callable
+ Callback accepting ``(module, target)`` and returning a pipeline tuple
or
+ None to try the next selector. Callbacks accepting ``target_to_bind``
pass
+ it to the selected factory; other callbacks inspect the bound module,
+ which is reused without binding it a second time.
"""
_DEFAULT_PIPELINE_SELECTORS.append(selector)
@@ -219,7 +290,23 @@ def get_default_tir_pipeline(
*,
target_to_bind: tvm_tirx.target.Target | None = None,
):
- """Get the default pipeline for the input module and target."""
+ """Select the default pipeline for the input module and target.
+
+ Parameters
+ ----------
+ target : tvm_tirx.target.Target or None
+ Target supplied to dialect selectors to choose the pipeline.
+ mod : tvm_tirx.ir.IRModule or None, optional
+ Module inspected by dialect selectors. None uses the default
registration.
+ target_to_bind : tvm_tirx.target.Target or None, optional
+ Target to bind during lowering; None preserves existing function
targets.
+
+ Returns
+ -------
+ stages : tuple
+ Selected lowering pass, host-finalization factory, and
device-finalization
+ factory. Falls back to the ``"default"`` registration if no selector
matches.
+ """
original_mod = mod
prepared_mod = None
remaining_target = target_to_bind
diff --git a/tirx/python/tvm_tirx/tirx/function.py
b/tirx/python/tvm_tirx/tirx/function.py
index 8b06d4678d..b1da7e8ca1 100644
--- a/tirx/python/tvm_tirx/tirx/function.py
+++ b/tirx/python/tvm_tirx/tirx/function.py
@@ -31,24 +31,23 @@ from . import _ffi_api
@tvm_ffi.register_object("tirx.Function")
class Function(BaseFunc, Scriptable):
- """A function declaration expression.
+ """A TIRx function definition or declaration.
+
+ The function's ``ty`` is a FuncType derived from parameter types and
+ ``ret_type``. A present body is normalized to a SeqStmt.
Parameters
----------
- params: List[Union[tvm_tirx.tirx.Var, tvm_tirx.tirx.Var]]
- List of input parameters to the function.
-
- body: Stmt | Sequence[Stmt] | None
- The body of the function, or None for a declaration.
-
- ret_type: tvm_tirx.ir.Type
- The return type annotation of the function.
-
- attrs: Optional[tvm.Attrs]
- Attributes of the function, can be None
-
+ params : list[tvm_tirx.ir.Var]
+ Input parameters, in calling order.
+ body : Stmt or Sequence[Stmt] or None
+ The function body, or None for a declaration without an implementation.
+ ret_type : tvm_tirx.ir.Type, optional
+ The return type; defaults to MissingType until supplied or inferred.
+ attrs : tvm_tirx.ir.DictAttrs, optional
+ Function attributes, empty by default.
loc : Location, optional
- The location of this itervar in the source code.
+ The source location of the function.
"""
body: SeqStmt | None
@@ -83,7 +82,7 @@ class Function(BaseFunc, Scriptable):
) # type: ignore
def with_body(self, new_body: Stmt | Sequence[Stmt] | None, loc: Location
= UNKNOWN_LOC):
- """Create a new Function with the same set signatures but a new body.
+ """Create a function with the same parameters, return type, and
attributes but a new body.
Parameters
----------
@@ -91,7 +90,7 @@ class Function(BaseFunc, Scriptable):
The new body, or None for a declaration.
loc : Location, optional
- The location of this itervar in the source code.
+ The source location of the function.
Returns
-------
diff --git a/tirx/python/tvm_tirx/tirx/index_map.py
b/tirx/python/tvm_tirx/tirx/index_map.py
index 9c7323b813..851082b994 100644
--- a/tirx/python/tvm_tirx/tirx/index_map.py
+++ b/tirx/python/tvm_tirx/tirx/index_map.py
@@ -33,20 +33,20 @@ from . import _ffi_api
@tvm_ffi.register_object("tirx.IndexMap")
class IndexMap(tvm_ffi.Object):
- """A mapping from multi-dimensional indices to another set of
multi-dimensional indices
+ """Map input index variables to expressions for output indices.
+
+ Use :meth:`from_func` to derive a map from a Python callable, or supply the
+ input variables and output expressions directly.
Parameters
----------
- initial_indices : List[Var]
- Variables representing the indices prior to remapping.
- final_indices : List[Expr]
- Expressions defining the indices after remapping.
- inverse_index_map : Union[Callable, Optional[IndexMap]]
- The optional pre-defined inverse index map.
- When this is defined, IndexMap::Inverse will return the pre-defined
inverse index map.
- Otherwise, the inverse index map will be computed on the fly.
- It is the user's responsibility to ensure the correctness of the
pre-defined inverse
- index map.
+ initial_indices : list[Var]
+ Bound variables representing the input coordinates, in dimension order.
+ final_indices : list[Expr]
+ Output coordinate expressions in terms of the input variables.
+ inverse_index_map : IndexMap or Callable or None
+ A supplied inverse, or None to compute the inverse when requested. The
+ caller is responsible for correctness of a supplied inverse.
"""
initial_indices: list[Var]
diff --git a/tirx/python/tvm_tirx/tirx/layout.py
b/tirx/python/tvm_tirx/tirx/layout.py
index 219e15e400..2e4606f3bc 100644
--- a/tirx/python/tvm_tirx/tirx/layout.py
+++ b/tirx/python/tvm_tirx/tirx/layout.py
@@ -32,6 +32,13 @@ from .thread_scope import ThreadScope
@tvm_ffi.register_object("tirx.Layout")
class Layout(tvm_ffi.Object):
+ """Base class for mappings from logical indices to named physical axes.
+
+ Concrete layouts describe memory offsets and thread ownership.
+ :meth:`apply` returns coordinates keyed by axis name; :meth:`size`
describes
+ the input size and :meth:`span` the covered output span.
+ """
+
def check_well_formed(self) -> None:
"""Check whether the layout is well formed.
@@ -280,7 +287,22 @@ class _AxisMeta(type(Enum)):
class Axis(Enum, metaclass=_AxisMeta, type_key="tirx.Axis", init=False):
- """Layout axis wrapper."""
+ """A named physical coordinate used by a layout.
+
+ Use ``Axis.get(name)`` or a registered attribute such as ``Axis.m``. Axis
+ metadata identifies memory or thread coordinates and any owning backend;
+ ``value @ axis`` attaches a stride or offset to that coordinate.
+
+ Examples
+ --------
+ .. code:: python
+
+ from tvm_tirx.script import tirx as T
+
+ memory = T.Axis.get("m") # Also available as T.Axis.m.
+ layout = T.TileLayout(T.S[8 : 2 @ memory])
+ assert int(layout.apply(3)["m"]) == 6
+ """
@property
def name(self) -> str:
@@ -504,7 +526,17 @@ def _to_offset_expr(x: _OffsetExprLike) -> _OffsetExpr:
@tvm_ffi.register_object("tirx.Iter")
class Iter(tvm_ffi.Object):
- """A memory layout that tiles data across devices."""
+ """One extent-and-stride factor contributing to a layout axis.
+
+ Parameters
+ ----------
+ extent : Expr
+ Number of positions in this factor.
+ stride : Expr
+ Increment on the output axis for each position.
+ axis : Axis or str
+ Output coordinate receiving the contribution, resolved by name for
strings.
+ """
extent: Expr
stride: Expr
@@ -543,7 +575,29 @@ def _spec_to_iters(pair) -> list:
@tvm_ffi.register_object("tirx.TileLayout")
class TileLayout(Layout):
- """A memory layout that tiles data across devices."""
+ """Map logical indices to memory and ownership coordinates using strided
factors.
+
+ ``shard`` holds the factors for distinct logical elements; ``replica``
holds
+ factors for copies of those elements. ``offset`` gives an origin per axis.
+ Use :meth:`from_iters` to construct these fields directly.
+
+ Parameters
+ ----------
+ spec : layout specification
+ Compose ``S[shape:stride]`` shard factors, optional ``R[shape:stride]``
+ replica factors, and offsets with ``+``. Integer strides use the memory
+ axis ``m``; ``stride @ Axis.name`` selects another axis. Omitted
strides
+ use contiguous row-major memory order.
+
+ Examples
+ --------
+ .. code:: python
+
+ from tvm_tirx.script import tirx as T
+
+ layout = T.TileLayout(T.S[4, 8]) # Contiguous 4-by-8 tile.
+ assert int(layout.apply(1, 2)["m"]) == 10
+ """
shard: list[Iter]
replicate: list[Iter]
@@ -640,12 +694,34 @@ class TileLayout(Layout):
@tvm_ffi.register_object("tirx.ComposeLayout")
class ComposeLayout(Layout):
- """A memory layout that swizzles a tile layout.
-
- ``per_element`` / ``swizzle_len`` / ``atom_len`` / ``swizzle_inner`` carry
the
- swizzle (formerly the standalone ``SwizzleLayout``); ``tile_layout`` is the
- tiled memory map the swizzle is applied to. A bare swizzle is a
- ``ComposeLayout`` over a trivial identity tile.
+ """Apply an XOR swizzle to the memory coordinate of a tile layout.
+
+ All three bit counts must be nonnegative.
+
+ Parameters
+ ----------
+ per_element : int
+ Number of low address bits left unchanged by the swizzle.
+ swizzle_len : int
+ Number of bits in each XOR group; zero disables swizzling.
+ atom_len : int
+ Separation between the bit groups; must be at least ``swizzle_len``.
+ tile_layout : TileLayout
+ The underlying mapping to memory coordinates.
+ swizzle_inner : bool, optional
+ If true, XOR the higher group into the lower group; otherwise XOR the
+ lower group into the higher group. Defaults to true.
+
+ Examples
+ --------
+ .. code:: python
+
+ from tvm_tirx.script import tirx as T
+
+ tile = T.TileLayout(T.S[8, 8])
+ layout = T.ComposeLayout(0, 2, 3, tile)
+ # Memory offset 10 becomes 10 ^ ((10 >> 3) & 3) = 11.
+ assert int(layout.apply(1, 2)["m"]) == 11
"""
per_element: int
diff --git a/tirx/python/tvm_tirx/tirx/tensor.py
b/tirx/python/tvm_tirx/tirx/tensor.py
index b67b3019b8..a7ced4fdf1 100644
--- a/tirx/python/tvm_tirx/tirx/tensor.py
+++ b/tirx/python/tvm_tirx/tirx/tensor.py
@@ -44,12 +44,32 @@ def _tensor_type_field(name):
@tvm_ffi.register_object("tirx.TensorType")
class TensorType(Type):
- """The structural type carried by an ordinary TIRx tensor variable.
-
- Declared expression methods receive the tensor variable as ``expr`` after
- the type receiver. Access through ``expr.method`` binds both receivers.
- Read-only expression properties receive the same type and operand through
- the getters in ``__expr_properties__``.
+ """The shape, element type, and storage contract of a TIRx tensor
expression.
+
+ Tensor variables expose the type's properties and operations, such as
+ ``tensor.shape`` and ``tensor.view(...)``. Strides and element offsets are
+ described by ``layout``.
+
+ Parameters
+ ----------
+ shape : sequence[Expr] or Expr
+ Logical dimension extents; a single extent denotes a one-dimensional
tensor.
+ dtype : PrimType or str
+ Primitive element type.
+ scope : MemoryScope or str, optional
+ Storage scope, defaulting to global memory.
+ align : int, optional
+ Data-pointer alignment in bytes; zero selects the runtime default.
+ layout : Layout or str, optional
+ Explicit storage mapping. ``"default"`` creates a contiguous row-major
+ layout; None leaves the layout unspecified.
+ data : Expr, optional
+ Keyword-only pointer expression whose PtrType supplies the storage
scope.
+ The pointer itself is not stored in TensorType.
+ loc : Location, optional
+ The source location of the type.
+ data_alignment : int, optional
+ Keyword-only spelling of ``align``; both must agree when explicitly
set.
"""
dtype: PrimType
diff --git a/tirx/python/tvm_tirx/tirx/thread_scope.py
b/tirx/python/tvm_tirx/tirx/thread_scope.py
index d9b6ac862f..b435507926 100644
--- a/tirx/python/tvm_tirx/tirx/thread_scope.py
+++ b/tirx/python/tvm_tirx/tirx/thread_scope.py
@@ -28,8 +28,10 @@ from . import _ffi_api
class ThreadScope(Enum, Scriptable, type_key="tirx.ThreadScope", init=False):
"""A registered thread scope, queried by its qualified backend.scope name.
- Backend and hierarchy rank are enum attributes. The registration index
- does not determine the hierarchy.
+ Obtain a tag with ``ThreadScope.get("cuda.warp")``, for example; unknown
+ names raise an error. ``backend`` identifies its owner and ``rank`` orders
+ the hierarchy, with smaller ranks denoting wider scopes. The registration
+ index does not determine the hierarchy.
"""
@classmethod
diff --git a/tirx/python/tvm_tirx/tirx/tile/dispatch.py
b/tirx/python/tvm_tirx/tirx/tile/dispatch.py
index 48e476d700..7a89f806cc 100644
--- a/tirx/python/tvm_tirx/tirx/tile/dispatch.py
+++ b/tirx/python/tvm_tirx/tirx/tile/dispatch.py
@@ -33,7 +33,11 @@ from ..thread_scope import ThreadScope
@register_object("tirx.DispatchContext")
class DispatchContext(tvm_ffi.Object, Scriptable):
- """DispatchContext node.
+ """Execution context supplied to a tile instruction's lowering callback.
+
+ Carries the target, execution scope and variable bounds used to choose
+ an implementation, together with accumulated allocations and setup
+ statements. Instances are normally supplied by
:func:`~tvm_tirx.tirx.transform.TileDispatch`.
Parameters
----------
@@ -44,16 +48,29 @@ class DispatchContext(tvm_ffi.Object, Scriptable):
The execution scope of the dispatch context.
launch_params : Dict[str, Tuple[Var, Expr]]
- The launch parameters of the dispatch context.
+ Hardware launch tag to ``(index variable, extent)`` pairs.
var_range_map : Dict[Var, Range]
A map from loop variables to their ranges.
+ alloc_only : bool, optional
+ Whether dispatch is collecting allocations rather than emitting the
+ instruction body. Defaults to False.
+
callbacks : Dict[str, tvm_ffi.Object]
- The callbacks of the dispatch context.
+ Named callbacks used to request allocations and initialization
statements.
shared_state : Dict[str, tvm_ffi.Object]
Shared state persisting across dispatch calls within a single lowering
pass.
+
+ inter, intra : Dict[str, list], optional
+ Per-axis ``[extent, offset]`` pairs for the execution context's
inter-team
+ and intra-team views. An empty map means no execution-context tracking
+ is available.
+
+ scope_kind : str, optional
+ Backend-local execution-scope name, such as CUDA's ``kernel``, ``cta``
+ or ``warp``. Empty when unspecified.
"""
target: Target
diff --git a/tirx/python/tvm_tirx/tirx/transform/__init__.py
b/tirx/python/tvm_tirx/tirx/transform/__init__.py
index 498be623a7..0850184e6e 100644
--- a/tirx/python/tvm_tirx/tirx/transform/__init__.py
+++ b/tirx/python/tvm_tirx/tirx/transform/__init__.py
@@ -15,7 +15,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Namespace of all TIR transformations"""
+"""Transformation passes for lowering and optimizing tensor-level IR.
+
+Pass factories return callable passes: ``after = PassName()(before)``, where
+``before`` is an IRModule. Each pass example includes a complete input module,
+expected output and invocation.
+"""
# pylint: disable=wildcard-import, invalid-name
from .function_pass import function_pass, FunctionPass
diff --git a/tirx/python/tvm_tirx/tirx/transform/function_pass.py
b/tirx/python/tvm_tirx/tirx/transform/function_pass.py
index 9b26a303b6..3dfa4bfe15 100644
--- a/tirx/python/tvm_tirx/tirx/transform/function_pass.py
+++ b/tirx/python/tvm_tirx/tirx/transform/function_pass.py
@@ -29,8 +29,10 @@ from . import _ffi_api
@tvm_ffi.register_object("tirx.FunctionPass")
class FunctionPass(Pass):
- """A pass that works on each :py:func:`tvm_tirx.tirx.Function` in a
module. A function
- pass class should be created through
py:func:`tvm_tirx.tirx.transform.function_pass`.
+ """A pass that transforms each :py:class:`tvm_tirx.tirx.Function` in a
module.
+
+ Create instances or subclasses with :py:func:`function_pass`, which
supplies
+ pass metadata and adapts a Python callback to the pass manager.
"""
@@ -72,69 +74,27 @@ def function_pass(
opt_level: int | None = None,
name: str | None = None,
) -> Callable | FunctionPass:
- """Decorate a function pass.
+ """Create a TIRx function pass from a callback or class.
- This function returns a callback when pass_func
- is provided. Otherwise, it returns the created function pass using the
- given optimization function.
+ The callback receives ``(function, module, context)`` and returns the
+ transformed Function. A decorated class instead implements
+ ``transform_function(self, function, module, context)``. With no
+ ``pass_func``, returns a decorator for either form.
Parameters
----------
- pass_func : Optional[Callable[
- (tvm_tirx.tirx.Function, IRModule, PassContext) ->
tvm_tirx.tirx.Function]]
- The transformation function or class.
-
+ pass_func : Callable or type, optional
+ Transformation callback, or class implementing ``transform_function``.
opt_level : int
- The optimization level of this module pass.
-
- name : Optional[str]
- The name of the function pass. The name could be empty. In this case,
the
- name of the optimization function will be used as the pass name.
+ Required optimization level for the pass.
+ name : str, optional
+ Pass name. Defaults to the callback or class name.
Returns
-------
- create_function_pass : Union[Callable, FunctionPass]
-
- A decorator will be returned if pass_func is not provided,
- otherwise return the decorated result.
- The returned decorator has two behaviors depending on the input:
- A new FunctionPass will be returned when we decorate a pass function.
- A new FunctionPass class will be returned when we decorate a class
type.
-
- Examples
- --------
- The following code block decorates a function pass class.
-
- .. code-block:: python
-
- @tvm_tirx.tirx.transform.function_pass(opt_level=1)
- class TestReplaceFunc:
- def __init__(self, new_func):
- self.new_func = new_func
-
- def transform_function(self, func, mod, ctx):
- # just for demo purposes
- # transform func to new_func
- return self.new_func
-
- The following code creates a function pass by decorating
- a user defined transform function.
-
- .. code-block:: python
-
- @tvm_tirx.tirx.transform.function_pass(opt_level=2)
- def transform(func, mod, ctx):
- # my transformations here.
- return func
-
- function_pass = transform
- assert isinstance(function_pass, transform.FunctionPass)
- assert function_pass.info.opt_level == 2
-
- # Given a module m, the optimization could be invoked as the following:
- updated_mod = function_pass(m)
- # Now constant folding should have been applied to every function in
- # the provided module m. And the updated module will be returned.
+ result : FunctionPass, type, or Callable
+ A pass instance for a callback, a pass class for a class, or a
decorator
+ when ``pass_func`` is omitted.
"""
if opt_level is None:
diff --git a/tirx/python/tvm_tirx/tirx/transform/transform.py
b/tirx/python/tvm_tirx/tirx/transform/transform.py
index 49213a9482..2d1135f0a7 100644
--- a/tirx/python/tvm_tirx/tirx/transform/transform.py
+++ b/tirx/python/tvm_tirx/tirx/transform/transform.py
@@ -24,14 +24,42 @@ from .function_pass import function_pass
def DecorateDeviceScope():
- """Wrap a function body in a device scope region."""
+ """Wrap each function body in a ``tirx.device_scope`` region.
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(x: T.int32):
+ T.evaluate(x)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(x: T.int32):
+ with T.device_scope():
+ T.evaluate(x)
+
+ def run_pass():
+ return tr.DecorateDeviceScope()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
+ """
return _ffi_api.DecorateDeviceScope()
def Apply(ftransform):
"""Apply ftransform to each function in the Module.
- This function is a thin wrapper around
tvm_tirx.tirx.transform.function_pass
+ This function is a thin wrapper around :py:func:`function_pass`
Parameters
----------
@@ -42,6 +70,32 @@ def Apply(ftransform):
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ pass
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ T.func_attr({"example": 1})
+
+ def run_pass():
+ return tr.Apply(lambda f: f.with_attr("example", 1))(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
# pylint: disable=unused-argument
@@ -69,59 +123,209 @@ def VectorizeLoop(enable_vectorize: bool = True):
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.float32), B: T.Tensor((4,),
T.float32)):
+ for i in T.vectorized(0, 4):
+ B[i] = A[i]
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.float32), B: T.Tensor((4,),
T.float32)):
+ B[0:4] = A[T.Ramp(0, 1, 4)]
+
+ def run_pass():
+ return tr.VectorizeLoop()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.VectorizeLoop(enable_vectorize) # type: ignore
def StorageRewrite():
- """Rewrite storage allocation pattern.
+ """Plan storage by moving allocations and reusing compatible storage.
- Moves the allocation to outer most possible scope.
- Trying to share space between allocations to make
- a static allocation plan when possible.
+ Hoists allocations where valid and shares storage when lifetimes and
+ memory-scope allocation properties permit it.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(output: T.Tensor((2,), "float32")):
+ A = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32")])
+ for i in T.serial(4):
+ A[i] = T.float32(i)
+ output[0] = A[3]
+ B = T.alloc_tensor(ty_args=[T.Tensor((4,), "float32")])
+ for i in T.serial(4):
+ B[i] = T.float32(i + 1)
+ output[1] = B[3]
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(output: T.Tensor((2,), T.float32)):
+ A = T.alloc_tensor(ty_args=[T.Tensor((4,), T.float32)])
+ for i in range(4):
+ A[i] = T.Cast("float32", i)
+ output[0] = A[3]
+ A_1 = T.decl_tensor(A.data, ty_args=[T.Tensor((4,),
T.float32)])
+ for i_1 in range(4):
+ A_1[i_1] = T.Cast("float32", i_1 + 1)
+ output[1] = A_1[3]
+
+ def run_pass():
+ return tr.StorageRewrite()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.StorageRewrite() # type: ignore
def PointerValueTypeRewrite():
- """
- Rewrite the pointer content type of arguments, as well as Alloc internal
to the function to use
- the most frequently accessed type for load/store to avoid pointer casting
in backend when
- possible.
+ """Rewrite pointer and tensor element types to match their accesses.
+
+ Adjusts argument and allocation types, shapes and accesses where possible
+ to reduce backend casts, including packing scalar storage into vectors.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((16,), T.float32,
scope="local")])
+ A[0:4] = T.Broadcast(T.float32(1.0), 4)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((4,), T.float32x4,
scope="local")])
+ A[0] = T.Broadcast(T.float32(1.0), 4)
+
+ def run_pass():
+ return tr.PointerValueTypeRewrite()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.PointerValueTypeRewrite() # type: ignore
def UnrollLoop():
- """Unroll the constant loop marked by unroll.
+ """Unroll loops with constant extents marked ``T.unroll``.
- This pass also automatically attach pragma unroll tag to loops which meets
the standard.
+ The ``tirx.UnrollLoop`` PassContext configuration also controls automatic
+ unrolling and whether to expand loops or retain unroll annotations.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((2,), T.int32)):
+ for i in T.unroll(0, 2):
+ A[i] = 1
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((2,), T.int32)):
+ A[0] = 1
+ A[1] = 1
+
+ def run_pass():
+ return tr.UnrollLoop()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.UnrollLoop() # type: ignore
def RemoveNoOp():
- """Remove No Op from the Stmt.
+ """Remove effect-free statements and provably empty control flow.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.int32)):
+ T.evaluate(1)
+ A[0] = 1
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.int32)):
+ A[0] = 1
+
+ def run_pass():
+ return tr.RemoveNoOp()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.RemoveNoOp() # type: ignore
@@ -133,23 +337,89 @@ def HoistIf():
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass.
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.int32), n: T.int32):
+ for i in range(4):
+ if n > 0:
+ A[i] = 1
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.int32), n: T.int32):
+ if 0 < n:
+ for i in range(4):
+ A[i] = 1
+
+ def run_pass():
+ return tr.HoistIf()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.HoistIf()
def BF16ComputeLegalize():
- """Legalize bf16 compute Ops.
+ """Promote unsupported BF16 computations and internal temporaries to
float32.
+
+ Inserts conversions at BF16 boundaries. Leaves functions unchanged when
+ their CUDA target reports native BF16 support.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((1,), "bfloat16")])
+ A[0] = T.bfloat16(1)
+ T.evaluate(A[0] + T.bfloat16(2))
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((1,), T.float32)])
+ A[0] = T.float32(1.0)
+ T.evaluate(A[0] + T.float32(2.0))
+
+ def run_pass():
+ return tr.BF16ComputeLegalize()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.BF16ComputeLegalize() # type: ignore
def FP8ComputeLegalize(promote_dtype: str = "float32"):
- """Legalize fp8 compute Ops.
+ """Promote unsupported FP8 computations and internal temporaries.
+
+ Inserts conversions at FP8 boundaries. Leaves functions unchanged when
+ their CUDA target reports native FP8 support.
Parameters
----------
@@ -160,28 +430,118 @@ def FP8ComputeLegalize(promote_dtype: str = "float32"):
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((1,), "float8_e5m2")])
+ A[0] = T.float8_e5m2(1)
+ T.evaluate(A[0] + T.float8_e5m2(2))
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((1,), T.float16)])
+ A[0] = T.float16(1.0)
+ T.evaluate(A[0] + T.float16(2.0))
+
+ def run_pass():
+ return tr.FP8ComputeLegalize("float16")(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.FP8ComputeLegalize(promote_dtype) # type: ignore
def BF16StorageLegalize():
- """Legalize bf16 storage types to u16.
+ """Represent BF16 storage as uint16 while preserving its bit patterns.
+
+ Run after compute legalization. Functions whose CUDA target reports
+ native BF16 support are unchanged.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+ After packed-API lowering, pointer storage uses integer elements.
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Ptr(T.bfloat16)):
+ T.evaluate(0)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Ptr(T.uint16)):
+ pass
+
+ def run_pass():
+ return tr.BF16StorageLegalize()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.BF16StorageLegalize() # type: ignore
def FP8StorageLegalize():
- """Legalize fp8 storage types to u8.
+ """Represent FP8 storage as uint8 while preserving its bit patterns.
+
+ Run after compute legalization. Functions whose CUDA target reports
+ native FP8 support are unchanged.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+ After packed-API lowering, pointer storage uses integer elements.
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Ptr(T.float8_e5m2)):
+ T.evaluate(0)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Ptr(T.uint8)):
+ pass
+
+ def run_pass():
+ return tr.FP8StorageLegalize()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.FP8StorageLegalize() # type: ignore
@@ -193,6 +553,35 @@ def CommonSubexprElim():
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(B: T.Tensor((2,), T.int32), x: T.int32, y: T.int32):
+ B[0] = x + y
+ B[1] = x + y
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(B: T.Tensor((2,), T.int32), x: T.int32, y: T.int32):
+ cse_v1: T.let[T.int32] = x + y
+ B[0] = cse_v1
+ B[1] = cse_v1
+
+ def run_pass():
+ return tr.CommonSubexprElim()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.CommonSubexprElim() # type: ignore
@@ -204,6 +593,32 @@ def RenormalizeSplitPattern():
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.int32), i: T.int32):
+ A[0] = i % 128 // 32
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.int32), i: T.int32):
+ A[0] = i // 32 % 4
+
+ def run_pass():
+ return tr.RenormalizeSplitPattern()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.RenormalizeSplitPattern() # type: ignore
@@ -215,23 +630,94 @@ def StmtSimplify():
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.int32)):
+ for i in range(4):
+ A[i] = i // 4
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.int32)):
+ for i in range(4):
+ A[i] = 0
+
+ def run_pass():
+ return tr.StmtSimplify()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.StmtSimplify() # type: ignore
def LowerThreadAllreduce():
- """Lower cross-thread reductions to target-specific communication.
+ """Lower ``gpu_thread_allreduce`` to target-specific communication.
+
+ Requires a function ``target`` attribute and explicit thread extents.
+ Selects warp shuffles or shared-memory reductions as supported.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass.
+
+ Examples
+ --------
+ A single participating thread reduces to a direct store.
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32)):
+ T.func_attr({"target": T.target({"kind": "cuda", "arch":
"sm_75"})})
+ R = T.alloc_tensor(ty_args=[T.Tensor((1,), T.float32,
scope="local")])
+ with T.launch_thread("threadIdx.x", 1) as tx:
+ T.gpu_thread_allreduce(
+ T.Lambda([T.float32, T.float32], lambda arg0, arg1:
(arg0 + arg1,)),
+ (T.float32(0.0),),
+ (A[tx],),
+ T.bool(True),
+ (R[0],),
+ (tx,),
+ )
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.float32)):
+ T.func_attr({"target": T.target({"kind": "cuda", "arch":
"sm_75"})})
+ R = T.alloc_tensor(ty_args=[T.Tensor((1,), T.float32,
scope="local")])
+ with T.launch_thread("threadIdx.x", 1) as tx:
+ R[0] = A[tx]
+
+ def run_pass():
+ return tr.LowerThreadAllreduce()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.LowerThreadAllreduce()
def ConvertSSA():
- """Convert an IRModule to be SSA form.
+ """Give every variable definition in an IRModule a distinct identity.
This pass handles cases where the same `tirx.Var` appears in
multiple functions within the same module. For example, after
@@ -244,35 +730,104 @@ def ConvertSSA():
fpass : tvm_tirx.ir.transform.Pass
The result pass
+ Examples
+ --------
+ The input deliberately reuses a variable and disables the SSA check.
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ v = T.dynamic("v", "int32")
+
+ @I.ir_module(check_well_formed=False)
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.bind(16, var=v)
+ T.evaluate(v)
+ T.bind(32, var=v)
+ T.evaluate(v)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ v: T.let[T.int32] = 16
+ T.evaluate(v)
+ v_1: T.let[T.int32] = 32
+ T.evaluate(v_1)
+
+ def run_pass():
+ return tr.ConvertSSA()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.ConvertSSA() # type: ignore
def MakePackedAPI():
- """Transform the Functions in the module to a packed func API.
-
- Prior to this pass, the Function may have parameters annotated with
- `TensorType`. This pass consumes those annotations to generate
- arguments that implement the packed based TVM FFI API.
-
- For static shapes, the tensor shape and layout are used to
- generate runtime checks on the corresponding member variables in
- the user-provided `DLTensor*` or `tvm_tirx.runtime.tensor` argument.
(e.g. A
- Function that accepts a tensor of shape `[16,32]` validates that
- the `DLTensor::shape` array is `[16,32]`.)
-
- For dynamic Tensors, in which one or more of these `TensorType` fields
- use `tirx.Var` that are not defined by other Function
- parameters, these are instead used to define the variables based on
- the corresponding `DLTensor` members. (e.g. A Function that accepts a
- tensor of shape `[tirx.Var("n", "int64"), tirx.Var("m", "int64")]`,
- when passed a `DLTensor` of shape `[16, 32]`, will define `n = 16` and
- `m = 32`, based on the argument's shape.
+ """Expose exported functions through the TVM FFI packed-call ABI.
+
+ Requires a ``target`` attribute with a host target on functions with a
+ ``global_symbol`` and the default calling convention. Private functions,
+ functions with other calling conventions, and functions whose target has
+ no host are left unchanged.
+
+ Tensor parameter types generate runtime checks for shape, dtype and
+ layout. Symbolic dimensions not defined by other arguments are bound from
+ the incoming DLTensor metadata. For example, a ``T.Tensor((16, 32),
+ "float32")`` argument is checked against that shape, while symbolic
+ dimensions are read from the argument's shape. The resulting function
+ uses the host target and an exported ``__tvm_ffi_`` symbol.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+ An exported no-argument function illustrates the packed-call ABI.
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function
+ def main():
+ T.func_attr({"target": T.target("llvm", host="llvm")})
+ T.evaluate(0)
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(
+ self_handle: T.Ptr(T.void),
+ args: T.Ptr(T.void),
+ num_args: T.int32,
+ result: T.Ptr(T.void),
+ ) -> T.int32:
+ T.func_attr({"global_symbol": "__tvm_ffi_main"})
+ T.func_attr({"calling_conv": 1, "target": T.target("llvm")})
+ assert num_args == 0, (
+ "TypeError",
+ ["Expected ", "0", " arguments", " when calling:\\n `",
"main()", "`"],
+ )
+ with T.compute_scope("main_compute_"):
+ pass
+ return 0
+
+ def run_pass():
+ return tr.MakePackedAPI()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.MakePackedAPI() # type: ignore
@@ -288,45 +843,203 @@ def SplitHostDevice():
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(n: T.int32):
+ T.func_attr({"target": T.target("cuda", host="llvm")})
+ with T.region("tirx.device_scope", []):
+ T.evaluate(n)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(n: T.int32):
+ T.func_attr(
+ {
+ "target": T.target(
+ {
+ "host": {
+ "keys": ["cpu"],
+ "kind": "llvm",
+ "mtriple": "",
+ "tag": "",
+ },
+ "keys": ["cuda", "gpu"],
+ "kind": "cuda",
+ "max_num_threads": 1024,
+ "tag": "",
+ "thread_warp_size": 32,
+ }
+ )
+ }
+ )
+ T.call_ffi_kernel("main_kernel", n)
+
+ @T.function
+ def main_kernel(n: T.int32):
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "target": T.target("cuda"),
+ "tirx.is_global_func": True,
+ "tirx.kernel_launch_params": [],
+ "tirx.noalias": True,
+ }
+ )
+ T.evaluate(n)
+
+ def run_pass():
+ return tr.SplitHostDevice()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.SplitHostDevice() # type: ignore
def SkipAssert():
- """Skip assert stmt.
+ """Replace assertion statements with no-ops.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.int32), n: T.int32):
+ assert n > 0, ("RuntimeError", ["positive extent required"])
+ A[0] = 1
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.int32), n: T.int32):
+ A[0] = 1
+
+ def run_pass():
+ return tr.SkipAssert()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.SkipAssert() # type: ignore
def LowerTVMBuiltin():
- """Lower tvm builtin intrinsics.
+ """Lower runtime builtins to explicit stack storage and ABI operations.
+
+ For example, packed calls gain argument/result slots. Requires the
+ function target attribute for target-dependent lowering.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.func_attr({"target": T.target("llvm")})
+ T.call_packed("callback", 1)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ T.func_attr({"target": T.target("llvm")})
+ stack_ffi_any: T.let[T.Ptr(T.void)] =
T.stack_alloca("tvm_ffi_any", 2)
+ T.abi_field_set(stack_ffi_any, 0, 13, 1)
+ T.abi_field_set(stack_ffi_any, 0, 14, 0)
+ T.abi_field_set(stack_ffi_any, 0, 15, T.Cast("int64", 1))
+ T.abi_field_set(stack_ffi_any, 1, 13, 0)
+ T.abi_field_set(stack_ffi_any, 1, 14, 0)
+ T.abi_field_set(stack_ffi_any, 1, 15, T.int64(0))
+ T.call_packed_lowered("callback", stack_ffi_any, 0, 1)
+
+ def run_pass():
+ return tr.LowerTVMBuiltin()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.LowerTVMBuiltin() # type: ignore
def LowerIntrin():
- """Lower target specific intrinsic calls.
+ """Lower intrinsic calls using rules registered for the function target.
+
+ Requires a function ``target`` attribute. Also lowers logical tensor
+ pointer operations to address operations.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.int32), x: T.int32):
+ T.func_attr({"target": T.target("llvm")})
+ A[0] = x // 8
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((1,), T.int32), x: T.int32):
+ T.func_attr({"target": T.target("llvm")})
+ A[0] = x >> 3
+
+ def run_pass():
+ return tr.LowerIntrin()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.LowerIntrin() # type: ignore
def NarrowDataType(target_bits: int):
- """Narrow down Expr datatype in stmt to target_bits.
+ """Narrow integer indices and loop variables when bounds prove it safe.
+
+ Keeps wider types when the analyzed range does not fit ``target_bits``.
Parameters
----------
@@ -341,14 +1054,46 @@ def NarrowDataType(target_bits: int):
Note
----
Run this pass after FlattenTensor.
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.float32)):
+ for i in range(T.int64(0), T.int64(4)):
+ A[i] = T.float32(1.0)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.float32)):
+ for i in range(4):
+ A[i] = T.float32(1.0)
+
+ def run_pass():
+ return tr.NarrowDataType(32)(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.NarrowDataType(target_bits) # type: ignore
def ForceNarrowIndexToInt32():
- """Force narrow down indexing expressions and integer tensors to int32
dtype.
+ """Force index expressions to int32 without range proofs.
- The function must not contain S-TIR blocks.
+ Use only when the caller can guarantee representability. Unlike
+ :py:func:`NarrowDataType`, this pass may change overflowing computations.
+ Rejects signed integer tensor parameters and ordinary allocations wider
+ than int32. Explicitly typed local scalars retain their storage type and
+ are cast at narrowed index uses.
Returns
-------
@@ -358,41 +1103,137 @@ def ForceNarrowIndexToInt32():
Note
----
This pass should not be used in default cases.
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.float32)):
+ for i in range(T.int64(0), T.int64(4)):
+ A[i] = T.float32(1.0)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.float32)):
+ for i in range(4):
+ A[i] = T.float32(1.0)
+
+ def run_pass():
+ return tr.ForceNarrowIndexToInt32()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.ForceNarrowIndexToInt32() # type: ignore
def VerifyMemory():
- """Verify if func contains illegal host side direct memory access.
+ """Reject illegal host-side access to device memory.
+
+ Returns the module unchanged on success and raises an error on invalid
+ access. Functions without a target attribute are skipped, as are
+ functions with a non-default calling convention.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+ Valid device memory accesses are checked without changing the IR.
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.float32)):
+ T.func_attr({"target": T.target({"kind": "cuda", "arch":
"sm_75"})})
+ with T.launch_thread("threadIdx.x", 4) as tx:
+ A[tx] = T.float32(1.0)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main(A: T.Tensor((4,), T.float32)):
+ T.func_attr({"target": T.target({"kind": "cuda", "arch":
"sm_75"})})
+ with T.launch_thread("threadIdx.x", 4) as tx:
+ A[tx] = T.float32(1.0)
+
+ def run_pass():
+ return tr.VerifyMemory()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.VerifyMemory() # type: ignore
def FlattenTensor():
- """Flatten the multi-dimensional TensorLoad and TensorStore to single
dimensional
- TensorLoad/TensorStore for the TIR not contains opaque block.
+ """Flatten tensor allocations and accesses to one-dimensional storage.
+
+ Uses shapes, strides and offsets to compute linear indices. Tensor
+ parameters retain their interface types and gain flat views as needed.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((2, 4), T.float32,
scope="local")])
+ A[1, 2] = T.float32(1.0)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((8,), T.float32,
scope="local")])
+ A[6] = T.float32(1.0)
+
+ def run_pass():
+ return tr.FlattenTensor()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.FlattenTensor() # type: ignore
def BindTarget(target):
- """Annotate a Function with a given target.
+ """Bind targets and classify host/device helper functions.
+
+ Preserves existing device targets and updates host information on exported
+ functions. Helpers used on both sides may be cloned for the host.
Device helper calls must be inside explicit launch_thread or device_entry
regions. Lower S-TIR thread bindings before classifying host/device calls.
Parameters
- -------
+ ----------
target : tvm_tirx.target.Target
target
@@ -400,40 +1241,167 @@ def BindTarget(target):
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.evaluate(0)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ T.func_attr({"target": T.target("cuda")})
+
+ def run_pass():
+ return tr.BindTarget(tvm.target.Target("cuda"))(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.BindTarget(target) # type: ignore
def AnnotateEntryFunc():
- """Set a Function as the entry point if it is only function in IRModule.
+ """Infer and mark an unambiguous TIRx entry function.
+
+ Sets ``tirx.is_entry_func`` when the module contains only one function
+ and it is TIRx, or when exactly one function is externally exposed and
+ it is a TIRx function.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ pass
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ T.func_attr({"tirx.is_entry_func": True})
+
+ def run_pass():
+ return tr.AnnotateEntryFunc()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.AnnotateEntryFunc() # type: ignore
def Filter(fcond: Callable):
- """Filter out Functions that does not satisfy the given condition.
- `fcond` should be a function that takes a function and returns boolean.
+ """Remove TIRx functions for which ``fcond(function)`` is false.
+
+ The predicate receives a Function and returns a boolean.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ T.evaluate(0)
+
+ @I.ir_module
+ class Expected:
+ pass
+
+ def run_pass():
+ return tr.Filter(lambda f: False)(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.Filter(fcond) # type: ignore
def TileDispatch():
- """Lower TIRx tile op calls through the active backend dispatch table.
+ """Lower TIRx tile calls using the active backend dispatch table.
+
+ Resolves tile instructions and device launch/execution scopes. Select
+ the backend with a target context; follow with
:py:func:`LowerTensorLayouts`
+ or the backend's own layout-lowering pass.
Returns
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+ Dispatch a tile-wide move into a loop of scalar PTX moves.
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_local((4, 1), "float32",
layout=T.TileLayout(T.S[(4, 1) : (1, 1)]))
+ T.cuda.tile.apply(T.ptx["mov.b32"], A[:, :], T.float32(0),
scope="thread")
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_local((4, 1), "float32",
layout=T.TileLayout(T.S[(4, 1) : (1, 1)]))
+ tile_d = T.decl_tensor(
+ A.data,
+ ty_args=[
+ T.Tensor(
+ (4,),
+ T.float32,
+ scope="local",
+ layout=T.TileLayout(T.S[(4,) : (1,)]),
+ )
+ ],
+ )
+ for repeat_0 in T.unroll(0, 4):
+ T.ptx.mov(tile_d[repeat_0], T.float32(0.0), "b32", "")
+
+ def run_pass():
+ with tvm.target.Target({"kind": "cuda", "arch": "sm_75"}):
+ return tr.TileDispatch()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.TileDispatch() # type: ignore
@@ -449,5 +1417,43 @@ def LowerTensorLayouts():
-------
fpass : tvm_tirx.ir.transform.Pass
The result pass
+
+ Examples
+ --------
+
+ .. code-block:: python
+
+ import tvm_tirx as tvm
+ from tvm_tirx.script import ir as I, tirx as T
+ from tvm_tirx.tirx import transform as tr
+
+ @I.ir_module
+ class Before:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(
+ ty_args=[
+ T.Tensor(
+ (3, 4),
+ T.uint32,
+ scope="local",
+ layout=T.TileLayout(T.S[(3, 4) : (16, 2)]),
+ )
+ ]
+ )
+ A[2, 3] = T.uint32(7)
+
+ @I.ir_module
+ class Expected:
+ @T.function(private=True)
+ def main():
+ A = T.alloc_tensor(ty_args=[T.Tensor((39,), T.uint32,
scope="local")])
+ A[38] = T.uint32(7)
+
+ def run_pass():
+ with tvm.target.Target({"kind": "cuda", "arch": "sm_75"}):
+ return tr.LowerTensorLayouts()(Before)
+
+ tvm.ir.assert_structural_equal(run_pass(), Expected)
"""
return _ffi_api.LowerTensorLayouts() # type: ignore
diff --git a/tirx/python/tvm_tirx/tirx/type.py
b/tirx/python/tvm_tirx/tirx/type.py
index 9e6443645b..d6d44e5006 100644
--- a/tirx/python/tvm_tirx/tirx/type.py
+++ b/tirx/python/tvm_tirx/tirx/type.py
@@ -26,12 +26,16 @@ from . import _ffi_api
@tvm_ffi.register_object("tirx.TensorMapType")
class TensorMapType(Type):
- """TensorMapType used in the low-level TIR.
+ """The type of a tensor-map descriptor used by backend memory operations.
+
+ A pointer to this type represents descriptor storage, rather than a
tensor's
+ element data. The descriptor's shape and layout are defined by the
operations
+ that create it, not by fields on this type.
Parameters
----------
- loc : tvm_tirx.ir.Location
- The loc information.
+ loc : tvm_tirx.ir.Location, optional
+ The source location of the type.
"""
def __init__(self, loc: Location = UNKNOWN_LOC):
