tvm.script.ir_builder

tvm.script.ir_builder is a generic IR builder for TVM.

class tvm.script.ir_builder.IRBuilder

A dialect-agnostic IRBuilder that constructs any IR of TVM.

Examples

An idiomatic use of this class is to put this inside the with-scope, call dialect-specific methods accordingly. Upon exiting the scope.

from tvm.script.ir_builder import tirx as T
from tvm.script.ir_builder import IRBuilder

with IRBuilder() as builder:
    with T.prim_func(...):  # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
                        # to `builder`'s stack of frames
        buffer = T.match_buffer(...)

return builder.get()        # returns the constructed IR, i.e. tirx.PrimFunc

static current() → IRBuilder

Get the current IRBuilder put in the with-scope.

Returns:: builder – The current IRBuilder.
Return type:: IRBuilder

static is_in_scope() → bool

See if the current thread-local scope has an IRBuilder.

Returns:: Whether the current thread-local scope has an IRBuilder
Return type:: bool

get() → Object: Get the constructed IR.

static name(s: str, v: Any) → Any

Set the name of an object.

Parameters:

s (str) – The name of the object.
v (Any) – The object to name.

Returns:

v – The same object with the name set.

Return type:

Any

static name_many(s: list[str], vs: list[Any]) → list[Any]

Set the name of a list of objects.

Parameters:

s (List[str]) – The names of the objects.
vs (List[Any]) – The objects to name.

Returns:

vs – The same objects with the names set.

Return type:

List[Any]

tvm.script.ir_builder.ir

Package tvm.script.ir_builder.ir

class tvm.script.ir_builder.ir.IRModuleFrame(global_vars, functions, attrs, global_infos)

tvm.script.ir_builder.ir.decl_function(func_name: str, func_signature: BaseFunc) → GlobalVar

Declare a Function without given the specific function implementation.

Parameters:

func_name (str) – The function unique name.
func_signature (BaseFunc) – A Function w/o body, which used to specify the function signature (i.e. func params and func return type/shape).

Note

It is usually used in cross-function call. And we can specify the function by DefFunction

Returns:: gv – The corresponding GlobalVar.
Return type:: GlobalVar

tvm.script.ir_builder.ir.def_function(func_name: str, func: BaseFunc) → None: Define the function which is declared before. :param func_name: The function unique name. :type func_name: str :param func: The given function implementation :type func: BaseFunc

tvm.script.ir_builder.ir.ir_module() → IRModuleFrame: Start a ir_module frame. :returns: frame – The constructed frame. :rtype: IRModuleFrame

tvm.script.ir_builder.ir.module_attrs(attrs: dict[str, Object], allow_overwrite=False) → None: Specify the attrs of the ir_module frame. :param attrs: The module attrs. :type attrs: Dict[str, Object] :param allow_overwrite: Whether allow overwrite the existing attrs. :type allow_overwrite: bool

tvm.script.ir_builder.ir.module_get_attr(attr_key: str) → Object | None

Get the specified attr of the ir_module frame. :param attr_key: The key of the attr to be retrieved. :type attr_key: str

Returns:: attr – The specified module attr or None if not found.
Return type:: Optional[Object]

tvm.script.ir_builder.ir.module_set_attr(attr_key: str, attr_value: Object | None, allow_overwrite: bool = False) → None: Set the specified attr of the ir_module frame. :param attr_key: The key of the attr to be set. :type attr_key: str :param attr_value: The value of the attr to be set. :type attr_value: Optional[Object] :param allow_overwrite: Whether allow overwrite the existing attr. :type allow_overwrite: bool

tvm.script.ir_builder.ir.module_global_infos(global_infos: dict[str, list[GlobalInfo]]) → None: Specify the global infos of the ir_module frame. :param global_infos: The module global infos. :type global_infos: Dict[str, List[GlobalInfo]]

tvm.script.ir_builder.ir.lookup_vdevice(target_kind: str | None = None, device_index: int = -1) → VDevice

Retrieve a virtual device from the globalinfo vdevice list. :param target_kind: The target device kind, for example ‘llvm’ or ‘cuda’. :type target_kind: str :param device_index: The virtual device index. :type device_index: int

Returns:: res – The result virtual device.
Return type:: VDevice

tvm.script.ir_builder.ir.vdevice(target=None, vdevice_id: int = 0, memory_scope: str = 'global') → VDevice

Create a virtual device global info. :param target: The target. :param vdevice_id: The virtual device index. :type vdevice_id: int :param memory_scope: The memory scope, default is “global” :type memory_scope: str

Returns:: res – The result virtual device.
Return type:: VDevice

tvm.script.ir_builder.ir.dummy_global_info() → DummyGlobalInfo: Create a dummy global info expression. :returns: res – The result dummy global info. :rtype: DummyGlobalInfo

tvm.script.ir_builder.relax

Package tvm.script.ir_builder.relax

tvm.script.ir_builder.relax.distributed

Package tvm.script.ir_builder.relax.distributed

tvm.script.ir_builder.tirx

Package tvm.script.ir_builder.tirx