[FlattenCFG] Fixup Phi nodes during CFG flattening

llvm/lib/Transforms/Utils/FlattenCFG.cpp

@@ -134,10 +134,6 @@ class FlattenCFGOpt {
 ///  its predecessor.  In Case 2, BB (BB3) only has conditional branches
 ///  as its predecessors.
 bool FlattenCFGOpt::FlattenParallelAndOr(BasicBlock *BB, IRBuilder<> &Builder) {
-  PHINode *PHI = dyn_cast<PHINode>(BB->begin());
-  if (PHI)
-    return false; // For simplicity, avoid cases containing PHI nodes.
-
   BasicBlock *LastCondBlock = nullptr;
   BasicBlock *FirstCondBlock = nullptr;
   BasicBlock *UnCondBlock = nullptr;
@@ -208,8 +204,10 @@ bool FlattenCFGOpt::FlattenParallelAndOr(BasicBlock *BB, IRBuilder<> &Builder) {
 
     if (Idx == -1)
       Idx = CIdx;
-    else if (CIdx != Idx)
-      return false;
+    else if (CIdx != Idx) {
+      // Inverse Branch Condition
+      InvertBranch(PBI, Builder);
+    }
 
     // PS is the successor which is not BB. Check successors to identify
     // the last conditional branch.
@@ -269,11 +267,6 @@ bool FlattenCFGOpt::FlattenParallelAndOr(BasicBlock *BB, IRBuilder<> &Builder) {
   if (!PBI1 || !PBI1->isUnconditional())
     return false;
 
-  // PS2 should not contain PHI node.
-  PHI = dyn_cast<PHINode>(PS2->begin());
-  if (PHI)
-    return false;
-
   // Do the transformation.
   BasicBlock *CB;
   BranchInst *PBI = cast<BranchInst>(FirstCondBlock->getTerminator());
@@ -291,17 +284,45 @@ bool FlattenCFGOpt::FlattenParallelAndOr(BasicBlock *BB, IRBuilder<> &Builder) {
     // Merge conditions.
     Builder.SetInsertPoint(PBI);
     Value *NC;
-    if (Idx == 0)
-      // Case 2, use parallel or.
-      NC = Builder.CreateOr(PC, CC);
-    else
+    if (UnCondBlock)
       // Case 1, use parallel and.
       NC = Builder.CreateAnd(PC, CC);
+    else
+      // Case 2, use parallel or.
+      NC = Builder.CreateOr(PC, CC);
+
+    // Fixup PHI node if needed
+    for (BasicBlock *CBS : successors(PBI)) {
+      for (PHINode &Phi : CBS->phis()) {
+        Value *origPhi0 = nullptr;
+        Value *newPhi = nullptr;
+        if (llvm::is_contained(Phi.blocks(), FirstCondBlock)) {
+          origPhi0 = Phi.removeIncomingValue(FirstCondBlock, false);
+          newPhi = origPhi0;
+        }
+        if (llvm::is_contained(Phi.blocks(), CB)) {
+          Value *origPhi1 = Phi.removeIncomingValue(CB, false);
+          newPhi = origPhi1;
+
+          if (origPhi0) {
+            // Swap branch given the conditions
+            if (PBI->getSuccessor(0) == CBS) {
+              newPhi = Builder.CreateSelect(PC, origPhi0, origPhi1);
+            } else {
+              newPhi = Builder.CreateSelect(PC, origPhi1, origPhi0);
+            }
+          }
+        }
+        if (newPhi)
+          Phi.addIncoming(newPhi, FirstCondBlock);
+      }
+    }
 
     PBI->replaceUsesOfWith(CC, NC);
     PC = NC;
     if (CB == LastCondBlock)
       Iteration = false;
+
     // Remove internal conditional branches.
     CB->dropAllReferences();
     // make CB unreachable and let downstream to delete the block.

llvm/test/Transforms/SimplifyCFG/flatten-cfg-with-phi.ll

@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --tool ./opt --version 5
+; RUN: opt < %s -passes=flatten-cfg -S | FileCheck %s
+
+define i1 @_Z7compareRK1SS1_(ptr %a, ptr %b) {
+; CHECK-LABEL: @_Z7compareRK1SS1_(
+; CHECK-SAME: ptr [[A:%.*]], ptr [[B:%.*]]) {
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = load i32, ptr %a, align 4, !tbaa !3
+; CHECK-NEXT:   %1 = load i32, ptr %b, align 4, !tbaa !3
+; CHECK-NEXT:   %cmp.i = icmp sge i32 %0, %1
+; CHECK-NEXT:   %cmp.i19 = icmp eq i32 %0, %1
+; CHECK-NEXT:   %2 = and i1 %cmp.i, %cmp.i19
+; CHECK-NEXT:   %3 = select i1 %cmp.i, i1 false, i1 true
+; CHECK-NEXT:   br i1 %2, label %land.rhs, label %lor.end
+; CHECK-LABEL: lor.end:                                          ; preds = %land.rhs, %entry
+; CHECK-NEXT:   %6 = phi i1 [ %cmp, %land.rhs ], [ %3, %entry ]
+; CHECK-NEXT:   ret i1 %6
+entry:
+  %0 = load i32, ptr %a, align 4, !tbaa !3
+  %1 = load i32, ptr %b, align 4, !tbaa !3
+  %cmp.i = icmp slt i32 %0, %1
+  br i1 %cmp.i, label %lor.end, label %lor.rhs
+
+lor.rhs:                                          ; preds = %entry
+  %cmp.i19 = icmp eq i32 %0, %1
+  br i1 %cmp.i19, label %land.rhs, label %lor.end
+
+land.rhs:                                         ; preds = %lor.rhs
+  %y = getelementptr inbounds nuw i8, ptr %a, i64 4
+  %2 = load i32, ptr %y, align 4, !tbaa !8
+  %y14 = getelementptr inbounds nuw i8, ptr %b, i64 4
+  %3 = load i32, ptr %y14, align 4, !tbaa !8
+  %cmp = icmp slt i32 %2, %3
+  br label %lor.end
+
+lor.end:                                          ; preds = %lor.rhs, %land.rhs, %entry
+  %4 = phi i1 [ true, %entry ], [ false, %lor.rhs ], [ %cmp, %land.rhs ]
+  ret i1 %4
+}
+
+!llvm.module.flags = !{!0, !1}
+!llvm.ident = !{!2}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 7, !"uwtable", i32 2}
+!2 = !{!"clang"}
+!3 = !{!4, !5, i64 0}
+!4 = !{!"_ZTS1S", !5, i64 0, !5, i64 4}
+!5 = !{!"int", !6, i64 0}
+!6 = !{!"omnipotent char", !7, i64 0}
+!7 = !{!"Simple C++ TBAA"}
+!8 = !{!4, !5, i64 4}
+!9 = !{!5, !5, i64 0}