Class IRModule

python/tvm/ir/module.py:31–283

IRModule that holds functions and type definitions. IRModule is the basic unit for all IR transformations across the stack. Parameters ---------- functions: Optional[dict]. Map of global var to BaseFunc

@tvm_ffi.register_object("ir.IRModule")
class IRModule(Node, Scriptable):
    """IRModule that holds functions and type definitions.

    IRModule is the basic unit for all IR transformations across the stack.

    Parameters
    ----------
    functions: Optional[dict].
        Map of global var to BaseFunc
    """

    def __init__(self, functions=None, attrs=None, global_infos=None):
        if functions is None:
            functions = {}
        elif isinstance(functions, dict):
            mapped_funcs = {}
            for k, v in functions.items():
                if isinstance(k, str):
                    k = _expr.GlobalVar(k)
                if not isinstance(k, _expr.GlobalVar):
                    raise TypeError("Expect functions to be Dict[GlobalVar, Function]")
                mapped_funcs[k] = v
            functions = mapped_funcs

        attrs = None if not attrs else attrs
        if attrs is not None:
            attrs = tvm.ir.make_node("ir.DictAttrs", **attrs)
        if global_infos is None:
            global_infos = {}
        self.__init_handle_by_constructor__(
            _ffi_api.IRModule,
            functions,
            attrs,
            global_infos,
        )
        self.pyfuncs = {}

    def clone(self) -> "IRModule":
        return _ffi_api.Module_Clone(self)

    def functions_items(self):
        """Get items in self.functions.items() in alphabetical order.

        Returns
        -------
        items: List[Tuple[GlobalVar, Function]]
            The functions items.
        """
        items = list(self.functions.items())
        items.sort(key=lambda item: str(item[0].name_hint))
        return items

    def __setitem__(self, var, val):
        """Add a mapping to the module.

        Parameters
        ---------
        var: GlobalVar