[Relax] Stabilize relax pass mutation order (#16883)

The current implementation of the relax pass is not stable, to be more
specific, the order of the mutation is not stable. This PR aims to
stabilize the mutation order of the relax pass, and further stabilize
the output of the relax pass.

Also fixes a minor doc typo in NN frontend

include/tvm/ir/module.h

@@ -249,7 +249,8 @@ class IRModuleNode : public Object {
   TVM_DLL GlobalVar GetGlobalVar(const String& str) const;
 
   /*!
-   * \brief Collect all global vars defined in this module.
+   * \brief Collect all global vars defined in this module, ordered by
+   *        the global variable name.
    * \returns An array of global vars
    */
   TVM_DLL Array<GlobalVar> GetGlobalVars() const;

python/tvm/relax/frontend/nn/core.py

@@ -475,10 +475,10 @@ def export_tvm(
         -------
         irmodule : tvm.ir.IRModule
             The converted tvm IR representation of the model.
-        params : Dict[str, tvm.nd.array]
-            A dictionary of parameters corresponding to the weights of
-            the model.
+        params : List[Tuple[str, Parameter]]
+            A list of Parameters corresponding to the weights of the model.
         ext_mods : List[nn.ExternModule]
+            A list of ExternModules that are used in the model.
         """
         # pylint: disable=import-outside-toplevel
         from . import spec as _spec

src/ir/module.cc

@@ -26,6 +26,7 @@
 #include <tvm/node/structural_equal.h>
 #include <tvm/runtime/registry.h>
 
+#include <algorithm>
 #include <fstream>
 #include <sstream>
 #include <unordered_set>
@@ -183,6 +184,9 @@ tvm::Array<GlobalVar> IRModuleNode::GetGlobalVars() const {
   for (const auto& pair : global_var_map_) {
     global_vars.push_back(pair.second);
   }
+  std::sort(global_vars.begin(), global_vars.end(), [](const GlobalVar& lhs, const GlobalVar& rhs) {
+    return lhs->name_hint < rhs->name_hint;
+  });
   return tvm::Array<GlobalVar>(global_vars);
 }
 

src/relax/transform/alter_op_impl.cc

@@ -89,7 +89,8 @@ class AlterOpImplMutator : public ExprMutator {
         op_buffer_axis_separators__(axis_separators_) {}
 
   IRModule Run() {
-    for (const auto& [gv, func] : mod_->functions) {
+    for (const auto& gv : mod_->GetGlobalVars()) {
+      const auto& func = mod_->Lookup(gv);
       if (func->IsInstance<relax::FunctionNode>()) {
         relax::Function update_func = Downcast<Function>(VisitExpr(func));
         builder_->UpdateFunction(gv, update_func);

src/relax/transform/dead_code_elimination.cc

@@ -148,7 +148,8 @@ IRModule DeadCodeElimination(const IRModule& arg_mod, Array<runtime::String> ent
   for (const auto& name : entry_function_names) {
     entry_functions.insert(mod->GetGlobalVar(name));
   }
-  for (const auto& [gv, func] : mod->functions) {
+  for (const auto& gv : mod->GetGlobalVars()) {
+    const auto& func = mod->Lookup(gv);
     if (func.as<ExternFuncNode>() || func->GetLinkageType() == LinkageType::kExternal) {
       entry_functions.insert(gv);
     }

src/relax/transform/fuse_ops.cc

@@ -691,7 +691,8 @@ class OperatorFusor : public ExprMutator {
    * \return The new IRModule after transformation
    */
   IRModule Transform() {
-    for (const auto& [gv, func] : mod_->functions) {
+    for (const auto& gv : mod_->GetGlobalVars()) {
+      const auto& func = mod_->Lookup(gv);
       // Only visit Relax function without attr kPrimitive.
       if (func->IsInstance<relax::FunctionNode>() && !func->HasNonzeroAttr(attr::kPrimitive)) {
         auto updated_func = Downcast<Function>(VisitExpr(func));
@@ -1196,9 +1197,9 @@ class CompositeFunctionAnnotator : public ExprMutator {
 
   IRModule Run() {
     auto mod = builder_->GetContextIRModule();
-    auto all_functions = mod->functions;
-    for (const auto& entry : all_functions) {
-      if (const auto* func = entry.second.as<FunctionNode>()) {
+    for (const auto& gv : mod->GetGlobalVars()) {
+      const auto& base_func = mod->Lookup(gv);
+      if (const auto* func = base_func.as<FunctionNode>()) {
         if (func->GetAttr<String>(attr::kComposite).defined() ||
             func->GetAttr<String>(attr::kCodegen).defined()) {
           continue;
@@ -1208,7 +1209,7 @@ class CompositeFunctionAnnotator : public ExprMutator {
         if (!new_body.same_as(func->body)) {
           auto new_func = Function(func->params, new_body, func->ret_struct_info, func->is_pure,
                                    func->attrs, func->span);
-          builder_->UpdateFunction(entry.first, new_func);
+          builder_->UpdateFunction(gv, new_func);
         }
       }
     }
@@ -1272,11 +1273,12 @@ IRModule FuseOpsByPattern(const tvm::Array<transform::FusionPattern>& patterns,
   support::Arena arena;
   for (const auto& pattern : patterns) {
     OperatorFusor::GroupMap group_map;
-    for (const auto& entry : mod->functions) {
-      if (entry.second->IsInstance<tir::PrimFuncNode>()) {
+    for (const auto& gv : mod->GetGlobalVars()) {
+      const auto& base_func = mod->Lookup(gv);
+      if (base_func->IsInstance<tir::PrimFuncNode>()) {
         continue;
       }
-      const FunctionNode* function = entry.second.as<FunctionNode>();
+      const FunctionNode* function = base_func.as<FunctionNode>();
       if (function->GetAttr<Integer>(attr::kPrimitive).defined() ||
           function->GetAttr<String>(attr::kComposite).defined() ||
           function->GetAttr<String>(attr::kCodegen).defined()) {
@@ -1285,8 +1287,8 @@ IRModule FuseOpsByPattern(const tvm::Array<transform::FusionPattern>& patterns,
 
       auto map = PatternBasedPartitioner::Run(pattern->name, pattern->pattern,
                                               pattern->annotation_patterns,
-                                              pattern->check.value_or(nullptr), entry.second,
-                                              &arena, pattern->attrs_getter.value_or(nullptr));
+                                              pattern->check.value_or(nullptr), base_func, &arena,
+                                              pattern->attrs_getter.value_or(nullptr));
       for (const auto& [key, value] : map) {
         CHECK(!group_map.count(key))
             << "ValueError: "

src/relax/transform/fuse_tir.cc

@@ -964,7 +964,8 @@ class TIRFuseMutator : public ExprMutator {
   static IRModule Transform(IRModule mod) {
     // Collect all primitive relax functions
     Map<GlobalVar, Function> primitive_relax;
-    for (const auto& [gvar, base_func] : mod->functions) {
+    for (const auto& gvar : mod->GetGlobalVars()) {
+      const auto& base_func = mod->Lookup(gvar);
       // Only fuse primitive relax functions
       if (base_func->HasNonzeroAttr(attr::kPrimitive)) {
         if (auto func = base_func.as<relax::Function>()) {

src/relax/transform/legalize_ops.cc

@@ -67,16 +67,18 @@ class LegalizeMutator : public ExprMutator {
   }
 
   IRModule Transform() {
-    for (const auto& [gv, func] : mod_->functions) {
+    for (const auto& gv : mod_->GetGlobalVars()) {
+      const auto& func = mod_->Lookup(gv);
       if (func->IsInstance<FunctionNode>()) {
         auto updated_func = Downcast<Function>(this->VisitExpr(func));
         builder_->UpdateFunction(gv, Downcast<BaseFunc>(updated_func));
       }
     }
     // Fill the "kTarget" attribute of PrimFunc
-    for (const auto& [gv, func] : builder_->GetContextIRModule()->functions) {
+    const auto& mod = builder_->GetContextIRModule();
+    for (const auto& gv : mod->GetGlobalVars()) {
       const tir::PrimFuncNode* prim_func;
-      if (tmap_.count(gv) && (prim_func = func.as<tir::PrimFuncNode>())) {
+      if (tmap_.count(gv) && (prim_func = mod->Lookup(gv).as<tir::PrimFuncNode>())) {
         auto f = WithAttr(GetRef<tir::PrimFunc>(prim_func), tvm::attr::kTarget, tmap_[gv]);
         builder_->UpdateFunction(gv, f);
       }