Function visit_tvm_declare_function

python/tvm/script/parser/ir/parser.py:144–171

Declare a Python function as an ExternFunc in the IRModule.

(self: Parser, node: doc.FunctionDef)

@dispatch.register(token="pyfunc", type_name="tvm_declare_function")
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
    """Declare a Python function as an ExternFunc in the IRModule."""
    # Check if Python functions are allowed in this context
    # We need to check if we're in a class that inherits from BasePyModule
    current_class = self._get_current_class_context()
    if current_class and not self._is_base_py_module_context():
        self.report_error(
            node,
            "@I.pyfunc are only allowed in classes that inherit from BasePyModule. "
            f"Class '{current_class}' does not inherit from BasePyModule.",
        )

    # 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 import ExternFunc  # pylint: disable=import-outside-toplevel

    # Create ExternFunc with proper attributes for Python functions
    func = ExternFunc(node.name)
    func = func.with_attr("is_pyfunc", True)
    func = func.with_attr("function_type", "python")
    func = func.with_attr("python_function_name", node.name)

    # Add placeholder attributes that will be filled in later
    func = func.with_attr("python_source", f"# Source will be filled for {node.name}")
    func = func.with_attr("python_packed_func", None)  # Will be filled in entry.py

    # Store the function name for later retrieval
    return I.decl_function(node.name, func)


@dispatch.register(token="pyfunc", type_name="FunctionDef")