[SimpleLoopUnswitch] Don't non-trivially unswitch loops with catchswitch exits

[aeubanks]

Summary

SplitBlock() can't handle catchswitch.

Fixes PR50973.

[aheejin]

Thank you!

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2784

Oh, I didn't find there's already a similar condition for cleanuppad...

[majnemer]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2778–2784

Why is splitting a block with a cleanuppad different from a block with a landingpad?

[aheejin]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2778–2784

The splitting here is to split the BB into BB with the first instruction and BB the rest and clone the BB with the first instruction, which happens when the EH pad BB is among a loop's exit blocks. SimpleUnswitch clones the loop and all its exit blocks, but instead of cloning the full exit blocks, it splits those blocks into the first instruction and the rest and clone only the first instruction. So

Before:

bb0:
  first instruction
  second instruction
  ...
  br %succ

succ:
  ...

If bb0 is among the exit blocks of a loop that's being unswitched,

bb0:
  first instruction
  br %bb0.rest

bb0.split:
  first instruction
  br %bb0.rest

bb0.rest:
  second instruction
  ...
  br %succ

succ:
  ...

If bb0 contains catchswitch or cleanuppad, it is cloned because it is the first instruction, and its token return value should be merged with a phi in bb0.rest. But tokens can't be phi'd, right?

This is the reason I thought why they are different from landingpad, but I can be mistaken, so please let me know if so.

[majnemer]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2778–2784

OK, so it does not have to do with SplitBlock. In that case, I'd adjust the comment to make it more clear.

So if the issue is that it needs to do something about tokens which will now be live-out of the block, this seems unrelated to ehpads and more about token producing instructions/intrinsics. IIRC, I tried to handle this here: https://github.com/llvm/llvm-project/blob/873ff5a72864fdf60614cca8adbd0d869fc9a9a2/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp#L2819-L2820

Maybe I got it wrong?

[aheejin]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2778–2784

In the attached test here, catchswitch is not inside the loop; it is an exit block of the loop. The code snippet you showed seems to only count for instructions within a loop.

[aeubanks]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2778–2784

The added assert in SplitBlockImpl() is triggering on the block with the catchswitch. SplitBlockImpl() assumes that there is a non-PHI/non-EHPad instruction after split point instruction, which is not true for catchswitch

[rnk]

By the way, the catchswitch instructions true purpose is really just to multiplex invokes. When we added WinEH, callbr didn't exist. We felt that extending invoke with multiple unwind edges would be too disruptive. If we were doing it today, we would probably implement this in such a way that blocks without insertion points do not exist. That's a good project that would simplify LLVM transforms in the long run. It's probably not that straightforward, since most passes think callbr is just a wrapper for asm goto.

The cleanuppad design is probably as good as it's going to get. If we want to unswitch a loop inside a destructor cleanup, we wouldn't want that pass to accidentally create two prologues for a cleanup funclet, that would be too challenging. Instead, passes that really need to insert code along edges entering a cleanuppad can be taught to create new trivial cleanup funclets. This isn't too different from splitting edges coming into a landingpad, which LLVM knows how to do.

[aeubanks]

any remaining objections?

[aheejin]

Can we land this? @majnemer, Do you have any remaining concerns?

[aheejin]

@aeubanks, Can we land this? Thank you!

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

static bool unswitchBestCondition(
   LoopBlocksRPO RPOT(&L);
   RPOT.perform(&LI);
   if (containsIrreducibleCFG<const BasicBlock *>(RPOT, LI))
     return false;
 
   SmallVector<BasicBlock *, 4> ExitBlocks;
   L.getUniqueExitBlocks(ExitBlocks);
 
-  // We cannot unswitch if exit blocks contain a cleanuppad instruction as we
-  // don't know how to split those exit blocks.
+  // We cannot unswitch if exit blocks contain a cleanuppad/catchswitch
+  // instruction as we don't know how to split those exit blocks.
   // FIXME: We should teach SplitBlock to handle this and remove this
   // restriction.
   for (auto *ExitBB : ExitBlocks) {
-    if (isa<CleanupPadInst>(ExitBB->getFirstNonPHI())) {
-      LLVM_DEBUG(
-          dbgs() << "Cannot unswitch because of cleanuppad in exit block\n");
+    auto *I = ExitBB->getFirstNonPHI();
+    if (isa<CleanupPadInst>(I) || isa<CatchSwitchInst>(I)) {
+      LLVM_DEBUG(dbgs() << "Cannot unswitch because of cleanuppad/catchswitch "
+                           "in exit block\n");
       return false;
     }
   }
 
   LLVM_DEBUG(
       dbgs() << "Considering " << UnswitchCandidates.size()
              << " non-trivial loop invariant conditions for unswitching.\n");
 

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

static BasicBlock *SplitBlockImpl(BasicBlock *Old, Instruction *SplitPt,
                                   const Twine &BBName, bool Before) {
   if (Before) {
     DomTreeUpdater LocalDTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
     return splitBlockBefore(Old, SplitPt,
                             DTU ? DTU : (DT ? &LocalDTU : nullptr), LI, MSSAU,
                             BBName);
   }
   BasicBlock::iterator SplitIt = SplitPt->getIterator();
-  while (isa<PHINode>(SplitIt) || SplitIt->isEHPad())
+  while (isa<PHINode>(SplitIt) || SplitIt->isEHPad()) {
     ++SplitIt;
+    assert(SplitIt != SplitPt->getParent()->end());
+  }
   std::string Name = BBName.str();
   BasicBlock *New = Old->splitBasicBlock(
       SplitIt, Name.empty() ? Old->getName() + ".split" : Name);
 
   // The new block lives in whichever loop the old one did. This preserves
   // LCSSA as well, because we force the split point to be after any PHI nodes.
   if (LI)
     if (Loop *L = LI->getLoopFor(Old))

llvm/test/Transforms/SimpleLoopUnswitch/catchswitch.ll

+; RUN: opt -passes=simple-loop-unswitch -enable-nontrivial-unswitch < %s -S | FileCheck %s
+
+; CHECK: if.end{{.*}}:
+; CHECK-NOT: if.end{{.*}}:
+declare i32 @__gxx_wasm_personality_v0(...)
+
+declare void @foo()
+
+define void @test(i1 %arg) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+entry:
+  br label %while.body
+
+while.body:                                       ; preds = %cleanup, %entry
+  br i1 %arg, label %if.end, label %if.then
+
+if.then:                                          ; preds = %while.body
+  br label %if.end
+
+if.end:                                           ; preds = %if.then, %while.body
+  invoke void @foo()
+          to label %cleanup unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %invoke.cont, %if.end
+  %0 = catchswitch within none [label %catch] unwind to caller
+
+catch:                                            ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* null]
+  unreachable
+
+cleanup:                                          ; preds = %invoke.cont
+  br label %while.body
+}
+