Add priority to copies on backedges (#791)

* add a coalesce-locals test for costly backedges

* add script to strip local names out, for convenient diffing

* prioritize backedge copies

2 files changed, 33 insertions, 1 deletions

diff --git a/src/cfg/cfg-traversal.h b/src/cfg/cfg-traversal.h
index 99bc25ea8..13bb0c744 100644
--- a/src/cfg/cfg-traversal.h
+++ b/src/cfg/cfg-traversal.h
@@ -54,6 +54,7 @@ struct CFGWalker : public ControlFlowWalker<SubType, VisitorType> {
   // internal details
 
   std::vector<std::unique_ptr<BasicBlock>> basicBlocks; // all the blocks
+  std::vector<BasicBlock*> loopTops; // blocks that are the tops of loops, i.e., have backedges to them
 
   // traversal state
   BasicBlock* currBasicBlock; // the current block in play during traversal. can be nullptr if unreachable,
@@ -132,6 +133,7 @@ struct CFGWalker : public ControlFlowWalker<SubType, VisitorType> {
   static void doStartLoop(SubType* self, Expression** currp) {
     auto* last = self->currBasicBlock;
     self->startBasicBlock();
+    self->loopTops.push_back(self->currBasicBlock); // a loop with no backedges would still be counted here, but oh well
     self->link(last, self->currBasicBlock);
     self->loopStack.push_back(self->currBasicBlock);
   }
diff --git a/src/passes/CoalesceLocals.cpp b/src/passes/CoalesceLocals.cpp
index 3127b429c..9a2f8d220 100644
--- a/src/passes/CoalesceLocals.cpp
+++ b/src/passes/CoalesceLocals.cpp
@@ -178,13 +178,19 @@ struct CoalesceLocals : public WalkerPass<CFGWalker<CoalesceLocals, Visitor<Coal
     self->currBasicBlock->contents.actions.emplace_back(Action::Set, curr->index, currp);
     // if this is a copy, note it
     auto* get = curr->value->dynCast<GetLocal>();
-    if (get) self->addCopy(curr->index, get->index);
+    if (get) {
+      // add 2 units, so that backedge prioritization can decide ties, but not much more
+      self->addCopy(curr->index, get->index);
+      self->addCopy(curr->index, get->index);
+    }
   }
 
   // main entry point
 
   void doWalkFunction(Function* func);
 
+  void increaseBackEdgePriorities();
+
   void flowLiveness();
 
   void calculateInterferences();
@@ -251,6 +257,8 @@ void CoalesceLocals::doWalkFunction(Function* func) {
   // ignore links to dead blocks, so they don't confuse us and we can see their stores are all ineffective
   liveBlocks = findLiveBlocks();
   unlinkDeadBlocks(liveBlocks);
+  // increase the cost of costly backedges
+  increaseBackEdgePriorities();
 #ifdef CFG_DEBUG
   dumpCFG("the cfg");
 #endif
@@ -273,6 +281,28 @@ void CoalesceLocals::doWalkFunction(Function* func) {
   applyIndices(indices, func->body);
 }
 
+// A copy on a backedge can be especially costly, forcing us to branch just to do that copy.
+// Add weight to such copies, so we prioritize getting rid of them.
+void CoalesceLocals::increaseBackEdgePriorities() {
+  for (auto* loopTop : loopTops) {
+    // ignore the first edge, it is the initial entry, we just want backedges
+    auto& in = loopTop->in;
+    for (Index i = 1; i < in.size(); i++) {
+      auto* arrivingBlock = in[i];
+      if (arrivingBlock->out.size() > 1) continue; // we just want unconditional branches to the loop top, true phi fragments
+      for (auto& action : arrivingBlock->contents.actions) {
+        if (action.what == Action::Set) {
+          auto* set = (*action.origin)->cast<SetLocal>();
+          if (auto* get = set->value->dynCast<GetLocal>()) {
+            // this is indeed a copy, add to the cost (default cost is 2, so this adds 50%, and can mostly break ties)
+            addCopy(set->index, get->index);
+          }
+        }
+      }
+    }
+  }
+}
+
 void CoalesceLocals::flowLiveness() {
   interferences.resize(numLocals * numLocals);
   std::fill(interferences.begin(), interferences.end(), 0);