Function ir_module

python/tvm/script/parser/ir/entry.py:32–154

The parsing method for ir module, by using `@ir_module` as decorator. Parameters ---------- mod : Type The class to be parsed as ir module. check_well_formed : bool Whether to check well-formedness during parsing. Returns ------- ir_module : IRModule

(
    mod: type | None = None, check_well_formed: bool = True, s_tir: bool = False
)

# appear as a decorator by itself or to have optional arguments
# like @I.ir_module(check_well_formed=False)
def ir_module(
    mod: type | None = None, check_well_formed: bool = True, s_tir: bool = False
) -> IRModule:
    """The parsing method for ir module, by using `@ir_module` as decorator.

    Parameters
    ----------
    mod : Type
        The class to be parsed as ir module.

    check_well_formed : bool
        Whether to check well-formedness during parsing.

    Returns
    -------
    ir_module : IRModule
        The parsed ir module.
    """

    # Capture stack outside wrapper (wrapper adds to the stack)
    outer_stack = inspect.stack()

    def decorator_wrapper(mod):
        if not inspect.isclass(mod):
            raise TypeError(f"Expect a class, but got: {mod}")

        # Check BasePyModule inheritance
        base_py_module_inherited = any(base.__name__ == "BasePyModule" for base in mod.__bases__)

        extra_vars = utils.inspect_class_capture(mod)
        # Resolve closure variables hidden by PEP 563 (annotation-only names)
        utils.resolve_closure_vars(mod, extra_vars, outer_stack)
        m = parse(mod, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)

        if base_py_module_inherited:
            # Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
            # because tvm.script is loaded before tvm.relax during tvm initialization.
            from tvm.relax.base_py_module import BasePyModule
            from tvm.relax.expr import ExternFunc  # pylint: disable=import-outside-toplevel

            # Collect pyfunc methods
            pyfunc_methods = [
                name
                for name, attr in mod.__dict__.items()
                if hasattr(attr, "dispatch_token") and attr.dispatch_token == "pyfunc"
            ]

            mod._pyfunc_methods = pyfunc_methods

            # Create ExternFunc nodes

            for method_name in pyfunc_methods:
                try:
                    existing_gvars = [
                        global_var
                        for global_var in m.get_global_vars()
                        if global_var.name_hint == method_name
                    ]