# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """IRModule that holds the functions and type definitions.""" import tvm_ffi import tvm from tvm.runtime import Object, Scriptable from . import _ffi_api from . import expr as _expr from .attrs import DictAttrs from .base import Node from .function import 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 The global variable. val: Union[Function, Type] The value. """ return self._add(var, val, True) def _add(self, var, val, update=True): if isinstance(val, BaseFunc): if isinstance(var, str): if _ffi_api.Module_ContainGlobalVar(self, var): var = _ffi_api.Module_GetGlobalVar(self, var) else: var = _expr.GlobalVar(var) _ffi_api.Module_Add(self, var, val, update) def __getitem__(self, var): """Lookup a global definition by name or by variable. Parameters ---------- var: Union[String, GlobalVar, GlobalTypeVar] The name or global variable. Returns ------- val: Union[Function, Type] The definition referenced by :code:`var` (either a function or type). """ if isinstance(var, str): return _ffi_api.Module_Lookup_str(self, var) assert isinstance(var, _expr.GlobalVar) return _ffi_api.Module_Lookup(self, var) def __delitem__(self, var: str | _expr.GlobalVar): _ffi_api.Module_Remove(self, var) def __contains__(self, var: str | _expr.GlobalVar) -> bool: return _ffi_api.Module_Contains(self, var) def update(self, other): """Insert functions in another Module to current one. Parameters ---------- other: IRModule The module to merge into the current Module. """ if isinstance(other, dict): other = IRModule(other) return _ffi_api.Module_Update(self, other) def update_func(self, var, func): """Update the function corresponding to a global variable in the module. Parameters ---------- var: GlobalVar The global variable. func: tvm.ir.BaseFunc The function to be inserted. """ return _ffi_api.Module_UpdateFunction(self, var, func) def update_global_info(self, 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. """ return _ffi_api.Module_UpdateGlobalInfo(self, name, global_info) def get_global_var(self, 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 :code:`name`. Raises ------ RuntimeError if we cannot find corresponding global var. """ return _ffi_api.Module_GetGlobalVar(self, name) def get_global_vars(self): """Collect all global vars defined in this module. Returns ------- global_vars: Array[GlobalVar] An array of global vars. """ return _ffi_api.Module_GetGlobalVars(self) @staticmethod def 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) """ funcs = functions if functions is not None else {} return _ffi_api.Module_FromExpr(expr, funcs) def get_attr(self, attr_key): """Get the IRModule attribute. Parameters ---------- attr_key : str The attribute key. Returns ------- attr_value : Any Attribute value """ return _ffi_api.Module_GetAttr(self, attr_key) def with_attr(self, attr_key, attr_value): """Copy the IRModule and add an attribute to it. Parameters ---------- attr_key : str The attribute key. attr_value : Object The new attribute value. Returns ------- mod : IRModule A new copy of the IRModule with the attribute """ return _ffi_api.Module_WithAttr(self, attr_key, attr_value) def without_attr(self, attr_key: str) -> "IRModule": """Copy the IRModule and remove an attribute key and its associated value. Parameters ---------- attr_key : str The attribute key. Returns ------- mod : IRModule A new copy of the IRModule without the attribute """ return _ffi_api.Module_WithoutAttr(self, attr_key) def with_attrs(self, attr_map: DictAttrs | dict[str, Object]) -> "IRModule": """Copy the IRModule and add the given attribute map to it. Parameters ---------- attr_map: Union[DictAttrs, Dict[str, Object]] The attribute map Returns ------- mod : IRModule A new copy of the IRModule with the attribute """ if isinstance(attr_map, tvm.ir.DictAttrs): attr_map = attr_map._dict() return _ffi_api.Module_WithAttrs(self, attr_map)