[Relax] Stabilize relax pass mutation order

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);
       }