.. Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at .. http://www.apache.org/licenses/LICENSE-2.0 .. Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Tile Primitives =============== .. note:: This page documents the tile-primitive surface and dispatch as it exists in the source today; signatures and variants may change. Tile primitives are the dispatchable, hardware-level operations a TIRx kernel issues — data movement (``copy``, ``copy_async``), matrix multiply (``gemm``, ``gemm_async``), reductions, elementwise math, and a few fused/compose forms. A primitive call is recorded as an **unresolved** ``TilePrimitiveCall`` IR node; the compiler later *dispatches* it — selecting a concrete lowering from the primitive, the execution scope, the operand layouts, the target, and an optional explicit hint — and replaces it with native IR (loops, address arithmetic, synchronization, and backend intrinsics). Calling convention ------------------ Tile primitives are called in TVMScript via ``from tvm.script import tirx as T``, on the injected ``Tx`` namespace. The namespace prefix selects the **cooperation scope**: - ``Tx.(...)`` — unqualified, runs at **thread** scope. - ``Tx.warp.`` / ``Tx.wg.`` (alias ``Tx.warpgroup``) / ``Tx.cta.`` / ``Tx.cluster.`` / ``Tx.thread.`` — bind a wider scope. Every primitive also accepts, besides its operands: ``scope`` (usually set by the namespace), ``workspace: dict[str, Buffer] | None``, ``dispatch: str | None`` (force a named lowering variant), and ``**kwargs`` collected into a ``config`` dict that tunes the chosen lowering. Operands are ``Buffer`` / ``BufferRegion`` values, each carrying a :doc:`TileLayout ` that dispatch reads. Wiring (three layers): the authoritative op list is the C++ registry (``src/tirx/op/tirx.cc``, 29 ops named ``tirx.tile.``); the IR wrapper classes are in ``python/tvm/tirx/operator/tile_primitive/ops.py``; the user-facing ``Tx.*`` builders are in ``python/tvm/tirx/script/builder/tirx.py``. Primitive catalog ----------------- The 29 primitives, grouped. Signatures show the operands plus the common ``workspace``/``dispatch``/``scope``/``**kwargs`` tail (abbreviated ``...``). Data movement ~~~~~~~~~~~~~~ :: copy(dst, src, ...) # synchronous element copy src -> dst copy_async(dst, src, ...) # asynchronous copy (caller commits/waits) permute_layout(dst, src, ...) # rearrange under a different layout (may alias) Matrix multiply ~~~~~~~~~~~~~~~~ :: gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0, ...) # D = alpha*A*B + beta*C (register mma) gemm_async(C, A, B, SFA=None, SFB=None, transA=False, transB=False, accum=False, ...) # async / block-scaled Fill / memset / zero ~~~~~~~~~~~~~~~~~~~~~~ :: fill(dst, value, ...) # fill region with a scalar memset(dst, value, ...) # set all elements to a value zero(dst, src=None, ...) # zero out (in place if src omitted) Cast and elementwise ~~~~~~~~~~~~~~~~~~~~~~ :: cast(dst, src=None, ...) # dtype cast (buffer form) sqrt / exp / exp2(dst, src=None, bias=None, scale=None, ...) reciprocal(dst, src=None, ...) # dst = 1/src silu(dst, src, ...) # dst = src*sigmoid(src) add / sub / mul / fdiv(dst, src1, src2, ...) # element-wise arithmetic maximum / minimum(dst, src1, src2, ...) # element-wise max / min fma(dst, src, scale, bias, ...) # dst = src*scale + bias select(dst, true_value, false_value, pred, scope=None) # dst = pred ? t : f Reductions ~~~~~~~~~~ :: sum / max / min(dst, src, axes=-1, accum=False, ...) # reduce over axes Fused / compose ~~~~~~~~~~~~~~~~ :: binary_reduce(...) # binary op then reduce, fused unary_reduce(...) # unary (with bias/scale) then reduce binary_chain(...) # chain two binary ops reduce_negate(...) # reduce then negate compose_op(...) # frame/context manager to group primitives Dispatch config --------------- A call is materialized as a ``TilePrimitiveCall`` node whose fields carry everything dispatch needs (``python/tvm/tirx/stmt.py``): .. list-table:: :header-rows: 1 :widths: 14 22 64 * - Field - Type - Meaning * - ``op`` - ``tvm.Op`` - primitive identity, e.g. ``tirx.tile.copy_async`` * - ``args`` - ``Array`` - operands (regions / scalars), in the order shown above * - ``workspace`` - ``Map[str, Buffer]`` - pre-allocated scratch buffers * - ``config`` - ``Map[str, Any]`` - open-ended tuning bag (table below) * - ``dispatch`` - ``Optional[str]`` - forced variant name; ``None`` = auto-select * - ``scope`` - ``ExecScope`` - cooperation scope (default ``thread``) ``config`` has **no central schema** — each key is read only by the dispatch variant(s) that need it (via ``config.get(...)``); a key meant for another primitive is simply ignored. Only ``dispatch`` is generic. The keys observed in the CUDA backend, by consumer: .. list-table:: :header-rows: 1 :widths: 18 30 18 34 * - Key - Used by - Type / values - Meaning * - ``dispatch`` - any primitive - variant name (str) - force a lowering variant (also settable via the ``dispatch=`` kwarg) * - ``vec_len`` - ``copy`` / ``copy_async`` (vectorized variants) - int | None - vectorization width for the copy * - ``mbar`` - ``copy_async``: ``tma`` (g2s), ``dsmem`` - mbarrier handle - completion barrier * - ``cta_group`` - ``copy_async``: ``tma``, ``smem→tmem``; ``gemm_async``: ``tcgen05`` - ``1`` | ``2`` - CTA-group; ``2`` routes completion to the cluster * - ``cta_mask`` - ``copy_async``: ``tma`` (g2s) - int | None - multicast CTA mask * - ``cache_hint`` - ``copy_async``: ``tma`` - ``"evict_normal"`` | ``""`` - L2 cache eviction hint * - ``oob`` - ``copy_async``: ``tma`` - ``"zero"`` | ``"nan"`` | None - out-of-bounds fill policy (``nan`` is float-only) * - ``use_tma_reduce`` - ``copy_async``: ``tma`` (s2g) - str (e.g. ``"add"``) | None - TMA store-with-reduction mode * - ``prefetch_tensormap`` - ``copy_async``: ``tma`` - bool - prefetch the tensor map at kernel entry * - ``remote_cta_id`` - ``copy_async``: ``dsmem`` - int | PrimExpr - target CTA for a cross-CTA shared→shared copy * - ``descI`` - ``gemm_async``: ``tcgen05`` - uint32 | None - pre-encoded MMA instruction descriptor * - ``thread_reduce`` - ``reduction``: ``local`` (warp scope) - bool - per-thread shuffle reduction * - ``rounding_mode`` - ``elementwise``: binary ops - ``"rn"`` | ``"rz"`` | ... - FP rounding mode for the packed form Three dispatch inputs are **implicit**, not config keys: the **execution scope** (set by the namespace, then refined against the active thread set tracked through control flow into ``inter``/``intra`` maps and a ``scope_kind``), the **operand layouts** (each ``Buffer.layout``), and the **target** (the dispatch table is keyed by its kind, e.g. ``"cuda"``). Dispatch mechanism ------------------ Pipeline ~~~~~~~~ Dispatch runs in the ``tirx.TilePrimitiveDispatch`` pass — the sole pass inside ``LowerTIRx()``, the first stage of the compile pipeline. The C++ mutator ``TilePrimitiveDispatcher`` walks the IR and, per call: #. resolves the ``(inter, intra)`` execution split for the call's scope from the active set tracked through control flow (``if wg_id == ...``, ``warp_id``, ``T.ptx.elect_sync()``); #. builds a ``DispatchContext`` carrying ``target``, scope, launch params, value ranges, and the encoded ``inter``/``intra`` + ``scope_kind``; #. invokes the global FFI hook ``tirx.f_op_dispatcher`` (Python) with the call and context, which returns a ``PrimFunc``; #. splices that ``PrimFunc`` body in place of the call and drains side-effect callbacks (private allocs, device/host init statements). If any ``TilePrimitiveCall`` survives lowering, a verifier makes it a fatal error. Selection (``run_dispatch``) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ The Python dispatcher holds a table ``_DISPATCH_TABLE`` keyed by ``(Op, target_kind)``. Each entry is a list of *cases*, registered by backends via ``@register_dispatch(op_name, target_kind, variant=..., priority=..., when=[preds])``. ``run_dispatch(op_call, sctx)``: #. ``key = (op_call.op, sctx.target.kind.name)``; look up cases. None → error. #. If ``op_call.dispatch`` is set, **filter** to that variant (error if unknown). #. Sort cases by ``(-priority, variant)`` — highest priority first. #. For each case, evaluate its predicates; if any fails, record the reason and continue. If all pass, run the impl; on success return its ``PrimFunc``. #. An impl may still **decline** by raising ``DispatchFail`` (e.g. a hardware constraint found while emitting) — the search continues. #. If every variant is rejected, raise a ``RuntimeError`` listing each variant's rejection reason. So dispatch is **keyed by (primitive, target)**, then a **priority-ordered, predicate-guarded** case list, with an optional ``dispatch=`` override. Two recurring predicate helpers: ``validate_copy_op`` (both operands have a layout, equal dtype, equal non-unit extents) and ``_all_threads_active`` (the exec scope is full — ``laneid`` spans 32, etc., none of it narrowed by an enclosing ``if``), so a partial-warp copy is rejected rather than mis-lowered. Dispatch by primitive --------------------- Each page below documents one primitive's dispatch in detail — the variants, how each is selected, the algorithm it runs, the IR it emits, and when it declines. .. toctree:: :maxdepth: 1 tile_primitives/copy tile_primitives/copy_async tile_primitives/gemm tile_primitives/gemm_async tile_primitives/elementwise tile_primitives/reduction tile_primitives/permute_layout See also -------- - :doc:`layout` — the ``TileLayout`` model dispatch reads from operands. - :doc:`overview` — execution scope, tensor layout, and tile primitive dispatch as the three core constructs.