tvm_tirx.ir.BaseFunc

tvm_tirx.ir.BaseFunc#

class tvm_tirx.ir.BaseFunc(attrs, loc=<MISSING>, ty=<MISSING>)#

Bases: _CallableExprWithOp

Base class for callable IR function definitions.

attrs stores function metadata as DictAttrs. Attribute update methods return a new function, preserving the original. Concrete function classes define the parameters, body, and function type in ty.

Parameters:
with_attr(attr_key_or_dict, attr_value=None)#

Create a new copy of the function and update the attribute.

Parameters:

  • attr_key_or_dict (Union[str, dict]) – The attribute key to use or a dict containing multiple key value pairs.

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

Returns:

func (BaseFunc) – A new copy of the function

Return type:

BaseFunc

with_attrs(attr_map)#

Copy the function 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:

func (BaseFunc) – A new copy of the function

Parameters:

attr_map (DictAttrs | dict[str, Object])

Return type:

BaseFunc

without_attr(attr_key)#

Copy the function and remove the given attribute key.

Parameters:

attr_key (str) – The attribute key to delete from the attribute pairs.

Returns:

func (BaseFunc) – A new copy of the function

Return type:

BaseFunc

property attrs#
BaseFunc.expr_ty()#

Return this expression’s primitive result type.

BaseFunc.equal(other, loc=UNKNOWN_LOC)#
Parameters:

loc (Location)

BaseFunc.astype(dtype, loc=UNKNOWN_LOC)#
Parameters:

loc (Location)

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

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

CallingConv#

class tvm_tirx.ir.CallingConv(*values)#

Bases: IntEnum

Possible kinds of calling conventions.

DEFAULT = 0#
C_PACKED_FUNC = 1#
DEVICE_KERNEL_LAUNCH = 2#