This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 72da54952cff6f9f47d3d0e88df280e398e06e1e 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):
