[SimplifyCFG] FoldBranchToCommonDest(): bonus instrns must only be us…

…ed by PHI nodes in successors (PR48450)

In particular, if the successor block, which is about to get a new
predecessor block, currently only has a single predecessor,
then the bonus instructions will be directly used within said successor,
which is fine, since the block with bonus instructions dominates that
successor. But once there's a new predecessor, the IR is no longer valid,
and we don't fix it, because we only update PHI nodes.

Which means, the live-out bonus instructions must be exclusively used
by the PHI nodes in successor blocks. So we have to form trivial PHI nodes.
which will then be successfully updated to recieve cloned bonus instns.

This all works fine, except for the fact that we don't have access to
the dominator tree, and we don't ignore unreachable code,
so we sometimes do end up having to deal with some weird IR.

Fixes https://bugs.llvm.org/show_bug.cgi?id=48450

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -68,6 +68,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/SSAUpdater.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <algorithm>
 #include <cassert>
@@ -2709,6 +2710,66 @@ static bool extractPredSuccWeights(BranchInst *PBI, BranchInst *BI,
   }
 }
 
+/// Given \p BB block, for each instruction in said block, insert trivial
+/// (single value) PHI nodes into each successor block of \p BB block, and
+/// rewrite all the the non-PHI (or PHI uses not in successors of \p BB block)
+/// uses of instructions of \p BB block to use newly-inserted PHI nodes.
+/// NOTE: even though it would be correct to not deal with multi-predecessor
+///       successor blocks, or uses within the \p BB block, we may be dealing
+///       with an unreachable IR, where many invariants don't hold...
+static void FormTrivialSSAPHI(BasicBlock *BB) {
+  SmallSetVector<BasicBlock *, 16> Successors(succ_begin(BB), succ_end(BB));
+
+  // Process instructions in reverse order. There is no correctness reason for
+  // that order, but it allows us to consistently insert new PHI nodes
+  // at the top of blocks, while maintaining their relative order.
+  for (Instruction &DefInstr : make_range(BB->rbegin(), BB->rend())) {
+    SmallVector<std::reference_wrapper<Use>, 16> UsesToRewrite;
+
+    // Cache which uses we'll want to rewrite.
+    copy_if(DefInstr.uses(), std::back_inserter(UsesToRewrite),
+            [BB, &DefInstr, &Successors](Use &U) {
+              auto *User = cast<Instruction>(U.getUser());
+              auto *UserBB = User->getParent();
+              // Generally, ignore users in the same block as the instruction
+              // itself, unless the use[r] either comes before, or is [by] the
+              // instruction itself, which means we are in an unreachable IR.
+              if (UserBB == BB)
+                return !DefInstr.comesBefore(User);
+              // Otherwise, rewrite all non-PHI users,
+              // or PHI users in non-successor blocks.
+              return !isa<PHINode>(User) || !Successors.contains(UserBB);
+            });
+
+    // So, do we have uses to rewrite?
+    if (UsesToRewrite.empty())
+      continue; // Check next remaining instruction.
+
+    SSAUpdater SSAUpdate;
+    SSAUpdate.Initialize(DefInstr.getType(), DefInstr.getName());
+
+    // Create a new PHI node in each successor block.
+    // WARNING: the iteration order is externally-observable,
+    //          and therefore must be stable!
+    for (BasicBlock *Successor : Successors) {
+      IRBuilder<> Builder(&Successor->front());
+      auto *PN = Builder.CreatePHI(DefInstr.getType(), pred_size(Successor),
+                                   DefInstr.getName());
+      // By default, have an 'undef' incoming value for each predecessor.
+      for (BasicBlock *PredsOfSucc : predecessors(Successor))
+        PN->addIncoming(UndefValue::get(DefInstr.getType()), PredsOfSucc);
+      // .. but receive the correct value when coming from the right block.
+      PN->setIncomingValueForBlock(BB, &DefInstr);
+      // And make note of that PHI.
+      SSAUpdate.AddAvailableValue(Successor, PN);
+    }
+
+    // And finally, rewrite all the problematic uses to use the new PHI nodes.
+    while (!UsesToRewrite.empty())
+      SSAUpdate.RewriteUseAfterInsertions(UsesToRewrite.pop_back_val());
+  }
+}
+
 /// If this basic block is simple enough, and if a predecessor branches to us
 /// and one of our successors, fold the block into the predecessor and use
 /// logical operations to pick the right destination.
@@ -2878,6 +2939,11 @@ bool llvm::FoldBranchToCommonDest(BranchInst *BI, MemorySSAUpdater *MSSAU,
       PBI->swapSuccessors();
     }
 
+    // Ensure that the bonus instructions are *only* used by the PHI nodes,
+    // because the successor basic block is about to get a new predecessor
+    // and non-PHI uses will become invalid.
+    FormTrivialSSAPHI(BB);
+
     // Before cloning instructions, notify the successor basic block that it
     // is about to have a new predecessor. This will update PHI nodes,
     // which will allow us to update live-out uses of bonus instructions.
@@ -2920,14 +2986,11 @@ bool llvm::FoldBranchToCommonDest(BranchInst *BI, MemorySSAUpdater *MSSAU,
       BonusInst.setName(BonusInst.getName() + ".old");
       BonusInst.replaceUsesWithIf(NewBonusInst, [BB](Use &U) {
         auto *User = cast<Instruction>(U.getUser());
-        // Ignore uses in the same block as the bonus instruction itself.
+        // Ignore the original bonus instructions themselves.
         if (User->getParent() == BB)
           return false;
-        // We can safely update external non-PHI uses.
-        if (!isa<PHINode>(User))
-          return true;
-        // For PHI nodes, don't touch incoming values for same block
-        // as the bonus instruction itself.
+        // Otherwise, we've got a PHI node. Don't touch incoming values
+        // for same block as the bonus instruction itself.
         return cast<PHINode>(User)->getIncomingBlock(U) != BB;
       });
     }

llvm/test/Transforms/SimplifyCFG/fold-branch-to-common-dest.ll

@@ -5,6 +5,7 @@ declare void @sideeffect0()
 declare void @sideeffect1()
 declare void @sideeffect2()
 declare void @use8(i8)
+declare i1 @gen1()
 
 ; Basic cases, blocks have nothing other than the comparison itself.
 
@@ -726,3 +727,55 @@ final_right:
   call void @sideeffect1()
   ret void
 }
+
+define void @pr48450() {
+; CHECK-LABEL: @pr48450(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[COUNTDOWN:%.*]] = phi i16 [ 128, [[ENTRY:%.*]] ], [ [[DEC2:%.*]], [[FOR_BODYTHREAD_PRE_SPLIT:%.*]] ]
+; CHECK-NEXT:    [[C:%.*]] = call i1 @gen1()
+; CHECK-NEXT:    br i1 [[C]], label [[FOR_INC:%.*]], label [[IF_THEN:%.*]]
+; CHECK:       for.inc:
+; CHECK-NEXT:    [[DOTOLD:%.*]] = add i16 [[COUNTDOWN]], -1
+; CHECK-NEXT:    [[DOTOLD3:%.*]] = icmp eq i16 [[COUNTDOWN]], 0
+; CHECK-NEXT:    br i1 [[DOTOLD3]], label [[IF_END_LOOPEXIT:%.*]], label [[FOR_BODYTHREAD_PRE_SPLIT]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[C2:%.*]] = call i1 @gen1()
+; CHECK-NEXT:    [[C2_NOT:%.*]] = xor i1 [[C2]], true
+; CHECK-NEXT:    [[DEC:%.*]] = add i16 [[COUNTDOWN]], -1
+; CHECK-NEXT:    [[CMP_NOT:%.*]] = icmp eq i16 [[COUNTDOWN]], 0
+; CHECK-NEXT:    [[OR_COND:%.*]] = or i1 [[C2_NOT]], [[CMP_NOT]]
+; CHECK-NEXT:    br i1 [[OR_COND]], label [[IF_END_LOOPEXIT]], label [[FOR_BODYTHREAD_PRE_SPLIT]]
+; CHECK:       for.bodythread-pre-split:
+; CHECK-NEXT:    [[DEC2]] = phi i16 [ [[DOTOLD]], [[FOR_INC]] ], [ [[DEC]], [[IF_THEN]] ]
+; CHECK-NEXT:    call void @sideeffect0()
+; CHECK-NEXT:    br label [[FOR_BODY]]
+; CHECK:       if.end.loopexit:
+; CHECK-NEXT:    [[DEC1:%.*]] = phi i16 [ undef, [[IF_THEN]] ], [ [[DOTOLD]], [[FOR_INC]] ]
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %for.body
+
+for.body:
+  %countdown = phi i16 [ 128, %entry ], [ %dec, %for.bodythread-pre-split ]
+  %c = call i1 @gen1()
+  br i1 %c, label %for.inc, label %if.then
+
+for.inc:
+  %dec = add i16 %countdown, -1
+  %cmp.not = icmp eq i16 %countdown, 0
+  br i1 %cmp.not, label %if.end.loopexit, label %for.bodythread-pre-split
+
+if.then:
+  %c2 = call i1 @gen1()
+  br i1 %c2, label %for.inc, label %if.end.loopexit
+
+for.bodythread-pre-split:
+  call void @sideeffect0()
+  br label %for.body
+
+if.end.loopexit:
+  ret void
+}