[Unity][IR][UX] Privacy annotation in Relax

include/tvm/script/ir_builder/relax/frame.h

@@ -99,6 +99,8 @@ class FunctionFrameNode : public SeqExprFrameNode {
   Optional<tvm::relax::StructInfo> ret_struct_info;
   /*! \brief Whether the function is annotated as pure */
   Optional<Bool> is_pure;
+  /*! \brief Whether the function is annotated as private */
+  Optional<Bool> is_private;
   /*! \brief The function attributes. */
   Map<String, ObjectRef> attrs;
   /*! \brief The block builder to create Relax function. */

include/tvm/script/ir_builder/relax/ir.h

@@ -34,9 +34,10 @@ namespace relax {
 /*!
  * \brief Start a function frame.
  * \param is_pure Whether the function is annotated as pure.
+ * \param is_private Whether the function is annotated as private.
  * \return The created ir_builder Function frame.
  */
-TVM_DLL FunctionFrame Function(const Bool& is_pure);
+TVM_DLL FunctionFrame Function(const Bool& is_pure, const Bool& is_private);
 
 /*!
  * \brief Add a parameter to the last function frame.

python/tvm/relax/block_builder.py

@@ -199,6 +199,7 @@ def function(
         name: str,
         params: Optional[Union[Var, Tuple, List[Var]]] = None,
         attrs: Optional[Dict[str, Object]] = None,
+        private: bool = False,
     ) -> FunctionScope:
         """Annotate a Relax function.
 
@@ -215,6 +216,12 @@ def function(
         attrs : Dict[str, Object], optional
             The function attrs
 
+        private : bool, optional
+            Whether the function is annotated as private.
+            If the function is private, it will not have a global symbol attribute.
+            If it is not private and not an inner function, then it will have
+            a global symbol attribute (mapped to the function's name)
+
         Returns
         -------
         ret: FunctionScope
@@ -233,6 +240,11 @@ def function(
                     )
         if attrs is None:
             attrs = {}
+        # The block builder does not permit nesting functions, per above comment,
+        # so no further check should be needed
+        if not private:
+            attrs["global_symbol"] = name
+
         return FunctionScope(self, name, params, attrs)
 
     def testing_scope(self, def_vars: List[tir.Var]) -> TestingScope:

python/tvm/relax/frontend/torch/dynamo.py

@@ -154,7 +154,8 @@ def _capture(graph_module: fx.GraphModule, example_inputs):
             keep_params_as_input=keep_params_as_input,
             unwrap_unit_return_tuple=True,
         )
-        mod[f"subgraph_{len(mod.get_global_vars())}"] = mod_["main"]
+        new_name = f"subgraph_{len(mod.get_global_vars())}"
+        mod[new_name] = mod_["main"].with_attr("global_symbol", new_name)
         return graph_module.forward
 
     dynamo.reset()

python/tvm/relax/training/setup_trainer.py

@@ -198,7 +198,10 @@ def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRM
 
         # Add optimizer function.
         self._optimizer.init(params)
-        mod[self.OPTIMIZER_FUNC] = self._optimizer.get_function()
+        # Need the global symbol to match the function's name
+        mod[self.OPTIMIZER_FUNC] = self._optimizer.get_function().with_attr(
+            "global_symbol", self.OPTIMIZER_FUNC
+        )
 
         # Module attrs
         mod = mod.with_attrs(

python/tvm/script/ir_builder/relax/ir.py

@@ -165,19 +165,24 @@
 ############################### Function ################################
 
 
-def function(is_pure: bool = True) -> frame.FunctionFrame:
+def function(is_pure: bool = True, is_private: bool = False) -> frame.FunctionFrame:
     """Start a function frame.
     Parameters
     ----------
     is_pure: bool
         Whether the function is annotated as pure.
 
+    is_private : bool
+        Whether the function is annotated as private.
+
     Returns
     -------
     frame: FunctionFrame
         The constructed function frame.
     """
-    return _ffi_api.Function(is_pure)  # type: ignore[attr-defined] # pylint: disable=no-member
+    return _ffi_api.Function(  # type: ignore[attr-defined]  # pylint: disable=no-member
+        is_pure, is_private
+    )
 
 
 def arg(name: py_str, struct_info: StructInfo) -> Var:

python/tvm/script/parser/core/parser.py

@@ -240,6 +240,7 @@ class Parser(doc.NodeVisitor):
     dispatch_tokens: List[str]
     function_annotations: Optional[Dict[str, Dict[str, Any]]]
     var_table: VarTable
+    inside_function: bool  # whether we are within a function
 
     def __init__(
         self,
@@ -250,6 +251,7 @@ def __init__(
         self.dispatch_tokens = ["default"]
         self.function_annotations = function_annotations
         self.var_table = VarTable()
+        self.inside_function = False
 
     def parse(self, extra_vars: Optional[Dict[str, Any]] = None) -> Any:
         """The main parse method for parser.

python/tvm/script/parser/relax/entry.py

@@ -45,9 +45,11 @@
 # this formulation allows us to support having @R.function
 # appear as a decorator by itself or to have optional arguments
 # like @R.function(pure=False)
-def function(f: Optional[FType] = None, pure: bool = True) -> Union[Function, FType]:
+def function(
+    f: Optional[FType] = None, pure: bool = True, private: bool = False
+) -> Union[Function, FType]:
     # pylint: disable=unused-argument
-    # (pure isn't used here, but is used later in parsing)
+    # (pure and private aren't used here, but are used later in parsing)
 
     # need to inspect the stack first because is_defined_in_class expects the outer class
     # to be in a particular position in the stack

python/tvm/script/parser/relax/parser.py

@@ -160,6 +160,9 @@ def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) -> Non
 
 @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):
@@ -178,11 +181,14 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
         local_func_var = relax.Var(node.name, relax.FuncStructInfo(params_sinfo, ret_sinfo))
         self.var_table.add(node.name, local_func_var)
 
-    purity = find_purity_annotation(node)
+    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):
+            with R.function(is_pure=purity, is_private=privacy):
                 R.func_name(node.name)
                 collect_symbolic_var_from_params(self, node)
 
@@ -202,22 +208,22 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
                             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_purity_annotation(node: doc.FunctionDef) -> bool:
+def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
     """
-    Check the value of `pure` in the function decorator.
-    Returns the annotated purity if present, otherwise defaulting to True.
-    This allows for specifying the purity in the function signature.
+    Check the value of given annotation (argument name) in the function decorator.
+    Returns the value of the annotation if present, otherwise giving the default value.
     """
-    # look for the pure argument in the function decorator
+    # look for the named argument in the function decorator
     for dec in node.decorator_list:
         if not isinstance(dec, doc.Call) or dec.func.attr != "function":
             continue
         for keyword in dec.keywords:
-            if keyword.arg == "pure":
+            if keyword.arg == annotation:
                 return keyword.value.value
-    return True
+    return default
 
 
 @dispatch.register(token="relax", type_name="tvm_declare_function")
@@ -238,7 +244,8 @@ def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar
             param_sinfo = eval_struct_info(self, arg.annotation, eval_str=True)
             params.append(relax.Var(arg.arg, param_sinfo))
 
-    is_pure = find_purity_annotation(node)
+    is_pure = find_decorator_annotation(node, "pure")
+
     func_signature = relax.Function.create_empty(params, ret_sinfo, is_pure=is_pure)
     return I.decl_function(node.name, func_signature)
 

src/relax/training/utils.cc

@@ -48,7 +48,9 @@ class AppendLossMutator : private ExprMutator {
     Function new_loss_func = CopyWithNewVars(loss_function);
 
     AppendLossMutator mutator(mod, new_loss_func, num_backbone_outputs);
-    auto new_func_transformed = Downcast<Function>(mutator.VisitExpr(new_func));
+    auto new_func_transformed =
+        WithAttr(Downcast<Function>(mutator.VisitExpr(new_func)), tvm::attr::kGlobalSymbol,
+                 new_func_name.value_or(func_name + "_loss"));
 
     auto new_module = GetRef<IRModule>(mod.CopyOnWrite());
     auto new_var = GlobalVar(new_func_name.value_or(func_name + "_loss"));

src/relax/transform/gradient.cc

@@ -333,10 +333,14 @@ class GradientMutator : private ExprMutator {
     }
 
     GradientMutator mutator(mod, require_grads_value, target_index);
-    Function new_func_transformed = Downcast<Function>(mutator.VisitExpr(new_func));
+
+    // make the adjoint public
+    auto new_name = func_name + "_adjoint";
+    Function new_func_transformed = WithAttr(Downcast<Function>(mutator.VisitExpr(new_func)),
+                                             tvm::attr::kGlobalSymbol, new_name);
 
     IRModule new_module = GetRef<IRModule>(mod.CopyOnWrite());
-    new_module->Add(GlobalVar(func_name + "_adjoint"), new_func_transformed);
+    new_module->Add(GlobalVar(new_name), new_func_transformed);
     return new_module;
   }
 

src/relax/transform/lift_transform_params.cc

@@ -251,7 +251,10 @@ class TransformParamsLifter : public ExprMutator {
     lift_plan_ = planner.Plan(func, num_input);
 
     // Step 2: Add the lifted function to the module
-    builder_->AddFunction(lift_plan_.f_transform_params, new_func_name);
+    // (The lifted function should be public so we add a global symbol to it)
+    auto lift_func =
+        WithAttr(lift_plan_.f_transform_params, tvm::attr::kGlobalSymbol, new_func_name);
+    builder_->AddFunction(lift_func, new_func_name);
 
     // Step 3: Update the current function.
 

src/script/ir_builder/relax/frame.cc

@@ -56,6 +56,11 @@ void FunctionFrameNode::ExitWithScope() {
                              "`return` to return an Expr";
   this->block_builder->BeginScope(params);
   Expr body = this->block_builder->Normalize(tvm::relax::SeqExpr(binding_blocks, output.value()));
+  // if the function is not private, add a global symbol to its attributes
+  if (!is_private.value_or(Bool(false))->value && name.defined() &&
+      !attrs.count(tvm::attr::kGlobalSymbol)) {
+    attrs.Set(tvm::attr::kGlobalSymbol, name.value());
+  }
   auto dict_attrs = attrs.empty() ? NullValue<DictAttrs>() : DictAttrs(attrs);
   this->block_builder->EndScope();
   tvm::relax::Function func(/*params=*/params,

src/script/ir_builder/relax/ir.cc

@@ -52,7 +52,7 @@ TVM_STATIC_IR_FUNCTOR(Namer, vtable)
 
 /////////////////////////////// Function ////////////////////////////////
 
-FunctionFrame Function(const Bool& is_pure) {
+FunctionFrame Function(const Bool& is_pure, const Bool& is_private) {
   ObjectPtr<FunctionFrameNode> n = make_object<FunctionFrameNode>();
   const IRBuilder& ir_builder = IRBuilder::Current();
   Optional<tvm::IRModule> mod = NullOpt;
@@ -61,6 +61,7 @@ FunctionFrame Function(const Bool& is_pure) {
   }
   n->block_builder = tvm::relax::BlockBuilder::Create(/*mod=*/mod);
   n->is_pure = is_pure;
+  n->is_private = is_private;
   return FunctionFrame(n);
 }
 

src/script/printer/relax/function.cc

@@ -22,13 +22,36 @@ namespace tvm {
 namespace script {
 namespace printer {
 
+bool AtTopLevelFunction(const IRDocsifier& d) {
+  // fewer than 2 frames: not in a function at all
+  if (d->frames.size() < 2) {
+    return false;
+  }
+  // if the first frame is a RelaxFrame, then this is not inside a module.
+  // 2 frames => we are at a function (more than 2 => nested function)
+  if (d->frames[0]->IsInstance<RelaxFrameNode>()) {
+    return d->frames.size() == 2;
+  }
+  // otherwise the first two frames pertain to an IR module,
+  // so 3 frames => we are at a top-level function (more than 3 => nested function)
+  return d->frames.size() == 3;
+}
+
 TVM_REGISTER_NODE_TYPE(RelaxFrameNode);
 
 TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
     .set_dispatch<relax::Function>("", [](relax::Function n, ObjectPath n_p, IRDocsifier d) -> Doc {
       std::unordered_set<const tir::VarNode*> func_vars;
       With<RelaxFrame> f(d);
-      IdDoc func_name = d->Define(n, f(), FindFunctionName(d, n).value_or("main"));
+
+      IdDoc func_name("");
+      // if we are binding a local definition, then calling d->Define
+      // will result in a repeated definition and an incorrect displayed name
+      if (Optional<String> name = GetBindingName(d)) {
+        func_name = std::move(IdDoc(name.value()));
+      } else {
+        func_name = std::move(d->Define(n, f(), FindFunctionName(d, n).value_or("main")));
+      }
       (*f)->AddDispatchToken(d, "relax");
       (*f)->is_func = true;
       (*f)->func_vars = &func_vars;
@@ -52,17 +75,49 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
       (*f)->func_vars = nullptr;
       // Step 4. Print attributes
       if (n->attrs.defined() && !n->attrs->dict.empty()) {
-        (*f)->stmts.push_back(
-            ExprStmtDoc(Relax(d, "func_attr")  //
-                            ->Call({d->AsDoc<ExprDoc>(n->attrs, n_p->Attr("attrs"))})));
+        // If the function is a global function and has a global symbol,
+        // then don't print the global symbol (it will be implicit from not being private).
+        // For a function without an IR module whose global symbol
+        // doesn't match the function name, we should still print the global symbol attribute.
+        if (AtTopLevelFunction(d) && n->attrs->dict.count(tvm::attr::kGlobalSymbol) &&
+            Downcast<String>(n->attrs->dict.at(tvm::attr::kGlobalSymbol)) == func_name->name) {
+          Map<String, ObjectRef> new_attrs;
+          for (auto kv : n->attrs->dict) {
+            if (kv.first != tvm::attr::kGlobalSymbol) {
+              new_attrs.Set(kv.first, kv.second);
+            }
+          }
+          if (!new_attrs.empty()) {
+            (*f)->stmts.push_back(ExprStmtDoc(
+                Relax(d, "func_attr")  //
+                    ->Call({d->AsDoc<ExprDoc>(DictAttrs(new_attrs), n_p->Attr("attrs"))})));
+          }
+        } else {
+          (*f)->stmts.push_back(
+              ExprStmtDoc(Relax(d, "func_attr")  //
+                              ->Call({d->AsDoc<ExprDoc>(n->attrs, n_p->Attr("attrs"))})));
+        }
       }
       // Step 5. Prepare the decorator (include purity if it's impure)
       ExprDoc decorator = Relax(d, "function");
+      Array<ExprDoc, void> pos_args = {};
+      Array<String, void> dec_keys;
+      Array<ExprDoc, void> dec_values;
       if (!n->is_pure) {
-        Array<ExprDoc> pos_args = {};
-        decorator = std::move(decorator->Call(
-            pos_args, {"pure"}, {LiteralDoc::Boolean(false, Optional<ObjectPath>())}));
+        dec_keys.push_back("pure");
+        dec_values.push_back(LiteralDoc::Boolean(false, Optional<ObjectPath>()));
       }
+      // if the function is global or is not in a module and does not have a global symbol,
+      // indicate that it's private
+      if (AtTopLevelFunction(d) &&
+          (!n->attrs.defined() || !n->attrs->dict.count(tvm::attr::kGlobalSymbol))) {
+        dec_keys.push_back("private");
+        dec_values.push_back(LiteralDoc::Boolean(true, Optional<ObjectPath>()));
+      }
+      if (dec_keys.size()) {
+        decorator = std::move(decorator->Call(pos_args, dec_keys, dec_values));
+      }
+
       // Step 6. Print body
       Array<StmtDoc> body =
           PrintSeqExpr(Downcast<relax::SeqExpr>(n->body), n_p->Attr("body"), d, /*use_ret=*/true);