[IROutliner] Ignore Regions where part of an outlined phi nodes incoming block is included, but the final branch instruction is not

[AndrewLitteken]

Summary

Reproduced: https://godbolt.org/z/roznP99P3

When outlining a phi node, if the the incoming branch is a block contained in the region and the branch from that block is not outlined, we create broken code. The fix is to recognize when that branch from the included incoming block is not contained, and ignore the region.

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

280

since we are unable to handle the severing of the phi node right now.

Should this be a TODO?

309

Nit: Pull PN->getIncomingBlock(i) into a variable?

310

Isn't this always true since EndBB = BackInst->getParent()? Could this check be removed?

311

I think this could be hoisted out of the loops and into a variable?

bool EndBBTermAndBackInstDifferent = EndBB->getTerminator() != BackInst;
while (PhiNode ...)

317–318

Nit: This can all be one if now:

if (isa<PHINode>(BackInst) && BackInst != &*std::prev(EndBB->getFirstInsertionPt()))
   return;

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

300

Pull this above the if (!BBSet.contains(PN->getIncomingBlock(i)))?

[paquette]

Also, which GitHub issue is this associated with?

[AndrewLitteken]

https://github.com/llvm/llvm-project/issues/54212

[paquette]

I think after the one nit is fixed, this is good to go.

llvm/lib/Transforms/IPO/IROutliner.cpp

void OutlinableRegion::splitCandidate() {
   StartBB = StartInst->getParent();
   PrevBB = StartBB;
 
   DenseSet<BasicBlock *> BBSet;
   Candidate->getBasicBlocks(BBSet);
 
   // We iterate over the instructions in the region, if we find a PHINode, we
   // check if there are predecessors outside of the region, if there are,
   // we ignore this region since we are unable to handle the severing of the
-  // phi node right now. 
+  // phi node right now.
+
+  // TODO: Handle extraneous inputs for PHINodes through variable number of
+  // inputs, similar to how outputs are handled.
   BasicBlock::iterator It = StartInst->getIterator();
+  EndBB = BackInst->getParent();
+  BasicBlock *IBlock;
+  bool EndBBTermAndBackInstDifferent = EndBB->getTerminator() != BackInst;
   while (PHINode *PN = dyn_cast<PHINode>(&*It)) {
     unsigned NumPredsOutsideRegion = 0;
-    for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
-      if (!BBSet.contains(PN->getIncomingBlock(i)))
+    for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i) {
+      if (!BBSet.contains(PN->getIncomingBlock(i))) {
         ++NumPredsOutsideRegion;
+        continue;
+      }
+
+      // We must consider the case there the incoming block to the PHINode is
+      // the same as the final block of the OutlinableRegion.  If this is the
+      // case, the branch from this block must also be outlined to be valid.
+      IBlock = PN->getIncomingBlock(i);
+      if (IBlock == EndBB && EndBBTermAndBackInstDifferent)
+        ++NumPredsOutsideRegion;
+    }
 
     if (NumPredsOutsideRegion > 1)
       return;
     
     It++;
   }
 
   // If the region starts with a PHINode, but is not the initial instruction of
   // the BasicBlock, we ignore this region for now.
   if (isa<PHINode>(StartInst) && StartInst != &*StartBB->begin())
     return;
   
   // If the region ends with a PHINode, but does not contain all of the phi node
   // instructions of the region, we ignore it for now.
-  if (isa<PHINode>(BackInst)) {
-    EndBB = BackInst->getParent();
-    if (BackInst != &*std::prev(EndBB->getFirstInsertionPt()))
-      return;
-  }
+  if (isa<PHINode>(BackInst) &&
+      BackInst != &*std::prev(EndBB->getFirstInsertionPt()))
 
   // The basic block gets split like so:
   // block:                 block:
   //   inst1                  inst1
   //   inst2                  inst2
   //   region1               br block_to_outline
   //   region2              block_to_outline:
   //   region3          ->    region1

llvm/test/Transforms/IROutliner/phi-nodes-parent-block-referential.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --include-generated-funcs
+; RUN: opt -S -verify -iroutliner -ir-outlining-no-cost < %s | FileCheck %s
+
+; Show that we do not outline when all of the phi nodes in the beginning
+; block are included not in the region.
+
+define void @function1(i32* %a, i32* %b) {
+entry:
+  %0 = alloca i32, align 4
+  %c = load i32, i32* %0, align 4
+  %y = add i32 %c, %c
+  br label %test1
+dummy:
+  ret void
+test1:
+  %1 = phi i32 [ %e, %test1 ], [ %y, %entry ]
+  %2 = phi i32 [ %e, %test1 ], [ %y, %entry  ]
+  %e = load i32, i32* %0, align 4
+  %3 = add i32 %c, %c
+  %4 = sub i32 %c, %c
+  br i1 true, label %first, label %test1
+first:
+  ret void
+}
+
+define void @function2(i32* %a, i32* %b) {
+entry:
+  %0 = alloca i32, align 4
+  %c = load i32, i32* %0, align 4
+  %y = mul i32 %c, %c
+  br label %test1
+dummy:
+  ret void
+test1:
+  %1 = phi i32 [ %e, %test1 ], [ %y, %entry ]
+  %2 = phi i32 [ %y, %entry ], [ %e, %test1 ]
+  %e = load i32, i32* %0, align 4
+  %3 = add i32 %c, %c
+  %4 = mul i32 %c, %c
+  br i1 true, label %first, label %test1
+first:
+  ret void
+}
+; CHECK-LABEL: @function1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[E_LOC:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[C:%.*]] = load i32, i32* [[TMP0]], align 4
+; CHECK-NEXT:    [[Y:%.*]] = add i32 [[C]], [[C]]
+; CHECK-NEXT:    br label [[TEST1:%.*]]
+; CHECK:       dummy:
+; CHECK-NEXT:    ret void
+; CHECK:       test1:
+; CHECK-NEXT:    [[TMP1:%.*]] = phi i32 [ [[E_RELOAD:%.*]], [[TEST1]] ], [ [[Y]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[TMP2:%.*]] = phi i32 [ [[E_RELOAD]], [[TEST1]] ], [ [[Y]], [[ENTRY]] ]
+; CHECK-NEXT:    [[LT_CAST:%.*]] = bitcast i32* [[E_LOC]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32* [[TMP0]], i32 [[C]], i32* [[E_LOC]])
+; CHECK-NEXT:    [[E_RELOAD]] = load i32, i32* [[E_LOC]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    [[TMP3:%.*]] = sub i32 [[C]], [[C]]
+; CHECK-NEXT:    br i1 true, label [[FIRST:%.*]], label [[TEST1]]
+; CHECK:       first:
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: @function2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[E_LOC:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[C:%.*]] = load i32, i32* [[TMP0]], align 4
+; CHECK-NEXT:    [[Y:%.*]] = mul i32 [[C]], [[C]]
+; CHECK-NEXT:    br label [[TEST1:%.*]]
+; CHECK:       dummy:
+; CHECK-NEXT:    ret void
+; CHECK:       test1:
+; CHECK-NEXT:    [[TMP1:%.*]] = phi i32 [ [[E_RELOAD:%.*]], [[TEST1]] ], [ [[Y]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[TMP2:%.*]] = phi i32 [ [[Y]], [[ENTRY]] ], [ [[E_RELOAD]], [[TEST1]] ]
+; CHECK-NEXT:    [[LT_CAST:%.*]] = bitcast i32* [[E_LOC]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32* [[TMP0]], i32 [[C]], i32* [[E_LOC]])
+; CHECK-NEXT:    [[E_RELOAD]] = load i32, i32* [[E_LOC]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    [[TMP3:%.*]] = mul i32 [[C]], [[C]]
+; CHECK-NEXT:    br i1 true, label [[FIRST:%.*]], label [[TEST1]]
+; CHECK:       first:
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: define internal void @outlined_ir_func_0(
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label [[TEST1_TO_OUTLINE:%.*]]
+; CHECK:       test1_to_outline:
+; CHECK-NEXT:    [[E:%.*]] = load i32, i32* [[TMP0:%.*]], align 4
+; CHECK-NEXT:    [[TMP3:%.*]] = add i32 [[TMP1:%.*]], [[TMP1]]
+; CHECK-NEXT:    br label [[TEST1_AFTER_OUTLINE_EXITSTUB:%.*]]
+; CHECK:       test1_after_outline.exitStub:
+; CHECK-NEXT:    store i32 [[E]], i32* [[TMP2:%.*]], align 4
+; CHECK-NEXT:    ret void
+;