Function visit_function_def

python/tvm/relax/script/parser/parser.py:200–249

(self: Parser, node: doc.FunctionDef)

@dispatch.register(token="relax", type_name="FunctionDef")
def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
    is_inner_function = self.inside_function
    self.inside_function = True

    # reserve a var for local function
    func_val = self.var_table.get().get(node.name)
    if not func_val and is_recursive(node):
        collect_symbolic_var_from_params(self, node)
        if node.returns is None:
            ret_sinfo = relax.TupleStructInfo([])
        else:
            ret_sinfo = eval_struct_info(self, node.returns, eval_str=True)
        params_sinfo = []
        for arg in node.args.args:
            if arg.annotation is None:
                self.report_error(arg, "Type annotation is required for function parameters.")
            param_sinfo = eval_struct_info(self, arg.annotation, eval_str=True)
            params_sinfo.append(param_sinfo)
        # created a var for the local function, the same var could be used for recursive call
        local_func_var = relax.Var(node.name, relax.FuncStructInfo(params_sinfo, ret_sinfo))
        self.var_table.add(node.name, local_func_var)

    purity = find_decorator_annotation(node, "pure")
    # treat the function as private if we are inside another function
    # or if it has a privacy annotation
    privacy = is_inner_function or find_decorator_annotation(node, "private", default=False)

    with self.var_table.with_frame():
        with self.with_dispatch_token("relax"):
            with R.function(is_pure=purity, is_private=privacy):
                R.func_name(node.name)
                collect_symbolic_var_from_params(self, node)

                if node.returns is not None:
                    ann_sinfo = eval_struct_info(self, node.returns, eval_str=True)
                    R.func_ret_struct_info(ann_sinfo)

                self.visit(node.args)

                for stmt in node.body:
                    if isinstance(stmt, doc.FunctionDef):
                        if not stmt.decorator_list:
                            self.report_error(stmt, "Function must be decorated")
                        dec = self.eval_expr(stmt.decorator_list[-1])
                        # inline prim_func was found
                        if dec.dispatch_token == "tirx":
                            self.report_error(stmt, "inline prim_func is disallowed in Relax IR")

                self.visit_body(node.body)
    self.inside_function = is_inner_function


def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool: