Tile Primitive Dispatch

Tile Primitive Dispatch#

This chapter documents how unresolved TilePrimitiveCall nodes are selected and lowered. Kernel authors should start with the Tile Primitives programming guide; extension authors can find the callable registration interfaces in Tile Dispatch Extension API.

Implementation surfaces#

The authoritative operation list is the C++ registry (src/tirx/op/tirx.cc, with operations named tirx.tile.<name>). IR wrapper classes live in python/tvm/tirx/operator/tile_primitive/ops.py. Raw TilePrimitiveCall constructors live in python/tvm/tirx/script/builder/tirx.py, while the validated authoring facade is in python/tvm/tirx/script/tile.py. Both Python construction surfaces produce the same IR node type.

Dispatch pipeline#

Dispatch runs in the tirx.TilePrimitiveDispatch pass, the first phase of LowerTIRx(), before layout and execution-scope cleanup. The C++ mutator TilePrimitiveDispatcher walks the IR and, for each call:

  1. resolves the (inter, intra) execution split for the call’s scope from the active set tracked through control flow (if wg_id == ..., warp_id, and Tx.cuda.elect_sync());

  2. builds a DispatchContext carrying the target, scope, launch parameters, value ranges, and encoded inter/intra maps plus scope_kind;

  3. invokes the global FFI hook tirx.f_op_dispatcher with the call and context, which returns a PrimFunc;

  4. splices that PrimFunc body in place of the call and drains side-effect callbacks for private allocations and device or host initialization.

If a TilePrimitiveCall survives lowering, the verifier reports a fatal error.

Variant selection#

The Python dispatcher holds a table keyed by (Op, target_kind). Backends register each case with register_dispatch, including its variant name, priority, predicates, and implementation. run_dispatch(op_call, sctx):

  1. looks up the primitive and target pair;

  2. filters to op_call.dispatch when the caller explicitly requests a variant;

  3. sorts candidates by descending priority and then by variant name;

  4. evaluates each candidate’s predicates and runs the first implementation that accepts the call; a predicate exception is recorded as a rejection reason;

  5. continues searching when an implementation raises DispatchFail or another exception; and

  6. reports every rejection reason when no candidate accepts the call. If an implementation raised an unexpected exception, the final RuntimeError is chained from the last such exception so its traceback is retained.

Dispatch is therefore target-specific, priority-ordered, and predicate-guarded, with an optional dispatch= override. Common predicates validate matching operand shapes and layouts or require a complete active thread group, preventing a partial warp or warpgroup from being lowered by an implementation that needs full participation.

Variants by primitive#

The following chapters document the registered variants, their selection conditions, the IR they emit, and the conditions under which they decline. Their URLs remain under tile_primitives for compatibility with existing links.

Contents