tvm_tirx.ir.IRModule

tvm_tirx.ir.IRModule#

class tvm_tirx.ir.IRModule(functions=None, attrs=None, global_infos=None)#

Bases: Node, Scriptable

A collection of global functions and module-level metadata.

Modules are the input and output of compiler passes. Functions can be looked up by their GlobalVar or name.

Parameters:

  • functions (dict[GlobalVar or str, BaseFunc], optional) – Global function definitions. String keys are converted to GlobalVar objects.

  • attrs (dict, optional) – Module attributes, stored as DictAttrs.

  • global_infos (dict[str, list[GlobalInfo]], optional) – Named groups of global information referenced by the IR.

clone()#

Create a module object with the same functions and metadata.

The clone initially shares its function objects and other contents with the original. Adding, replacing, or removing function-map entries on the clone preserves the original module’s entries. Contents are not recursively copied.

Return type:

IRModule

functions_items()#

Get items in self.functions.items() in alphabetical order.

Returns:

items (List[Tuple[GlobalVar, Function]]) – The functions items.

update(other)#

Insert function definitions from another module or dictionary in place.

Definitions are merged using their GlobalVar keys. A dictionary’s string keys create fresh GlobalVar objects; a fresh key with the same name as an existing key raises RuntimeError. To replace an existing definition, reuse its stored GlobalVar or use assignment by name.

Parameters:

other (IRModule or dict[GlobalVar or str, BaseFunc]) – Function definitions to merge into this module. Module attributes and global_infos are not merged.

update_func(var, func)#

Update the function corresponding to a global variable in the module.

Parameters:
update_global_info(name, global_info)#

Update global info in the module

Parameters:

  • name (str) – The name for the global info.

  • global_info (List[GlobalInfo]) – The global info to be updated.

get_global_var(name)#

Get a global variable in the function by name.

Parameters:

name (str) – The name of the global variable.

Returns:

global_var (GlobalVar) – The global variable mapped to name.

Raises:

RuntimeError if we cannot find corresponding global var. –

get_global_vars()#

Collect all global vars defined in this module.

Returns:

global_vars (Array[GlobalVar]) – An array of global vars.

static from_expr(expr, functions=None)#

Construct a module from a standalone expression.

Parameters:

  • expr (Expr) – The starting expression

  • global_funcs (Optional[dict]) – Map of global vars to function definitions

Returns:

mod (Module) – A module containing the passed definitions, where expr is set as the entry point (wrapped in a function if necessary)

get_attr(attr_key)#

Get the IRModule attribute.

Parameters:

attr_key (str) – The attribute key.

Returns:

attr_value (Any) – Attribute value

with_attr(attr_key, attr_value)#

Copy the IRModule and add an attribute to it.

Parameters:

  • attr_key (str) – The attribute key.

  • attr_value (tvm_ffi.Object) – The new attribute value.

Returns:

mod (IRModule) – A new copy of the IRModule with the attribute

without_attr(attr_key)#

Copy the IRModule and remove an attribute key and its associated value. :param attr_key: The attribute key. :type attr_key: str

Returns:

mod (IRModule) – A new copy of the IRModule without the attribute

Parameters:

attr_key (str)

Return type:

IRModule

with_attrs(attr_map)#

Copy the IRModule and add the given attribute map to it. :param attr_map: The attribute map :type attr_map: Union[DictAttrs, Dict[str, tvm_ffi.Object]]

Returns:

mod (IRModule) – A new copy of the IRModule with the attribute

Parameters:

attr_map (DictAttrs | dict[str, Object])

Return type:

IRModule

property attrs#
property functions#
property global_infos#
property global_var_map_#
IRModule.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

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