tvm_tirx.tirx.tile.DispatchContext

tvm_tirx.tirx.tile.DispatchContext#

class tvm_tirx.tirx.tile.DispatchContext(target, thread_scope, launch_params, var_range_map, alloc_only=False, callbacks={}, shared_state={}, inter=None, intra=None, scope_kind='')#

Bases: Object, Scriptable

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 TileDispatch().

Parameters:

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

  • thread_scope (ThreadScope) – The execution scope of the dispatch context.

  • launch_params (Dict[str, Tuple[Var, Expr]]) – 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]) – 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 (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.

  • 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.

property target#
property thread_scope#
property launch_params#
property var_range_map#
property alloc_only#
property callbacks#
property shared_state#
property inter#
property intra#
property scope_kind#
kPrivateAlloc = 'private_alloc'#
kDeviceInitStmt = 'device_init_stmt'#
kHostInitStmt = 'host_init_stmt'#
kPostTensorDefStmt = 'post_tensor_def_stmt'#
add_alloc_tensor(tensor, allocated_addr=None, annotations=None)#
Add an allocated tensor to the dispatch context.

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

Parameters:

  • tensor (Var) – The tensor to be added.

  • allocated_addr (Sequence[PrimExpr], optional) – Physical placement operands on the producing allocation call.

  • annotations (dict, optional) – Annotations on the producing allocation call.

Return type:

None

add_init_stmt(stmt, host=False)#
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).

Return type:

None

add_post_tensor_def_stmt(tensor, stmt)#

Add a statement to be inserted after a tensor’s definition (DeclTensor/AllocTensor).

Parameters:

  • tensor (Var) – The tensor whose definition scope the statement should appear in.

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

Return type:

None

cache_get(key)#

Look up a cached value by key.

Parameters:

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

Returns:

Optional[tvm_ffi.Object] – The cached value, or None on miss.

Return type:

Object | None

cache_set(key, value)#

Store a value in the cross-dispatch cache.

Parameters:

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

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

Return type:

None

is_cuda()#

Check if the target is CUDA.

Return type:

bool

is_trn()#

Check if the target is Trainium.

Return type:

bool

is_target(name)#

Check if the target kind matches name.

Parameters:

name (str)

Return type:

bool

property is_thread: bool#
property is_warp: bool#
property is_warpgroup: bool#
property is_cta: bool#
property is_cluster: bool#
DispatchContext.script(*, name=None, show_meta=False, int_dtype='int32', float_dtype='void', verbose_expr=False, indent_spaces=4, print_line_numbers=False, num_context_lines=-1, path_to_underline=None, path_to_annotate=None, obj_to_underline=None, obj_to_annotate=None, extra_config=None)#

Print TVM IR into TVMScript text format

Parameters:

  • name (Optional[str] = None) – The name of the object

  • show_meta (bool = False) – Whether to print the meta data of the object

  • int_dtype (str = "int32") – The default data type of integer

  • float_dtype (str = "void") – The default data type of float

  • verbose_expr (bool = False) – Whether to print the detailed definition of each variable in the expression

  • indent_spaces (int = 4) – The number of spaces for indentation

  • print_line_numbers (bool = False) – Whether to print line numbers

  • num_context_lines (int = -1) – The number of lines of context to print before and after the line to underline.

  • path_to_underline (Optional[List[AccessPath]] = None) – Object path to be underlined

  • path_to_annotate (Optional[Dict[AccessPath, str]] = None) – Object path to be annotated

  • obj_to_underline (Optional[List[tvm_ffi.Object]] = None) – Object to be underlined

  • obj_to_annotate (Optional[Dict[tvm_ffi.Object, str]] = None) – Object to be annotated

  • extra_config (Optional[dict] = None) – Dialect-specific configuration passed through to PrinterConfig.extra_config, including ir.prefix for the IR namespace alias (default "I"). Keys are conventionally namespaced as “<dialect>.<knob>”, e.g. {"tirx.prefix": "T"}.

Returns:

script (str) – The TVM Script of the given TVM IR

Return type:

str

DispatchContext.show(style=None, *, name=None, show_meta=False, int_dtype='int32', float_dtype='void', verbose_expr=False, indent_spaces=4, print_line_numbers=False, num_context_lines=-1, path_to_underline=None, path_to_annotate=None, obj_to_underline=None, obj_to_annotate=None, extra_config=None)#

Print highlighted TVMScript.

Parameters:

  • style (str, optional) – Pygmentize printing style, auto-detected if None. See tvm_tirx.script.printer.highlight.cprint for more details.

  • name (Optional[str] = None) – The name of the object

  • show_meta (bool = False) – Whether to print the meta data of the object

  • int_dtype (str = "int32") – The default data type of integer

  • float_dtype (str = "void") – The default data type of float

  • verbose_expr (bool = False) – Whether to print the detailed definition of each variable in the expression

  • indent_spaces (int = 4) – The number of spaces for indentation

  • print_line_numbers (bool = False) – Whether to print line numbers

  • num_context_lines (int = -1) – The number of lines of context to print before and after the line to underline.

  • path_to_underline (Optional[List[AccessPath]] = None) – Object path to be underlined

  • path_to_annotate (Optional[Dict[AccessPath, str]] = None) – Object path to be annotated

  • obj_to_underline (Optional[List[tvm_ffi.Object]] = None) – Object to be underlined

  • obj_to_annotate (Optional[Dict[tvm_ffi.Object, str]] = None) – Object to be annotated

  • extra_config (Optional[dict] = None) – Dialect-specific configuration passed through to PrinterConfig.extra_config, including ir.prefix for the IR namespace alias (default "I").

Return type:

None