Optimize ReorderGlobals ordering

src/passes/ReorderGlobals.cpp

@@ -35,6 +35,7 @@
 
 #include "ir/find_all.h"
 #include "pass.h"
+#include "support/topological_sort.h"
 #include "wasm.h"
 
 namespace wasm {
@@ -78,7 +79,9 @@ struct ReorderGlobals : public Pass {
   // each one moves, which is logically a mapping between indices.
   using IndexIndexMap = std::vector<Index>;
 
-  // We will also track counts of uses for each global.
+  // We will also track counts of uses for each global. We use floating-point
+  // values here since while the initial counts are integers, we will be
+  // considering fractional sums of them later.
   using IndexCountMap = std::vector<double>;
 
   // We must take into account dependencies, so that globals appear before
@@ -148,7 +151,8 @@ struct ReorderGlobals : public Pass {
     struct SortAndSize {
       IndexIndexMap sort;
       double size;
-      SortAndSize(IndexIndexMap&& sort, double size) : sort(sort), size(size) {}
+      SortAndSize(IndexIndexMap&& sort, double size)
+        : sort(std::move(sort)), size(size) {}
     };
     std::vector<SortAndSize> options;
     auto addOption = [&](const IndexCountMap& customCounts) {
@@ -185,44 +189,39 @@ struct ReorderGlobals : public Pass {
     // simple sum may be too naive).
     double const EXPONENTIAL_FACTOR = 0.095;
     IndexCountMap sumCounts(globals.size()), exponentialCounts(globals.size());
-    // We add items to |sumCounts, exponentialCounts| after they are computed,
-    // so anything not there must be pushed to a stack before we can handle it.
-    std::vector<Index> stack;
-    // Do a simple recursion to compute children before parents, while tracking
-    // which have been computed.
-    std::vector<bool> computed(globals.size());
-    for (Index i = 0; i < globals.size(); i++) {
-      stack.push_back(i);
-    }
-    while (!stack.empty()) {
-      auto global = stack.back();
-      if (computed[global]) {
-        // We've already computed this.
-        stack.pop_back();
-        continue;
-      }
-      // Leave ourselves on the stack and push any unresolved deps.
-      auto pushed = false;
-      for (auto dep : deps.dependedUpon[global]) {
-        if (!computed[dep]) {
-          stack.push_back(dep);
-          pushed = true;
+
+    struct Sort : public TopologicalSort<Index, Sort> {
+      const Dependencies& deps;
+
+      Sort(Index numGlobals, const Dependencies& deps) : deps(deps) {
+        for (Index i = 0; i < numGlobals; i++) {
+          push(i);
         }
       }
-      if (pushed) {
-        continue;
+
+      void pushPredecessors(Index global) {
+        auto iter = deps.dependedUpon.find(global);
+        if (iter == deps.dependedUpon.end()) {
+          return;
+        }
+        for (auto dep : iter->second) {
+          push(dep);
+        }
       }
-      // We can handle this now, as all deps are resolved.
-      stack.pop_back();
-      // Start with the self-count, then add the deps.
+    } sort(globals.size(), deps);
+
+    for (auto global : sort) {
+      // We can compute this global's count as in the sorted order all the
+      // values it cares about are resolved. Start with the self-count, then
+      // add the deps.
       sumCounts[global] = exponentialCounts[global] = counts[global];
       for (auto dep : deps.dependedUpon[global]) {
         sumCounts[global] += sumCounts[dep];
         exponentialCounts[global] +=
           EXPONENTIAL_FACTOR * exponentialCounts[dep];
       }
-      computed[global] = true;
     }
+
     addOption(sumCounts);
     addOption(exponentialCounts);
 
@@ -237,14 +236,12 @@ struct ReorderGlobals : public Pass {
     }
 
     // Apply the indices we computed.
-    std::sort(
-      globals.begin(),
-      globals.end(),
-      [&](const std::unique_ptr<Global>& a, const std::unique_ptr<Global>& b) {
-        return (*best)[originalIndices[a->name]] <
-               (*best)[originalIndices[b->name]];
-      });
-
+    std::vector<std::unique_ptr<Global>> old;
+    old.swap(globals);
+    globals.resize(old.size());
+    for (Index i = 0; i < old.size(); i++) {
+      globals[(*best)[i]] = std::move(old[i]);
+    }
     module->updateMaps();
   }
 
@@ -308,11 +305,8 @@ struct ReorderGlobals : public Pass {
       // Sort by the counts.
       auto aCount = counts[a];
       auto bCount = counts[b];
-      if (aCount < bCount) {
-        return true;
-      }
-      if (aCount > bCount) {
-        return false;
+      if (aCount != bCount) {
+        return aCount < bCount;
       }
 
       // Break ties using the original order, which means just using the
@@ -333,7 +327,7 @@ struct ReorderGlobals : public Pass {
       }
     }
 
-    // Pop off the heap: Emit the global and it its final, sorted index. Keep
+    // Pop off the heap: Emit the global and its final, sorted index. Keep
     // doing that until we finish processing all the globals.
     IndexIndexMap sortedindices(globals.size());
     Index numSortedindices = 0;
@@ -380,7 +374,7 @@ struct ReorderGlobals : public Pass {
 
     if (always) {
       // In this mode we gradually increase the cost of later globals, in an
-      // unrealistic but manner.
+      // unrealistic but smooth manner.
       double total = 0;
       for (Index i = 0; i < actualOrder.size(); i++) {
         // Multiply the count for this global by a smoothed LEB factor, which

test/lit/passes/reorder-globals.wast

@@ -742,7 +742,7 @@
 ;; was moved up to between $b and $c. Sum sort works well here because the
 ;; first comparison is $a and $other, and sum takes into account $b and $c in
 ;; $a's favor, so it wins. Likewise $b and $c win against $other as well, so
-;; the other is $a, $b, $c, $other which is optimal here.
+;; the order is $a, $b, $c, $other which is optimal here.
 (module
   ;; CHECK:      (global $a i32 (i32.const 0))
   (global $a i32 (i32.const 0))