Tile Dispatch Extension API

Tile Dispatch Extension API#

Raw node constructors#

tvm.tirx.script.builder.tirx constructs the same TilePrimitiveCall nodes as the validated Tx.tile facade. This surface is intended for parser, builder, and extension authors; kernel code should normally use the tile authoring API. Its public builder spelling is tvm.script.ir_builder.tirx.tile:

from tvm.script.ir_builder import tirx as Tx_builder

Tx_builder.tile.copy(dst, src)

Each callable tile operation has a same-name raw constructor. The authoring page lists the public facade without duplicating the full operation list here.

class ScopedOp(fn)#

Internal callable wrapper that supplies a default thread scope or binds a scope selected through a namespace.

class ScopeNamespace(scope, label)#

Internal resolver that binds the named execution scope to a tile operation.

Dispatch registration#

class tvm.tirx.operator.tile_primitive.DispatchContext(target: Target, exec_scope: ExecScope, launch_params: dict[str, IterVar], var_range_map: dict[Var, Range], alloc_only: bool = False, callbacks: dict[str, Object] = {}, shared_state: dict[str, Object] = {}, inter: dict[str, list] | None = None, intra: dict[str, list] | None = None, scope_kind: str = '')

DispatchContext node.

Parameters:

  • target (Target) – The target of the dispatch context.

  • exec_scope (ExecScope) – The execution scope of the dispatch context.

  • launch_params (Dict[str, Expr]) – The launch parameters of the dispatch context.

  • var_range_map (Dict[tirx.Var, Range]) – A map from loop variables to their ranges.

  • callbacks (Dict[str, Object]) – The callbacks of the dispatch context.

  • shared_state (Dict[str, Object]) – Shared state persisting across dispatch calls within a single lowering pass.

add_alloc_buffer(buffer: Var) → None
Add an allocated buffer to the dispatch context.

Can be called only if alloc_only is True. The buffer will be added to the workspace of operator (the key in the workspace is the buffer name).

Parameters:

buffer (Buffer) – The buffer to be added.

add_init_stmt(stmt: Stmt, host: bool = False) → None
Add an initialization statement to the dispatch context.

Device initialization statements is only allowed if alloc_only is True. Host initialization statements will be ignored if alloc_only is True. The statements will be added to the beginning of the kernel.

Parameters:

  • stmt (Stmt) – The initialization statement to be added.

  • host (bool) – Whether the statement is a host statement. If True, the statement will be added to the host code (before the kernel). If False, the statement will be added to the kernel body (at the beginning of the kernel).

add_post_buffer_def_stmt(buffer: Var, stmt: Stmt) → None

Add a statement to be inserted after a buffer’s definition (DeclBuffer/AllocBuffer).

Parameters:

  • buffer (Buffer) – The buffer whose definition scope the statement should appear in.

  • stmt (Stmt) – The statement to be inserted.

cache_get(key: str) → Object | None

Look up a cached value by key.

Parameters:

key (str) – Cache key (built by the caller from construction parameters).

Returns:

The cached value, or None on miss.

Return type:

Optional[Object]

cache_set(key: str, value: Object) → None

Store a value in the cross-dispatch cache.

Parameters:

  • key (str) – Cache key (built by the caller from construction parameters).

  • value (Object) – The object to cache (e.g. a Buffer or tirx.Var).

is_cuda() → bool

Check if the target is CUDA.

is_trn() → bool

Check if the target is Trainium.

is_target(name: str) → bool

Check if the target kind matches name.

tvm.tirx.operator.tile_primitive.fail(reason: str) → None

Helper for schedule variants to explain why they decline to handle the op.

tvm.tirx.operator.tile_primitive.list_registered_schedules() → dict[str, dict[str, list[str]]]

Return a mapping: op_name -> target_kind -> [variant names].

tvm.tirx.operator.tile_primitive.predicate(name: str, fn: Callable[[TilePrimitiveCall, DispatchContext], Any], **kwargs) → Predicate

Wrap a callable into a named predicate.

tvm.tirx.operator.tile_primitive.register_dispatch(op_name: str, target_kind: str, *, variant: str, priority: int = 0, when: list[Predicate] | None = None)

Decorator to add a dispatch case for an op/target pair.

Cases with higher priority run earlier. When list predicates must all pass. The impl must return a PrimFunc on success, and must NOT return None. To decline handling, raise fail(“reason”) (or DispatchFail).

See Tile Primitive Dispatch for the selection algorithm and the registered CUDA and Trainium variants.

Contents