[WebAssembly] Irreducible control flow rewrite

[kripken]

Summary

Rewrite WebAssemblyFixIrreducibleControlFlow to a simpler and cleaner design,
which directly computes reachability and other properties itself. This avoids
previous complexity and bugs. (The new graph analyses are very similar to how
the Relooper algorithm would find loop entries and so forth.)

This fixes a few bugs, including where we had a false positive and thought
fannkuch was irreducible when it was not, which made us much larger and
slower there, and a reverse bug where we missed irreducibility. On fannkuch,
we used to be 44% slower than asm2wasm and are now 4% faster.

[aheejin]

Sorry for the delayed reply!

What changes fixed the previous bug?

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

30

Why is this?

144

Real nit: We can merge these into a single line:

MachineBasicBlock *MBB, *Succ;

158

We used to have the description on the definitions of 'loopers' and 'non-loopers' in the comment at the beginning of the file, but not anymore. Can we have that again here maybe?

167

Does this definition cover when Entry is an entry (or header) of a loop? I'm wondering because the reachability computation here does not consider backedges to Entry, so it will not be discovered as a looper.

Also does it cover this case?

bb0:
  br bb1
bb1:
  br_if bb0
  br bb2
bb2:
  br bb3
bb3:
  br bb2

So here bb0 and bb1 comprise a loop and bb2 and bb3 comprise another loop. And all four blocks are loopers, but bb0 and bb2 are loop entries. But they are not reachable from any of non-loopers, because there are no non-loopers in this CFG.
(Also here bb0 is the header of a loop so it belongs to the first case where Entry is a header)

201

Why is this worklist a set?

222

Do we need to make this as a separate class now? Unlike the previous version, we create LoopFixer object only once in WebAssemblyFixIrreducibleControlFlow::runOnMachineFunction. I think we can add paste the body of LoopFixer::run into WebAssemblyFixIrreducibleControlFlow::runOnMachineFunction and make LoopFixer:: processRegion as a member function of WebAssemblyFixIrreducibleControlFlow that returns a bool.

270

I might be mistaken, but the algorithm seems to not care about whether loop entries found are within the same level loop. For example, the case where there are two-level nested loops and block A is one of the entries to the outer loop and B is one of the entries for the inner loop. Even in this case, we don't distinguish them and resolve these with a single pass with a single dispatch block. Is this OK and intended? And for the 'mutual entries' above, can't they also be not in the same loop but each in an outer loop and an inner loop?

275

• How about reducing depth of the subloop processing routine below by wrapping this loop here early, like

while (FoundIrreducibility) {
  ReachabilityGraph Graph(Entry, Blocks);

  bool FoundIrreducibility = false;
  for (auto *LoopEntry : Graph.getLoopEntries()) {
    ...
  }
}

for (auto *LoopEntry : Graph.getLoopEntries()) {
  ...
}

• Btw, is there any case we can find irreducibility again with the same region before we recurse into inner loops below? If there is, how is that case different from the case that's processed below on the inner loops?

286

In which case do we need to recurse into inner loops? (Given that we didn't distinguish inner loops and outer loops above) Could you add some comments?

303

Nit: We can do in a single line

BlockVector SortedEntries(Entries.begin(), Entries.end());

347

Is there any case !Pair.second is not set? Aren't all blocks in SortedEntries distinct?

371

This loop does not seem to check if a corresponding Split block for Entry has been already created. In case Entry has more than one predecessors, Split for Entry is gonna be created multiple times and Map[Entry] will be overwritten. We should do something like this inside the loop:

if (Map.count(Entry))
  continue;

378

Did you mean

Entry->isLayoutSuccessor(Pred)

? In the current state all new blocks are gonna be appended at the end, because this condition is hardly likely to be satisfied.

390

Nit: There is a [[ https://github.com/llvm/llvm-project/blob/a946997c2482e4386549ee38b4bb154eb58efbb6/llvm/include/llvm/CodeGen/MachineInstrBuilder.h#L405-L410 | BuildMI function ]] that defaults to append instruction at the end, so you can just do

BuildMI(Split, DebugLoc(), TII.get(WebAssembly::CONST_I32), Reg)
    .addImm(Indices[Entry]);                                              
BuildMI(Split, DebugLoc(), TII.get(WebAssembly::BR)).addMBB(Dispatch);

(Note that Split does not have * anymore)

test/CodeGen/WebAssembly/irreducible-cfg.ll

224

• If this is without irreducibility, shouldn't this be in non-irreducible-cfg.ll below?
• Can we remove #0?

[kripken]

What changes fixed the previous bug?

Basically the rewrite avoids the entire previous approach of "canonicalization", of representing an inner loop by its entry. When we got that wrong, we could reach very wrong results.

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

30

The NumLoops^2 part? To find mutual loop entries, that is, loop entries that can reach each other, we compare each one to the others.

This is actually just the square of the number of loops in a single scope, so nested loops don't count. Anyhow, even in the "worst case" which is a loop after a loop after a loop, this factor is going to be quite small, I'd guess.

167

When we look at a region, the entry block cannot be a looper, which simplifies things. LLVM doesn't let a function entry be in a loop (maybe for similar reasons as why it's convenient here?), and when we recurse into a loop we ignore backedges to Entry intentionally (so we see the inner part of the loop, and not the looping).

In other words, if a block is a looper, it will be seen as such in the outer scope it is in (where it is not the entry). When it is the entry into a region, and we are looking at just that region, it cannot be a looper there.

For the second issue (with that example): yeah, the comment was not clear enough, I'll update it. They key is that a loop enterer is a block not in the same loop, that enters it. So here bb2 is a looper, and bb1 is a predecessor of it, and not in the same loop (bb2 can't reach bb1), so we identify bb2 as a loop entry and bb1 as a loop enterer for it.

201

As a set, it won't have duplicate entries - so if we visit A that causes us to add C to the future work, and then visit B that also causes us to add C to the future work, then when we get to C we'll visit it once, and avoid visiting it again later (and not having anything to do the second time).

I haven't benchmarked this recently, but I remember it was beneficial in the Relooper.

222

Good idea, thanks!

270

Yeah, it's somewhat debatable (is an inner loop an "inner loop" if there's a branch going into it from several levels up?), but yes, when we see irreducibility we handle it now, instead of leaving it for the inner loop. We have to, I think - if we can branch to it now, we need to fix up that branch now - we wouldn't see it later when we recurse and ignore the outside. At least that would require a very different approach I suppose (which might be more optimal in some cases).

275

The problem is that the Graph is used outside of the loop in your code. I agree it's a little awkward as it is... another option is to use an std::unique_ptr to allow using it from the outside. But that's less clear I think.

Yes, we might find irreducibility again if we have two irreducible loops one after the other (i.e. not nested). It's *possible* we don't need to look at the irreducible one we've already fixed, however, I don't have a proof of that. And we do need to recompute the Graph again after fixing one loop (since it can affect others; and trying to update it is fairly complex an optimization), so starting another loop iteration is the simple thing. The cost is another iteration per irreducible loop, which is not that bad since they are rare.

I'll try to improve the comment here.

286

Good point, this is important to explain, adding.

347

This is code not changed in this patch (just some variable renaming), but I think you're right, fixing.

371

This is code that did not change in this patch (aside from names and clang-format), and I'm not sure I fully understand it (I think Dan wrote it?). But I don't think Map[Entry] can be overwritten? The shape of the code is

DenseMap<MachineBasicBlock *, MachineBasicBlock *> Map;
for (auto *Entry : Pred->successors()) {
  [..]
  Map[Entry] = Split;
}

378

As before, this is Dan's code that I didn't change here (aside from variable names). I'm not sure what it does...

test/CodeGen/WebAssembly/irreducible-cfg.ll

224

• The other file has -O0, and this one doesn't specify an opt level (so it's -O2 I guess?), so I kept them separate.
• Removed #0.

[kripken]

Review feedback changes, in particular:

• Remove LoopFixer class.
• Various comment additions and improvements.
• Various nits.

[aheejin]

Real nit: We usually start CL/commit headers with [WebAssembly]

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

64

• Not sure what's happened, but it says "Context not available" on the diff on these lines. Could you rebase your CL onto the tip of the repo?
• We usually include WebAssembly.h in every pass, and I guess we need to include (Maybe it's not shown here because of the diff problem?)

[aheejin]

• There are several places in the diff that says "Context not available". I think the diff has some problems. (I thought it needed rebase but maybe it's not relevant) Could you check?
• In this new algorithm, all tests in irreducible-cfg-nested.ll, irreducible-cfg-nested2.ll, and irreducible-cfg-exceptions.ll are reducible now, i.e. this pass does not do anything on them. But some of them are irreducible if we pass -O0. (I only checked func_2 in irreducible-cfg-nested2.ll) I guess we should check these tests and either delete the ones that are reducible or change the flags to llc to -O0 or something to make them irreducible again. Also test3 in irredducible-cfg.ll is now considered as reducible too.
• On second thought, wouldn't it better to just specify -O0 for all irreducibility tests? -O2 basically means we don't now what would happen to our carefully created irreducible CFG.

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

167

Oh I see. Thanks. I didn't know that function entries cannot be loop headers.

201

But isn't it possible that by the time you try to add C for the second time, the first C has been already erased from the set so you do the unnecessary work again? Usually what I've seen was making worklist as a vector and have a separate set called something like Visited, to which visited blocks are added and not erased.

230

Nit: Can we take out this and makeSingleEntryLoop from the class definition? I don't think we want to inline these long functions and we can have one less level of indentation outside. I mean, outside the class,

bool WebAssemblyFixIrreducibleControlFlow::processRegion(...) {
  ...
}

270

Is there a condition in the code that we only consider inner loops that have entries from several levels up? I don't think I found one. We just treat all blocks within the region the same way.

So, for example, an outer loop consists of blocks A, B, C, D, E, and F, and among them E and F comprises an inner loop. The outer loop has two entries: A and D from the blocks outside, and the inner loop also has two entries: E and F. E's enterer is C and F's enterer is D. In this case, can we solve this case in one pass with a single dispatch block that branches into all of A, D, E, and F?

I guess I don't understand the algorithm :( Sorry for that. Please correct me if I'm mistaken.

275

I see. I missed the graph computation was there too.

371

Ah right, it's not gonna be overwritten. So the code structure is

for (MachineBasicBlock *Pred : AllPreds) {
  DenseMap<MachineBasicBlock *, MachineBasicBlock *> Map;
  for (auto *Entry : Pred->successors()) {
    ...
    Map[Entry] = Split;
  }

And in case we have entries E0 and E1, and there are predecessors P0 and P1, and the edges are like P0->E0, P0->E1, P1->E0, and P1->E1. So the both preds point to the both entries. In this case, each entry (E0 or E1) is processed twice in the inner for loop. Because Map is created within the outer loop, it's not gonna be overwritten, but Split for each E0 and E1 is gonna be created twice unnecessarily.

(I know it's not the part of the code that's changed in this CL; but I thought while we're writing the major part of the pass, maybe we can fix other things too. If you don't prefer that maybe I can submit that as a separate patch. Let me know which one you prefer.)

378

(Because code parts have moved, this comment is not where it was anymore, so)
The code is

// This is a successor we need to rewrite.                               
MachineBasicBlock *Split = MF.CreateMachineBasicBlock();                 
MF.insert(Pred->isLayoutSuccessor(Entry)                                 
              ? MachineFunction::iterator(Entry)                         
              : MF.end(),                                                
          Split);

I think the intention of the original code was, when we insert the Split which is to serve as a stepping stone between Pred and Entry, if currently Pred is right before Entry in the BB order, we want to insert Split between them, which makes a natural order. If not, we just append the new Split at the end of the function. I think this bug was there all along for years, and this was all newly created blocks are appended at the end.

I know this is not related to this CL, but if you prefer you can fix it maybe, or I can submit that as a separate patch.

390

Ping

test/CodeGen/WebAssembly/irreducible-cfg.ll

224

• You can use two different commands and check them separately within a single file.

; RUN: llc < %s -O0 -asm-verbose=false -verify-machineinstrs -disable-block-placement -wasm-disable-explicit-locals -wasm-keep-registers | FileCheck %s
; RUN: llc < %s -O2 -asm-verbose=false -verify-machineinstrs -disable-block-placement -wasm-disable-explicit-locals -wasm-keep-registers | FileCheck -check-prefix=OPT %s

; OPT-NOT: br_table
define hidden void @ps_hints_apply() {
  ...
}

; CHECK-NOT: br_table
define hidden i32 @_Z15fannkuch_workerPv(i8* %_arg) {
  ...
}

Here is the manual for this -check-prefix option.

• Nit: Can we remove hidden too?

test/CodeGen/WebAssembly/non-irreducible-cfg.ll

13

Nit: Can we remove hidden and #0?

[aheejin]

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

259

Nit: How about renaming this to EntriesForSingleLoop or something and start by inserting LoopEntry first? (I don't really like the name EntriesForSingleLoop, but can't think of a better alternative now... I guess there can be better ones)

BlockSet EntriesForSingleLoop;
EntriesForSingleLoop.insert(LoopEntry);
for (auto *OtherLoopEntry : Graph.getLoopEntries()) {
  if (Graph.canReach(LoopEntry, OtherLoopEntry) &&
      Graph.canReach(OtherLoopEntry, LoopEntry))
    EntriesForSingleLoop.insert(OhterLoopEntry);
}

if (EntriesForSingleLoop.size() > 1) {
  makeSingleEntryLoop(EntriesForSingleLoop, Blocks, MF);
  ...
}
...

That way we don't need to std::move this later and don't need to check if (OtherLoopEntry != LoopEntry). And in my opinion if (EntriesForSingleLoop.size() > 1) makes it clear that we take a multiple entry loop and transform it into a single entry loop.

This is just a matter of style preference, so if you don't like it, please ignore it.

270

Oh, nevermind. I think I got it. :) Inner loop entries that are not reachable from outside will not be identified as loop entries when we examine the outer region, because those inner loop entries can also reach their preds, as long as the preds are not from outside and in an outer loop.

292

we may only alter branches 'in' existing..?

test/CodeGen/WebAssembly/irreducible-cfg.ll

224

On second thought, wouldn't it better to just specify -O0 for all irreducibility tests? -O2 basically means we don't now what would happen to our carefully created irreducible CFG.

[kripken]

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

64

Rebasing - hopefully that fixes the context issue.

Yeah, WebAssembly.h was included already, so might be the diff context problem.

201

The way you suggest is probably faster overall, I agree, I'll change it to that.

259

I agree the move is not nice. I think we can just add LoopEntry to MutualLoopEntries and simplify things that way, updating the patch with that, let me know what you think.

292

Thanks, I'll fix that.

371

I'd prefer we do it separately since this patch already rewrites all the other code in this file, aside from that one unchanged function. Seems safer to change it later.

390

Same issue as before - I'd prefer to not modify this one function in this patch (which changes everything aside from that one function), for safety.

(I did rename a bunch of variables, because I wanted to keep names consistent throughout the file - that seemed necessary in this patch.)

test/CodeGen/WebAssembly/irreducible-cfg.ll

224

Removed hidden, and changed to -O0 in irreducible-cfg. Yeah, that seems like the right thing to do for irreducibility tests.

I didn't want to change irreducible-cfg-exceptions.ll's optimization level though because I'm not sure that wouldn't invalidate it.

[kripken]

Various changes from the feedback:

• Change how the LoopBlocks work list works, using a set+vector instead of just a set.
• Move large functions out of the class definition.
• Avoid a move in MutualEntries, instead have them always contain all the mutual entries (which includes the original entry, as it is mutual to itself).
• Test cleanup; use -O0 for the main irreducibility tests; remove old irreducibility tests which were invalid.
• Comment improvements.

[aheejin]

• Sorry, but at least some of the tests in irreducible-cfg-nested.ll and irreducible-cfg-nested2.ll were still irreducible if we passed -O0. Do we want to delete all of them? Shouldn't we keep the ones that are still irreducible with -O0?
• Real nit: We usually start CL/commit headers with [WebAssembly]

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

259

I think this is better, thanks.

[aheejin]

test/CodeGen/WebAssembly/irreducible-cfg.ll

288

Nit: newline at the end

[kripken]

Updates:

• Thanks, indeed one of the tests removed before (func_2) was still irreducible at -O0, restored it into the main irreducible file test, which is now at -O0. The other (tre_parse) was actually not irreducible, so we mis-identified it, and I left it out.
• Fixed extra newline at the end of a test.
• Renamed the title to have [WebAssembly]

[aheejin]

LGTM, thank you. Do you want me to commit this for you? Or are you gonna get a commit access?

[kripken]

Great, and thanks for the comments!

I should probably get commit access, yeah :) Can you please commit this one, though, so it doesn't wait on that?

[aheejin]

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

378

Just letting you know, I was incorrect on this part; the code was correct. I realized that while I was working on D59462..

lib/Target/WebAssembly/WebAssemblyFixIrreducibleControlFlow.cpp

 //=- WebAssemblyFixIrreducibleControlFlow.cpp - Fix irreducible control flow -//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
-/// This file implements a pass that transforms irreducible control flow into
-/// reducible control flow. Irreducible control flow means multiple-entry
-/// loops; they appear as CFG cycles that are not recorded in MachineLoopInfo
-/// due to being unnatural.
+/// This file implements a pass that removes irreducible control flow.
+/// Irreducible control flow means multiple-entry loops, which this pass
+/// transforms to have a single entry.
 ///
 /// Note that LLVM has a generic pass that lowers irreducible control flow, but
 /// it linearizes control flow, turning diamonds into two triangles, which is
 /// both unnecessary and undesirable for WebAssembly.
 ///
-/// The big picture: Ignoring natural loops (seeing them monolithically), we
-/// find all the blocks which can return to themselves ("loopers"). Loopers
-/// reachable from the non-loopers are loop entries: if there are 2 or more,
-/// then we have irreducible control flow. We fix that as follows: a new block
-/// is created that can dispatch to each of the loop entries, based on the
-/// value of a label "helper" variable, and we replace direct branches to the
-/// entries with assignments to the label variable and a branch to the dispatch
-/// block. Then the dispatch block is the single entry in a new natural loop.
+/// The big picture: We recursively process each "region", defined as a group
+/// of blocks with a single entry and no branches back to that entry. A region
+/// may be the entire function body, or the inner part of a loop, i.e., the
+/// loop's body without branches back to the loop entry. In each region we fix
+/// up multi-entry loops by adding a new block that can dispatch to each of the
+/// loop entries, based on the value of a label "helper" variable, and we
+/// replace direct branches to the entries with assignments to the label
+/// variable and a branch to the dispatch block. Then the dispatch block is the
+/// single entry in the loop containing the previous multiple entries. After
+/// ensuring all the loops in a region are reducible, we recurse into them. The
+/// total time complexity of this pass is:
+///   O(NumBlocks * NumNestedLoops * NumIrreducibleLoops +
+///     NumLoops * NumLoops)
 ///
-/// This is similar to what the Relooper [1] does, both identify looping code
-/// that requires multiple entries, and resolve it in a similar way. In
-/// Relooper terminology, we implement a Multiple shape in a Loop shape. Note
+/// This pass is similar to what the Relooper [1] does. Both identify looping
+/// code that requires multiple entries, and resolve it in a similar way (in
+/// Relooper terminology, we implement a Multiple shape in a Loop shape). Note
 /// also that like the Relooper, we implement a "minimal" intervention: we only
 /// use the "label" helper for the blocks we absolutely must and no others. We
-/// also prioritize code size and do not perform node splitting (i.e. we don't
-/// duplicate code in order to resolve irreducibility).
-///
-/// The difference between this code and the Relooper is that the Relooper also
-/// generates ifs and loops and works in a recursive manner, knowing at each
-/// point what the entries are, and recursively breaks down the problem. Here
-/// we just want to resolve irreducible control flow, and we also want to use
-/// as much LLVM infrastructure as possible. So we use the MachineLoopInfo to
-/// identify natural loops, etc., and we start with the whole CFG and must
-/// identify both the looping code and its entries.
+/// also prioritize code size and do not duplicate code in order to resolve
+/// irreducibility. The graph algorithms for finding loops and entries and so
+/// forth are also similar to the Relooper. The main differences between this
+/// pass and the Relooper are:
+///  * We just care about irreducibility, so we just look at loops.
+///  * The Relooper emits structured control flow (with ifs etc.), while we
+///    emit a CFG.
 ///
 /// [1] Alon Zakai. 2011. Emscripten: an LLVM-to-JavaScript compiler. In
 /// Proceedings of the ACM international conference companion on Object oriented
 /// programming systems languages and applications companion (SPLASH '11). ACM,
 /// New York, NY, USA, 301-312. DOI=10.1145/2048147.2048224
 /// http://doi.acm.org/10.1145/2048147.2048224
 ///
 //===----------------------------------------------------------------------===//
 
 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
 #include "WebAssembly.h"
 #include "WebAssemblyMachineFunctionInfo.h"
 #include "WebAssemblySubtarget.h"
 #include "llvm/ADT/PriorityQueue.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "wasm-fix-irreducible-control-flow"
 
 namespace {
 
-class LoopFixer {
+using BlockVector = SmallVector<MachineBasicBlock *, 4>;
+using BlockSet = SmallPtrSet<MachineBasicBlock *, 4>;
+
+// Calculates reachability in a region. Ignores branches to blocks outside of
+// the region, and ignores branches to the region entry (for the case where
+// the region is the inner part of a loop).
+class ReachabilityGraph {
 public:
-  LoopFixer(MachineFunction &MF, MachineLoopInfo &MLI, MachineLoop *Loop)
-      : MF(MF), MLI(MLI), Loop(Loop) {}
+  ReachabilityGraph(MachineBasicBlock *Entry, const BlockSet &Blocks)
+      : Entry(Entry), Blocks(Blocks) {
+#ifndef NDEBUG
+    // The region must have a single entry.
+    for (auto *MBB : Blocks) {
+      if (MBB != Entry) {
+        for (auto *Pred : MBB->predecessors()) {
+          assert(inRegion(Pred));
+        }
+      }
+    }
+#endif
+    calculate();
+  }
+
+  bool canReach(MachineBasicBlock *From, MachineBasicBlock *To) {
+    assert(inRegion(From) && inRegion(To));
+    return Reachable[From].count(To);
+  }
+
+  // Get all blocks that are in a loop.
+  const BlockSet &getLoopers() { return Loopers; }
 
-  // Run the fixer on the given inputs. Returns whether changes were made.
-  bool run();
+  // Get all blocks that are loop entries.
+  const BlockSet &getLoopEntries() { return LoopEntries; }
+
+  // Get all blocks that enter a particular loop from outside.
+  const BlockSet &getLoopEnterers(MachineBasicBlock *LoopEntry) {
+    assert(inRegion(LoopEntry));
+    return LoopEnterers[LoopEntry];
+  }
 
 private:
-  MachineFunction &MF;
-  MachineLoopInfo &MLI;
-  MachineLoop *Loop;
+  MachineBasicBlock *Entry;
+  const BlockSet &Blocks;
 
-  MachineBasicBlock *Header;
-  SmallPtrSet<MachineBasicBlock *, 4> LoopBlocks;
+  BlockSet Loopers, LoopEntries;
+  DenseMap<MachineBasicBlock *, BlockSet> LoopEnterers;
 
-  using BlockSet = SmallPtrSet<MachineBasicBlock *, 4>;
+  bool inRegion(MachineBasicBlock *MBB) { return Blocks.count(MBB); }
+
+  // Maps a block to all the other blocks it can reach.
   DenseMap<MachineBasicBlock *, BlockSet> Reachable;
 
-  // The worklist contains pairs of recent additions, (a, b), where we just
-  // added a link a => b.
+  void calculate() {
+    // Reachability computation work list. Contains pairs of recent additions
+    // (A, B) where we just added a link A => B.
     using BlockPair = std::pair<MachineBasicBlock *, MachineBasicBlock *>;
     SmallVector<BlockPair, 4> WorkList;
 
-  // Get a canonical block to represent a block or a loop: the block, or if in
-  // an inner loop, the loop header, of it in an outer loop scope, we can
-  // ignore it. We need to call this on all blocks we work on.
-  MachineBasicBlock *canonicalize(MachineBasicBlock *MBB) {
-    MachineLoop *InnerLoop = MLI.getLoopFor(MBB);
-    if (InnerLoop == Loop) {
-      return MBB;
-    } else {
-      // This is either in an outer or an inner loop, and not in ours.
-      if (!LoopBlocks.count(MBB)) {
-        // It's in outer code, ignore it.
-        return nullptr;
-      }
-      assert(InnerLoop);
-      // It's in an inner loop, canonicalize it to the header of that loop.
-      return InnerLoop->getHeader();
+    // Add all relevant direct branches.
+    for (auto *MBB : Blocks) {
+      for (auto *Succ : MBB->successors()) {
+        if (Succ != Entry && inRegion(Succ)) {
+          Reachable[MBB].insert(Succ);
+          WorkList.emplace_back(MBB, Succ);
+        }
       }
     }
 
-  // For a successor we can additionally ignore it if it's a branch back to a
-  // natural loop top, as when we are in the scope of a loop, we just care
-  // about internal irreducibility, and can ignore the loop we are in. We need
-  // to call this on all blocks in a context where they are a successor.
-  MachineBasicBlock *canonicalizeSuccessor(MachineBasicBlock *MBB) {
-    if (Loop && MBB == Loop->getHeader()) {
-      // Ignore branches going to the loop's natural header.
-      return nullptr;
-    }
-    return canonicalize(MBB);
-  }
-
-  // Potentially insert a new reachable edge, and if so, note it as further
-  // work.
-  void maybeInsert(MachineBasicBlock *MBB, MachineBasicBlock *Succ) {
-    assert(MBB == canonicalize(MBB));
-    assert(Succ);
-    // Succ may not be interesting as a sucessor.
-    Succ = canonicalizeSuccessor(Succ);
-    if (!Succ)
-      return;
-    if (Reachable[MBB].insert(Succ).second) {
-      // For there to be further work, it means that we have
-      //   X => MBB => Succ
-      // for some other X, and in that case X => Succ would be a new edge for
-      // us to discover later. However, if we don't care about MBB as a
-      // successor, then we don't care about that anyhow.
-      if (canonicalizeSuccessor(MBB)) {
-        WorkList.emplace_back(MBB, Succ);
+    while (!WorkList.empty()) {
+      MachineBasicBlock *MBB;
+      MachineBasicBlock *Succ;
+      std::tie(MBB, Succ) = WorkList.pop_back_val();
+      assert(inRegion(MBB) && Succ != Entry && inRegion(Succ));
+      if (MBB != Entry) {
+        // We recently added MBB => Succ, and that means we may have enabled
+        // Pred => MBB => Succ.
+        for (auto *Pred : MBB->predecessors()) {
+          if (Reachable[Pred].insert(Succ).second) {
+            WorkList.emplace_back(Pred, Succ);
+          }
         }
       }
     }
-};
 
-bool LoopFixer::run() {
-  Header = Loop ? Loop->getHeader() : &*MF.begin();
+    // Blocks that can return to themselves are in a loop.
+    for (auto *MBB : Blocks) {
+      if (canReach(MBB, MBB)) {
+        Loopers.insert(MBB);
+      }
+    }
+    assert(!Loopers.count(Entry));
 
-  // Identify all the blocks in this loop scope.
-  if (Loop) {
-    for (auto *MBB : Loop->getBlocks()) {
-      LoopBlocks.insert(MBB);
+    // Find the loop entries - loopers reachable from non-loopers - and their
+    // loop enterers.
+    for (auto *Looper : Loopers) {
+      for (auto *Pred : Looper->predecessors()) {
+        // Pred can reach Looper. If Looper can reach Pred, it is in the loop;
+        // otherwise, it is a block that enters into the loop.
+        if (!canReach(Looper, Pred)) {
+          LoopEntries.insert(Looper);
+          LoopEnterers[Looper].insert(Pred);
         }
-  } else {
-    for (auto &MBB : MF) {
-      LoopBlocks.insert(&MBB);
       }
     }
+  }
+};
 
-  // Compute which (canonicalized) blocks each block can reach.
+// Finds the blocks in a single-entry loop, given the loop entry and the
+// list of blocks that enter the loop.
+class LoopBlocks {
+public:
+  LoopBlocks(MachineBasicBlock *Entry, const BlockSet &Enterers)
+      : Entry(Entry), Enterers(Enterers) {
+    calculate();
+  }
 
-  // Add all the initial work.
-  for (auto *MBB : LoopBlocks) {
-    MachineLoop *InnerLoop = MLI.getLoopFor(MBB);
+  BlockSet &getBlocks() { return Blocks; }
 
-    if (InnerLoop == Loop) {
-      for (auto *Succ : MBB->successors()) {
-        maybeInsert(MBB, Succ);
-      }
-    } else {
-      // It can't be in an outer loop - we loop on LoopBlocks - and so it must
-      // be an inner loop.
-      assert(InnerLoop);
-      // Check if we are the canonical block for this loop.
-      if (canonicalize(MBB) != MBB) {
-        continue;
-      }
-      // The successors are those of the loop.
-      SmallVector<MachineBasicBlock *, 2> ExitBlocks;
-      InnerLoop->getExitBlocks(ExitBlocks);
-      for (auto *Succ : ExitBlocks) {
-        maybeInsert(MBB, Succ);
-      }
+private:
+  MachineBasicBlock *Entry;
+  const BlockSet &Enterers;
+
+  BlockSet Blocks;
+
+  void calculate() {
+    // Going backwards from the loop entry, if we ignore the blocks entering
+    // from outside, we will traverse all the blocks in the loop.
+    BlockSet WorkList;
+    Blocks.insert(Entry);
+    for (auto *Pred : Entry->predecessors()) {
+      if (!Enterers.count(Pred)) {
+        WorkList.insert(Pred);
       }
     }
 
-  // Do work until we are all done.
     while (!WorkList.empty()) {
-    MachineBasicBlock *MBB;
-    MachineBasicBlock *Succ;
-    std::tie(MBB, Succ) = WorkList.pop_back_val();
-    // The worklist item is an edge we just added, so it must have valid blocks
-    // (and not something canonicalized to nullptr).
-    assert(MBB);
-    assert(Succ);
-    // The successor in that pair must also be a valid successor.
-    assert(MBB == canonicalizeSuccessor(MBB));
-    // We recently added MBB => Succ, and that means we may have enabled
-    // Pred => MBB => Succ. Check all the predecessors. Note that our loop here
-    // is correct for both a block and a block representing a loop, as the loop
-    // is natural and so the predecessors are all predecessors of the loop
-    // header, which is the block we have here.
+      auto *MBB = *WorkList.begin();
+      WorkList.erase(MBB);
+      assert(!Enterers.count(MBB));
+      if (Blocks.insert(MBB).second) {
         for (auto *Pred : MBB->predecessors()) {
-      // Canonicalize, make sure it's relevant, and check it's not the same
-      // block (an update to the block itself doesn't help compute that same
-      // block).
-      Pred = canonicalize(Pred);
-      if (Pred && Pred != MBB) {
-        maybeInsert(Pred, Succ);
+          WorkList.insert(Pred);
         }
       }
     }
+  }
+};
 
-  // It's now trivial to identify the loopers.
-  SmallPtrSet<MachineBasicBlock *, 4> Loopers;
-  for (auto MBB : LoopBlocks) {
-    if (Reachable[MBB].count(MBB)) {
-      Loopers.insert(MBB);
+class LoopFixer {
+public:
+  LoopFixer(MachineFunction &MF) : MF(MF) {}
+
+  // Run on the given input. Returns whether changes were made.
+  bool run() {
+    // Start the recursive process on the entire function body.
+    BlockSet AllBlocks;
+    for (auto &MBB : MF) {
+      AllBlocks.insert(&MBB);
     }
+    processRegion(&*MF.begin(), AllBlocks);
+    return Changed;
   }
-  // The header cannot be a looper. At the toplevel, LLVM does not allow the
-  // entry to be in a loop, and in a natural loop we should ignore the header.
-  assert(Loopers.count(Header) == 0);
-
-  // Find the entries, loopers reachable from non-loopers.
-  SmallPtrSet<MachineBasicBlock *, 4> Entries;
-  SmallVector<MachineBasicBlock *, 4> SortedEntries;
-  for (auto *Looper : Loopers) {
-    for (auto *Pred : Looper->predecessors()) {
-      Pred = canonicalize(Pred);
-      if (Pred && !Loopers.count(Pred)) {
-        Entries.insert(Looper);
-        SortedEntries.push_back(Looper);
+
+private:
+  MachineFunction &MF;
+
+  bool Changed = false;
+
+  void processRegion(MachineBasicBlock *Entry, BlockSet &Blocks) {
+    // Remove irreducibility before processing child loops, which may take
+    // multiple iterations.
+    while (true) {
+      ReachabilityGraph Graph(Entry, Blocks);
+
+      bool FoundIrreducibility = false;
+
+      for (auto *LoopEntry : Graph.getLoopEntries()) {
+        // Find mutual entries - other entries which can reach this one, and are
+        // reached by it. Such mutual entries must be in the same loop, and so
+        // indicate irreducible control flow.
+        BlockSet MutualLoopEntries;
+        for (auto *OtherLoopEntry : Graph.getLoopEntries()) {
+          if (OtherLoopEntry != LoopEntry &&
+              Graph.canReach(LoopEntry, OtherLoopEntry) &&
+              Graph.canReach(OtherLoopEntry, LoopEntry)) {
+            MutualLoopEntries.insert(OtherLoopEntry);
+          }
+        }
+
+        if (!MutualLoopEntries.empty()) {
+          auto AllLoopEntries = std::move(MutualLoopEntries);
+          AllLoopEntries.insert(LoopEntry);
+          makeSingleEntryLoop(AllLoopEntries, Blocks);
+          FoundIrreducibility = true;
           break;
         }
       }
+      // Only go on to actually process the inner loops when we are done
+      // removing irreducible control flow and changing the graph.
+      if (FoundIrreducibility) {
+        continue;
       }
 
-  // Check if we found irreducible control flow.
-  if (LLVM_LIKELY(Entries.size() <= 1))
-    return false;
+      for (auto *LoopEntry : Graph.getLoopEntries()) {
+        LoopBlocks InnerBlocks(LoopEntry, Graph.getLoopEnterers(LoopEntry));
+        // Each of these calls to processRegion may change the graph, but are
+        // guaranteed not to interfere with each other. The only changes we make
+        // to the graph are to add blocks on the way to a loop entry. As the
+        // loops are disjoint, that means we may only alter branches exiting
+        // another loop, which are ignored when recursing into that other loop
+        // anyhow.
+        processRegion(LoopEntry, InnerBlocks.getBlocks());
+      }
+
+      return;
+    }
+  }
+
+  // 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.
+  void makeSingleEntryLoop(BlockSet &Entries, BlockSet &Blocks) {
+    assert(Entries.size() >= 2);
 
     // Sort the entries to ensure a deterministic build.
+    BlockVector SortedEntries;
+    for (auto *Entry : Entries) {
+      SortedEntries.push_back(Entry);
+    }
     llvm::sort(SortedEntries,
                [&](const MachineBasicBlock *A, const MachineBasicBlock *B) {
                  auto ANum = A->getNumber();
                  auto BNum = B->getNumber();
                  return ANum < BNum;
                });
 
 #ifndef NDEBUG
     for (auto Block : SortedEntries)
       assert(Block->getNumber() != -1);
     if (SortedEntries.size() > 1) {
-    for (auto I = SortedEntries.begin(), E = SortedEntries.end() - 1;
-         I != E; ++I) {
+      for (auto I = SortedEntries.begin(), E = SortedEntries.end() - 1; I != E;
+           ++I) {
         auto ANum = (*I)->getNumber();
         auto BNum = (*(std::next(I)))->getNumber();
         assert(ANum != BNum);
       }
     }
 #endif
 
     // Create a dispatch block which will contain a jump table to the entries.
     MachineBasicBlock *Dispatch = MF.CreateMachineBasicBlock();
     MF.insert(MF.end(), Dispatch);
-  MLI.changeLoopFor(Dispatch, Loop);
+    Blocks.insert(Dispatch);
 
     // Add the jump table.
     const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
     MachineInstrBuilder MIB = BuildMI(*Dispatch, Dispatch->end(), DebugLoc(),
                                       TII.get(WebAssembly::BR_TABLE_I32));
 
     // Add the register which will be used to tell the jump table which block to
     // jump to.
     MachineRegisterInfo &MRI = MF.getRegInfo();
     unsigned Reg = MRI.createVirtualRegister(&WebAssembly::I32RegClass);
     MIB.addReg(Reg);
 
     // Compute the indices in the superheader, one for each bad block, and
     // add them as successors.
     DenseMap<MachineBasicBlock *, unsigned> Indices;
-  for (auto *MBB : SortedEntries) {
-    auto Pair = Indices.insert(std::make_pair(MBB, 0));
+    for (auto *Entry : SortedEntries) {
+      auto Pair = Indices.insert(std::make_pair(Entry, 0));
       if (!Pair.second) {
         continue;
       }
 
       unsigned Index = MIB.getInstr()->getNumExplicitOperands() - 1;
       Pair.first->second = Index;
 
-    MIB.addMBB(MBB);
-    Dispatch->addSuccessor(MBB);
+      MIB.addMBB(Entry);
+      Dispatch->addSuccessor(Entry);
     }
 
-  // Rewrite the problematic successors for every block that wants to reach the
-  // bad blocks. For simplicity, we just introduce a new block for every edge
-  // we need to rewrite. (Fancier things are possible.)
+    // Rewrite the problematic successors for every block that wants to reach
+    // the bad blocks. For simplicity, we just introduce a new block for every
+    // edge we need to rewrite. (Fancier things are possible.)
 
-  SmallVector<MachineBasicBlock *, 4> AllPreds;
-  for (auto *MBB : SortedEntries) {
-    for (auto *Pred : MBB->predecessors()) {
+    BlockVector AllPreds;
+    for (auto *Entry : SortedEntries) {
+      for (auto *Pred : Entry->predecessors()) {
         if (Pred != Dispatch) {
           AllPreds.push_back(Pred);
         }
       }
     }
 
-  for (MachineBasicBlock *MBB : AllPreds) {
+    for (MachineBasicBlock *Pred : AllPreds) {
       DenseMap<MachineBasicBlock *, MachineBasicBlock *> Map;
-    for (auto *Succ : MBB->successors()) {
-      if (!Entries.count(Succ)) {
+      for (auto *Entry : Pred->successors()) {
+        if (!Entries.count(Entry)) {
           continue;
         }
 
         // This is a successor we need to rewrite.
         MachineBasicBlock *Split = MF.CreateMachineBasicBlock();
-      MF.insert(MBB->isLayoutSuccessor(Succ) ? MachineFunction::iterator(Succ)
-                                             : MF.end(),
+        MF.insert(Pred->isLayoutSuccessor(Entry)
+                      ? MachineFunction::iterator(Entry)
                   Split);
-      MLI.changeLoopFor(Split, Loop);
+        Blocks.insert(Split);
 
         // 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, Split->end(), DebugLoc(), TII.get(WebAssembly::CONST_I32),
-              Reg)
-          .addImm(Indices[Succ]);
+        BuildMI(*Split, Split->end(), DebugLoc(),
+                TII.get(WebAssembly::CONST_I32), Reg)
+            .addImm(Indices[Entry]);
         BuildMI(*Split, Split->end(), DebugLoc(), TII.get(WebAssembly::BR))
             .addMBB(Dispatch);
         Split->addSuccessor(Dispatch);
-      Map[Succ] = Split;
+        Map[Entry] = Split;
       }
       // Remap the terminator operands and the successor list.
-    for (MachineInstr &Term : MBB->terminators())
+      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)
-      MBB->replaceSuccessor(Rewrite.first, Rewrite.second);
+        Pred->replaceSuccessor(Rewrite.first, Rewrite.second);
     }
 
     // Create a fake default label, because br_table requires one.
     MIB.addMBB(MIB.getInstr()
                    ->getOperand(MIB.getInstr()->getNumExplicitOperands() - 1)
                    .getMBB());
 
-  return true;
+    Changed = true;
   }
+};
 
 class WebAssemblyFixIrreducibleControlFlow final : public MachineFunctionPass {
   StringRef getPassName() const override {
     return "WebAssembly Fix Irreducible Control Flow";
   }
 
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.setPreservesCFG();
-    AU.addRequired<MachineDominatorTree>();
-    AU.addPreserved<MachineDominatorTree>();
-    AU.addRequired<MachineLoopInfo>();
-    AU.addPreserved<MachineLoopInfo>();
-    MachineFunctionPass::getAnalysisUsage(AU);
-  }
-
   bool runOnMachineFunction(MachineFunction &MF) override;
 
-  bool runIteration(MachineFunction &MF, MachineLoopInfo &MLI) {
-    // Visit the function body, which is identified as a null loop.
-    if (LoopFixer(MF, MLI, nullptr).run()) {
-      return true;
-    }
-
-    // Visit all the loops.
-    SmallVector<MachineLoop *, 8> Worklist(MLI.begin(), MLI.end());
-    while (!Worklist.empty()) {
-      MachineLoop *Loop = Worklist.pop_back_val();
-      Worklist.append(Loop->begin(), Loop->end());
-      if (LoopFixer(MF, MLI, Loop).run()) {
-        return true;
-      }
-    }
-
-    return false;
-  }
-
 public:
   static char ID; // Pass identification, replacement for typeid
   WebAssemblyFixIrreducibleControlFlow() : MachineFunctionPass(ID) {}
 };
+
 } // end anonymous namespace
 
 char WebAssemblyFixIrreducibleControlFlow::ID = 0;
 INITIALIZE_PASS(WebAssemblyFixIrreducibleControlFlow, DEBUG_TYPE,
                 "Removes irreducible control flow", false, false)
 
 FunctionPass *llvm::createWebAssemblyFixIrreducibleControlFlow() {
   return new WebAssemblyFixIrreducibleControlFlow();
 }
 
 bool WebAssemblyFixIrreducibleControlFlow::runOnMachineFunction(
     MachineFunction &MF) {
   LLVM_DEBUG(dbgs() << "********** Fixing Irreducible Control Flow **********\n"
                        "********** Function: "
                     << MF.getName() << '\n');
-
-  bool Changed = false;
-  auto &MLI = getAnalysis<MachineLoopInfo>();
-
-  // When we modify something, bail out and recompute MLI, then start again, as
-  // we create a new natural loop when we resolve irreducible control flow, and
-  // other loops may become nested in it, etc. In practice this is not an issue
-  // because irreducible control flow is rare, only very few cycles are needed
-  // here.
-  while (LLVM_UNLIKELY(runIteration(MF, MLI))) {
-    // We rewrote part of the function; recompute MLI and start again.
-    LLVM_DEBUG(dbgs() << "Recomputing loops.\n");
+  if (LLVM_UNLIKELY(LoopFixer(MF).run())) {
+    // We rewrote part of the function; recompute relevant things.
     MF.getRegInfo().invalidateLiveness();
     MF.RenumberBlocks();
-    getAnalysis<MachineDominatorTree>().runOnMachineFunction(MF);
-    MLI.runOnMachineFunction(MF);
-    Changed = true;
+    return true;
   }
 
-  return Changed;
+  return false;
 }

test/CodeGen/WebAssembly/irreducible-cfg.ll

 
 for.inc205.i:                                     ; preds = %for.cond65.i
   br label %for.cond57.i
 
 return:                                           ; preds = %entry
   ret void
 }
 
+; Complex control flow without irreducibility. This requires we properly identify
+; the blocks entering each nested loop properly (in particular, even if they
+; are the entry to a parent loop).
+; CHECK-NOT: br_table
+define hidden void @ps_hints_apply() #0 {
+entry:
+  br label %psh
+
+psh:                                            ; preds = %entry
+  br i1 undef, label %for.cond, label %for.body
+
+for.body:                                       ; preds = %psh
+  br label %do.body
+
+do.body:                                        ; preds = %do.cond, %for.body
+  %cmp118 = icmp eq i32* undef, undef
+  br i1 %cmp118, label %Skip, label %do.cond
+
+do.cond:                                        ; preds = %do.body
+  br label %do.body
+
+for.cond:                                       ; preds = %Skip, %psh
+  br label %for.body39
+
+for.body39:                                     ; preds = %for.cond
+  br i1 undef, label %Skip, label %do.body45
+
+do.body45:                                      ; preds = %for.body39
+  unreachable
+
+Skip:                                           ; preds = %for.body39, %do.body
+  br label %for.cond
+}
+

test/CodeGen/WebAssembly/non-irreducible-cfg.ll

+; RUN: llc < %s -O0 -asm-verbose=false -verify-machineinstrs -disable-block-placement -wasm-disable-explicit-locals -wasm-keep-registers | FileCheck %s
+
+; Test irreducible CFG handling.
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown"
+
+; A simple sequence of loops with blocks in between, that should not be
+; misinterpreted as irreducible control flow.
+
+; CHECK-NOT: br_table
+define hidden i32 @_Z15fannkuch_workerPv(i8* %_arg) #0 {
+for.cond:                                         ; preds = %entry
+  br label %do.body
+
+do.body:                                          ; preds = %do.cond, %for.cond
+  br label %for.cond1
+
+for.cond1:                                        ; preds = %for.body, %do.body
+  br i1 1, label %for.cond1, label %for.end
+
+for.end:                                          ; preds = %for.cond1
+  br label %do.cond
+
+do.cond:                                          ; preds = %for.end
+  br i1 1, label %do.body, label %do.end
+
+do.end:                                           ; preds = %do.cond
+  br label %for.cond2
+
+for.cond2:                                        ; preds = %for.end6, %do.end
+  br label %for.cond3
+
+for.cond3:                                        ; preds = %for.body5, %for.cond2
+  br i1 1, label %for.cond3, label %for.end6
+
+for.end6:                                         ; preds = %for.cond3
+  br label %for.cond2
+
+return:                                           ; No predecessors!
+  ret i32 1
+}
+