[WebAssembly] Optimize the number of routing blocks in FixIrreducibleCFG

[aheejin]

Summary

Currently we create a routing block to the dispatch block for every
predecessor of every entry. So the total number of routing blocks
created will be (# of preds) * (# of entries). But we don't need to do
this: we need at most 2 routing blocks per loop entry, one for when the
predecessor is inside the loop and one for it is outside the loop. (We
can't merge these into one because this will creates another loop cycle
between blocks inside and blocks outside) This patch fixes this and
creates at most 2 routing blocks per entry.

This also renames variable Split to Routing, which I think is a bit
clearer.

[aheejin]

@kripken I wonder how much this will reduce in code size... How can I test with malloc or fannkuch (or other benchmarks)? (Where is the source / and what are the command line options you use?)

[aheejin]

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

406 ↗

(On Diff #190978)

This was previously DenseMap, but I changed this to std::map because to use std::pair as a key to DenseMap I need to write DenseMapInfo routines and I don't think it's really worth it because this map will be quite small anyway.

[aheejin]

Newline in test

[riccibruno]

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

406 ↗

(On Diff #190978)

Note that DenseMapInfo is already implemented for std::pair

[aheejin]

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

406 ↗

(On Diff #190978)

Thank you. But I think I was not correct by saying std::pair is not supported; the real problem looks like we don't have DenseMapInfo for bool. For example, this does not compile (without extra implementation of DenseMapInfo for bool)

DenseMap<bool, int> Map;

[aheejin]

Comment grammar fix

[aheejin]

Fix fallthrough condition check

[kripken]

Nice idea, and overall looks good to me. I didn't quite understand the TII part though.

Btw, are there no backend passes that would do this optimization anyhow (merge identical branches)?

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

390 ↗

(On Diff #191011)

What is the advantage of a map to bool over a set?

[aheejin]

• Use a set instead for InLoop

[aheejin]

In D59462#1433718, @kripken wrote:

Nice idea, and overall looks good to me. I didn't quite understand the TII part though.

TargetInstrInfo::analyzeBranch function's output is a bit complicated; we have to check the return value and three parameters that are passed (TBB, FBB, and Cond). The detailed description on the meaning of its output is here. The default return value true means (not intuitively) the branch is not analyzable, and each target is supposed to override this function. Our overridden version is here.

Before, we added a routing block for each predecessor / entry combination, so if in the original code a predecessor fell through to an entry, we placed a routing block right there where the entry block was, maintaining the fall-through relationship, so that in the new code the predecessor fell through to the new routing block. But after this patch we don't create a routing block per pred/entry conbination; there can be only one routing block for multiple predecessors, and one of those predecessor might have fell through to the entry in the original code, i.e., it was the layout predecessor of the entry. But we create a routing block when we first encounter any predecessor, so it's possible the new routing block is not created right after the layout predecessor anymore. For example, when the original code looks like this:

pred0:
  ...
  br label %entry
pred1:
  ...
  (falls through to %entry)
entry:
  ...

Before, we created two routing blocks, each for pred0 and pred1, so the result will be like this:

pred0:
  ...
  br label %routing0
pred1:
  ...
  (falls through to %routing1)
routing1:   // This `routing1` is created right after `pred1`, because `pred1` was the layout predecessor of `entry` 
  i32.const 0
  br %dispatch
entry:
  ...
routing0:
  i32.const 0
  br %dispatch
dispatch:
 br_table ...

But now we create only one routing block in this case, so we end up with

pred0:
  ...
  br label %routing
pred1:
  ...
  br label %routing            <-- We need this new branch
entry:
  ...
routing:
  i32.const 0
  br %dispatch
dispatch:
 br_table ...

We need the new branch, because we already created routing when we first encountered pred0, and for pred1 we need a branch. Ideally we can do some more optimization here by first checking if there is any layout predecessor among all predecessors and create the routing block right after that, but all of these can change anyway in CFGSort, so I'm not sure how much benefit it will give.

Anyway, so for the code:

// In case the predecessor fell through to the successor in the original
// CFG and and the new routing block is not the predecessor's layout
// successor, we need to add a branch to the new routing block.
MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
SmallVector<MachineOperand, 4> Cond;
bool Analyzable = !TII.analyzeBranch(*Pred, TBB, FBB, Cond);
if (Analyzable &&
    ((Cond.empty() && !TBB && !FBB) ||
     (!Cond.empty() && !FBB && Pred->isLayoutSuccessor(Succ))))
  BuildMI(Pred, DebugLoc(), TII.get(WebAssembly::BR)).addMBB(Routing);

• Cond.empty() && !TBB && !FBB is for when we don't have any branches and just fall through to the next block.
• !Cond.empty() && !FBB && Pred->isLayoutSuccessor(Succ) is for when we only have a conditional branch in the block and if that condition is not taken we fall through to the next block.

Btw, are there no backend passes that would do this optimization anyhow (merge identical branches)?

There are, but all of them run before this pass, because this pass is the last pass in the wasm compilation pipeline that modifies CFG. (which is not strictly true, because we have to modify CFG in CFGStackify to resolve things for exception handling, but I mean in general.)

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

390 ↗

(On Diff #191011)

None. It should've been a set. I think I was changing the implementation this way and that and somehow ended up with that weird map. :( Thanks for the catch.

[aheejin]

• Rebase

[kripken]

Thanks for the explanation! I think I understand.

So this patch makes us only create a branch there if we don't happen to fall through to the right place. Does that mean that if we ran some CFG-related pass after this, it would not see a valid CFG? (I guess we don't run such a pass now, but I wonder if we may reorder passes in the future at some point.)

[aheejin]

In D59462#1435602, @kripken wrote:

Thanks for the explanation! I think I understand.

So this patch makes us only create a branch there if we don't happen to fall through to the right place. Does that mean that if we ran some CFG-related pass after this, it would not see a valid CFG? (I guess we don't run such a pass now, but I wonder if we may reorder passes in the future at some point.)

I don't think so.

1. This pass should be the last pass that modifies CFG anyway, considering what this pass does. (There's an exception for exception handling in D48345 though)
2. Even if for some reason we have to change CFG after this pass, that pass is responsible for making everything valid, as in other CFG-changing passes. That does not change whether the pass is before or after this pass.

But to think about it, because CFGSort tried to preserve the BB order if possible, actually it may be worthwhile to figure out if there is a predecessor that is the layout predecessor of the entry and try to place the routing block after that, which can save a branch. Let me try that.

[aheejin]

Layout predecessor optimization. Now if there's a layout predecessor to the entry, we generate the routing block right after the layout predecessor. So now there's no case we need to analyze branch and generate extra brs.

[aheejin]

• Comment grammar fix

[kripken]

Nice!

llvm/trunk/lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

class WebAssemblyFixIrreducibleControlFlow final : public MachineFunctionPass {
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   bool processRegion(MachineBasicBlock *Entry, BlockSet &Blocks,
                      MachineFunction &MF);
 
   void makeSingleEntryLoop(BlockSet &Entries, BlockSet &Blocks,
-                           MachineFunction &MF);
+                           MachineFunction &MF, const ReachabilityGraph &Graph);
 
 public:
   static char ID; // Pass identification, replacement for typeid
   WebAssemblyFixIrreducibleControlFlow() : MachineFunctionPass(ID) {}
 };
 
 bool WebAssemblyFixIrreducibleControlFlow::processRegion(
     MachineBasicBlock *Entry, BlockSet &Blocks, MachineFunction &MF) {
for (auto *LoopEntry : Graph.getLoopEntries()) {
         if (OtherLoopEntry != LoopEntry &&
             Graph.canReach(LoopEntry, OtherLoopEntry) &&
             Graph.canReach(OtherLoopEntry, LoopEntry)) {
           MutualLoopEntries.insert(OtherLoopEntry);
         }
       }
 
       if (MutualLoopEntries.size() > 1) {
-        makeSingleEntryLoop(MutualLoopEntries, Blocks, MF);
+        makeSingleEntryLoop(MutualLoopEntries, Blocks, MF, Graph);
         FoundIrreducibility = true;
         Changed = true;
         break;
       }
     }
     // Only go on to actually process the inner loops when we are done
     // removing irreducible control flow and changing the graph. Modifying
     // the graph as we go is possible, and that might let us avoid looking at
while (true) {
 
     return Changed;
   }
 }
 
 // Given a set of entries to a single loop, create a single entry for that
 // loop by creating a dispatch block for them, routing control flow using
 // a helper variable. Also updates Blocks with any new blocks created, so
-// that we properly track all the blocks in the region.
+// that we properly track all the blocks in the region. But this does not update
+// ReachabilityGraph; this will be updated in the caller of this function as
+// needed.
 void WebAssemblyFixIrreducibleControlFlow::makeSingleEntryLoop(
-    BlockSet &Entries, BlockSet &Blocks, MachineFunction &MF) {
+    BlockSet &Entries, BlockSet &Blocks, MachineFunction &MF,
+    const ReachabilityGraph &Graph) {
   assert(Entries.size() >= 2);
 
   // Sort the entries to ensure a deterministic build.
   BlockVector SortedEntries(Entries.begin(), Entries.end());
   llvm::sort(SortedEntries,
              [&](const MachineBasicBlock *A, const MachineBasicBlock *B) {
                auto ANum = A->getNumber();
                auto BNum = B->getNumber();
#endif
   for (auto *Entry : SortedEntries) {
     for (auto *Pred : Entry->predecessors()) {
       if (Pred != Dispatch) {
         AllPreds.push_back(Pred);
       }
     }
   }
 
-  for (MachineBasicBlock *Pred : AllPreds) {
-    DenseMap<MachineBasicBlock *, MachineBasicBlock *> Map;
+  // This set stores predecessors within this loop.
+  DenseSet<MachineBasicBlock *> InLoop;
+  for (auto *Pred : AllPreds) {
     for (auto *Entry : Pred->successors()) {
-      if (!Entries.count(Entry)) {
+      if (!Entries.count(Entry))
         continue;
-      }
+      if (Graph.canReach(Entry, Pred)) {
+        InLoop.insert(Pred);
+        break;
+      }
+    }
 
+  // Record if each entry has a layout predecessor. This map stores
+  // <<Predecessor is within the loop?, loop entry>, layout predecessor>
+  std::map<std::pair<bool, MachineBasicBlock *>, MachineBasicBlock *>
+      EntryToLayoutPred;
+  for (auto *Pred : AllPreds)
+    for (auto *Entry : Pred->successors())
+      if (Entries.count(Entry) && Pred->isLayoutSuccessor(Entry))
+        EntryToLayoutPred[std::make_pair(InLoop.count(Pred), Entry)] = Pred;
+
+  // We need to create at most two routing blocks per entry: one for
+  // predecessors outside the loop and one for predecessors inside the loop.
+  // This map stores
+  // <<Predecessor is within the loop?, loop entry>, routing block>
+  std::map<std::pair<bool, MachineBasicBlock *>, MachineBasicBlock *> Map;
+  for (auto *Pred : AllPreds) {
+    bool PredInLoop = InLoop.count(Pred);
+    for (auto *Entry : Pred->successors()) {
+      if (!Entries.count(Entry) ||
+          Map.count(std::make_pair(InLoop.count(Pred), Entry)))
+        continue;
+      // If there exists a layout predecessor of this entry and this predecessor
+      // is not that, we rather create a routing block after that layout
+      // predecessor to save a branch.
+      if (EntryToLayoutPred.count(std::make_pair(PredInLoop, Entry)) &&
+          EntryToLayoutPred[std::make_pair(PredInLoop, Entry)] != Pred)
+        continue;
+
       // This is a successor we need to rewrite.
-      MachineBasicBlock *Split = MF.CreateMachineBasicBlock();
+      MachineBasicBlock *Routing = MF.CreateMachineBasicBlock();
       MF.insert(Pred->isLayoutSuccessor(Entry)
                     ? MachineFunction::iterator(Entry)
                     : MF.end(),
-                Split);
-      Blocks.insert(Split);
+                Routing);
+      Blocks.insert(Routing);
 
       // Set the jump table's register of the index of the block we wish to
       // jump to, and jump to the jump table.
-      BuildMI(Split, DebugLoc(), TII.get(WebAssembly::CONST_I32), Reg)
+      BuildMI(Routing, DebugLoc(), TII.get(WebAssembly::CONST_I32), Reg)
           .addImm(Indices[Entry]);
-      BuildMI(Split, DebugLoc(), TII.get(WebAssembly::BR)).addMBB(Dispatch);
-      Split->addSuccessor(Dispatch);
-      Map[Entry] = Split;
-    }
+      BuildMI(Routing, DebugLoc(), TII.get(WebAssembly::BR)).addMBB(Dispatch);
+      Routing->addSuccessor(Dispatch);
+      Map[std::make_pair(PredInLoop, Entry)] = Routing;
+    }
+
+  for (auto *Pred : AllPreds) {
+    bool PredInLoop = InLoop.count(Pred);
     // Remap the terminator operands and the successor list.
     for (MachineInstr &Term : Pred->terminators())
       for (auto &Op : Term.explicit_uses())
         if (Op.isMBB() && Indices.count(Op.getMBB()))
-          Op.setMBB(Map[Op.getMBB()]);
-    for (auto Rewrite : Map)
-      Pred->replaceSuccessor(Rewrite.first, Rewrite.second);
+          Op.setMBB(Map[std::make_pair(PredInLoop, Op.getMBB())]);
+
+    for (auto *Succ : Pred->successors()) {
+      if (!Entries.count(Succ))
+        continue;
+      auto *Routing = Map[std::make_pair(PredInLoop, Succ)];
+      Pred->replaceSuccessor(Succ, Routing);
+    }
   }
 
   // Create a fake default label, because br_table requires one.
   MIB.addMBB(MIB.getInstr()
                  ->getOperand(MIB.getInstr()->getNumExplicitOperands() - 1)
                  .getMBB());
 }
 

llvm/trunk/test/CodeGen/WebAssembly/irreducible-cfg.ll

bb13:                                             ; preds = %bb10, %bb3
   store double %tmp17, double* %tmp16, align 4
   %tmp18 = add nsw i32 %tmp15, 1
   br label %bb6
 
 bb19:                                             ; preds = %bb6
   ret void
 }
 
-; A simple loop 2 blocks that are both entries.
+; A simple loop 2 blocks that are both entries: A1 and A2.
+; Even though A1 and A2 both have 3 predecessors (A0, A1, and A2), not 6 but
+; only 4 new routing blocks to the dispatch block should be generated.
 
 ; CHECK-LABEL: test2:
 ; CHECK: br_if
 ; CHECK: i32.const $[[REG:[^,]+]]=
+; CHECK: i32.const $[[REG]]=
 ; CHECK: br_table  $[[REG]],
+; CHECK: i32.const $[[REG]]=
+; CHECK: i32.const $[[REG]]=
+; CHECK-NOT: i32.const $[[REG]]=
 define i32 @test2(i32) {
 entry:
   br label %A0
 
 A0:                                               ; preds = %entry
   %a0a = tail call i32 @test2(i32 1)
   %a0b = icmp eq i32 %a0a, 0
   br i1 %a0b, label %A1, label %A2

llvm/trunk/test/CodeGen/WebAssembly/irreducible-cfg.mir

+# RUN: llc -mtriple=wasm32-unknown-unknown -run-pass wasm-fix-irreducible-control-flow %s -o - | FileCheck %s
+
+# This tests if we correctly create at most 2 routing blocks per entry block,
+# and also whether those routing blocks are generated in the correct place. If
+# one of the predecessor is the layout predecessor of an entry, a routing block
+# for the entry should be generated right after the layout predecessor.
+
+--- |
+  target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+  target triple = "wasm32-unknown-unknown"
+
+  define void @test0() {
+  pred0:
+    ret void
+  pred1:
+    ret void
+  entry0:
+    ret void
+  entry1:
+    ret void
+  }
+...
+
+---
+# CHECK-LABEL: test0
+name: test0
+liveins:
+  - { reg: '$arguments' }
+body: |
+  bb.0.pred0:
+    successors: %bb.1, %bb.2
+    liveins: $arguments
+    %0:i32 = CONST_I32 100, implicit-def $arguments
+    BR_IF %bb.2, %0:i32, implicit-def $arguments
+  ; CHECK: bb.0.pred0:
+    ; CHECK: BR_IF %bb.2, %0, implicit-def $arguments
+
+  bb.1.pred1:
+  ; predecessors: %bb.1
+    successors: %bb.2, %bb.3
+    BR_IF %bb.3, %0:i32, implicit-def $arguments
+  ; CHECK: bb.1.pred1:
+    ; CHECK: BR_IF %bb.7, %0, implicit-def $arguments
+    ; This falls through to bb.2, so we don't need an additional BR here
+    ; CHECK-NOT: BR
+
+  ; Routing block for entry0, when predecessor is outside the loop
+  ; This routing block is shared between the two predecessors: pred0 and pred1.
+  ; CHECK: bb.2:
+    ; CHECK: %1:i32 = CONST_I32 0, implicit-def $arguments
+    ; CHECK: BR %bb.6, implicit-def $arguments
+
+  bb.2.entry0:
+  ; predecessors: %bb.0, %bb.1, %bb.1
+    successors: %bb.3
+    BR %bb.3, implicit-def $arguments
+  ; CHECK: bb.3.entry0:
+    ; CHECK: BR %bb.4, implicit-def $arguments
+
+  ; Routing block for entry1, when predecessor is inside the loop
+  ; CHECK: bb.4:
+    ; CHECK: %1:i32 = CONST_I32 1, implicit-def $arguments
+    ; CHECK: BR %bb.6, implicit-def $arguments
+
+  bb.3.entry1:
+  ; predecessors: %bb.1, %bb.2
+    successors: %bb.2
+    BR %bb.2, implicit-def $arguments
+  ; CHECK: bb.5.entry1:
+    ; CHECK: BR %bb.8, implicit-def $arguments
+
+  ; Dispatch block
+  ; CHECK: bb.6:
+    ; CHECK: BR_TABLE_I32 %1, %bb.3, %bb.5, %bb.5, implicit-def $arguments
+
+  ; Routing block for entry1, when predecessor is outside the loop
+  ; CHECK: bb.7:
+    ; CHECK: %1:i32 = CONST_I32 1, implicit-def $arguments
+    ; CHECK: BR %bb.6, implicit-def $arguments
+
+  ; Routing block for entry0, when predecessor is inside the loop
+  ; CHECK: bb.8:
+    ; CHECK: %1:i32 = CONST_I32 0, implicit-def $arguments
+    ; CHECK: BR %bb.6, implicit-def $arguments