Function visit_tvm_declare_function

python/tvm/tirx/script/parser/parser.py:871–910

The function declaration step for tirx Parameters ---------- self : Parser The visiting parser. node : doc.Return The doc AST return node.

(self: Parser, node: doc.FunctionDef)

@dispatch.register(token="tirx", type_name="tvm_declare_function")
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
    """The function declaration step for tirx

    Parameters
    ----------
    self : Parser
        The visiting parser.

    node : doc.Return
        The doc AST return node.
    """

    supplied_annotation = self.function_annotations
    func_annotation = supplied_annotation.get(node.name, {})

    ret_type = None
    with self.var_table.with_frame():
        if node.returns is not None:
            ret_type = self.eval_expr(node.returns)
            if callable(ret_type):
                ret_type = PrimType(ret_type().dtype)

        arg_annotations = []
        for arg in node.args.args:
            if arg.annotation is None:
                self.report_error(arg, "Type annotation required for function parameters.")
            try:
                ann = self.eval_expr(arg.annotation)
                if callable(ann):
                    ann = ann()
            except Exception:  # pylint: disable=broad-except
                ann = func_annotation.get(arg.arg, None)
                if ann is None:
                    raise

            IRBuilder.name(arg.arg, ann)
            arg_annotations.append(ann)

    func_signature = tvm.tirx.PrimFunc(arg_annotations, None, ret_type=ret_type)
    return I.decl_function(node.name, func_signature)