[IROutliner] IR Outliner can mix up order of incoming values when compressing phi nodes if contain the same values

[AndrewLitteken]

Summary

Github issue: https://github.com/llvm/llvm-project/issues/54211

Example: https://godbolt.org/z/KrKPojnqE

When matching PHINodes when margining functions the IROutliner only checks that an incoming value exists in phi node in overall function. It doesn't check the length, the order, or that the incoming block also matches. In the given example, we see that both phi nodes have the same incoming values, but from different blocks.

The fix is to to enforce stricter a match of the incoming value, and the incoming block as well when matching the created phi nodes.

[paquette]

llvm/lib/Transforms/IPO/IROutliner.cpp

1599–1603

• Nit: I don't think we need an initial value for SmallVector anymore
• Could use a comment explaining what the variable is? (Probably the same as findCanonNumsForPHI, but future readers may not see that.)

1629–1664

Could use a comment too?

1633

Comment on what you're trying to do here?

1681

findCorrespondingBlockIn mentions the value could be nullptr. Assert?

llvm/test/Transforms/IROutliner/different-order-phi-merges.ll

5

[AndrewLitteken]

Updating with nits

[AndrewLitteken]

Updating with context

[paquette]

LGTM

llvm/include/llvm/Transforms/IPO/IROutliner.h

struct OutlinableRegion {
   /// Find a corresponding value for \p V in similar OutlinableRegion \p Other.
   ///
   /// \param Other [in] - The OutlinableRegion to find the corresponding Value
   /// in.
   /// \param V [in] - The Value to look for in the other region.
   /// \return The corresponding Value to \p V if it exists, otherwise nullptr.
   Value *findCorrespondingValueIn(const OutlinableRegion &Other, Value *V);
 
+  /// Find a corresponding BasicBlock for \p BB in similar OutlinableRegion \p Other.
+  ///
+  /// \param Other [in] - The OutlinableRegion to find the corresponding
+  /// BasicBlock in.
+  /// \param BB [in] - The BasicBlock to look for in the other region.
+  /// \return The corresponding Value to \p V if it exists, otherwise nullptr.
+  BasicBlock *findCorrespondingBlockIn(const OutlinableRegion &Other,
+                                       BasicBlock *BB);
+
   /// Get the size of the code removed from the region.
   ///
   /// \param [in] TTI - The TargetTransformInfo for the parent function.
   /// \returns the code size of the region
   InstructionCost getBenefit(TargetTransformInfo &TTI);
 };
 
 /// This class is a pass that identifies similarity in a Module, extracts

llvm/lib/Transforms/IPO/IROutliner.cpp

Value *OutlinableRegion::findCorrespondingValueIn(const OutlinableRegion &Other,
   Optional<unsigned> GVN = Candidate->getGVN(V);
   assert(GVN.hasValue() && "No GVN for incoming value");
   Optional<unsigned> CanonNum = Candidate->getCanonicalNum(*GVN);
   Optional<unsigned> FirstGVN = Other.Candidate->fromCanonicalNum(*CanonNum);
   Optional<Value *> FoundValueOpt = Other.Candidate->fromGVN(*FirstGVN);
   return FoundValueOpt.getValueOr(nullptr);
 }
 
+BasicBlock *
+OutlinableRegion::findCorrespondingBlockIn(const OutlinableRegion &Other,
+                                           BasicBlock *BB) {
+  Instruction *FirstNonPHI = BB->getFirstNonPHI();
+  assert(FirstNonPHI && "block is empty?");
+  Value *CorrespondingVal = findCorrespondingValueIn(Other, FirstNonPHI);
+  if (!CorrespondingVal)
+    return nullptr;
+  BasicBlock *CorrespondingBlock =
+      cast<Instruction>(CorrespondingVal)->getParent();
+  return CorrespondingBlock;
+}
+
 /// Rewrite the BranchInsts in the incoming blocks to \p PHIBlock that are found
 /// in \p Included to branch to BasicBlock \p Replace if they currently branch
 /// to the BasicBlock \p Find.  This is used to fix up the incoming basic blocks
 /// when PHINodes are included in outlined regions.
 ///
 /// \param PHIBlock - The BasicBlock containing the PHINodes that need to be
 /// checked.
 /// \param Find - The successor block to be replaced.

 
 /// Find the canonical numbering for the incoming Values into the PHINode \p PN.
 ///
 /// \param PN [in] - The PHINode that we are finding the canonical numbers for.
 /// \param Region [in] - The OutlinableRegion containing \p PN.
 /// \param OutputMappings [in] - The mapping of output values from outlined
 /// region to their original values.
 /// \param CanonNums [out] - The canonical numbering for the incoming values to
-/// \p PN.
+/// \p PN paired with their incoming block.
 /// \param ReplacedWithOutlinedCall - A flag to use the extracted function call
 /// of \p Region rather than the overall function's call.
-static void
-findCanonNumsForPHI(PHINode *PN, OutlinableRegion &Region,
-                    const DenseMap<Value *, Value *> &OutputMappings,
-                    DenseSet<unsigned> &CanonNums,
-                    bool ReplacedWithOutlinedCall = true) {
+static void findCanonNumsForPHI(
+    PHINode *PN, OutlinableRegion &Region,
+    SmallVector<std::pair<unsigned, BasicBlock *>> &CanonNums,
   // Iterate over the incoming values.
   for (unsigned Idx = 0, EIdx = PN->getNumIncomingValues(); Idx < EIdx; Idx++) {
     Value *IVal = PN->getIncomingValue(Idx);
+    BasicBlock *IBlock = PN->getIncomingBlock(Idx);
     // If we have an argument as incoming value, we need to grab the passed
     // value from the call itself.
     if (Argument *A = dyn_cast<Argument>(IVal)) {
       if (ReplacedWithOutlinedCall)
         IVal = getPassedArgumentInAlreadyOutlinedFunction(A, Region);
       else
         IVal = getPassedArgumentAndAdjustArgumentLocation(A, Region);
     }
 
     // Get the original value if it has been replaced by an output value.
     IVal = findOutputMapping(OutputMappings, IVal);
 
     // Find and add the canonical number for the incoming value.
     Optional<unsigned> GVN = Region.Candidate->getGVN(IVal);
     assert(GVN.hasValue() && "No GVN for incoming value");
     Optional<unsigned> CanonNum = Region.Candidate->getCanonicalNum(*GVN);
     assert(CanonNum.hasValue() && "No Canonical Number for GVN");
-    CanonNums.insert(*CanonNum);
+    CanonNums.push_back(std::make_pair(*CanonNum, IBlock));
   }
 }
 
 /// Find, or add PHINode \p PN to the combined PHINode Block \p OverallPHIBlock
 /// in order to condense the number of instructions added to the outlined
 /// function.
 ///
 /// \param PN [in] - The PHINode that we are finding the canonical numbers for.
 /// \param Region [in] - The OutlinableRegion containing \p PN. 
 /// \param OverallPhiBlock [in] - The overall PHIBlock we are trying to find
 /// \p PN in.
 /// \param OutputMappings [in] - The mapping of output values from outlined
 /// region to their original values.
 /// \param UsedPhis [in, out] - The PHINodes in the block that have already been
 /// matched.
 /// \return the newly found or created PHINode in \p OverallPhiBlock.
 static PHINode*
 findOrCreatePHIInBlock(PHINode &PN, OutlinableRegion &Region,
                        BasicBlock *OverallPhiBlock,
                        const DenseMap<Value *, Value *> &OutputMappings,
                        DenseSet<PHINode *> &UsedPHIs) {
   OutlinableGroup &Group = *Region.Parent;
   
-  DenseSet<unsigned> PNCanonNums;
+  
+  // A list of the canonical numbering assigned to each incoming value, paired
+  // with the incoming block for the PHINode passed into this function.
+  SmallVector<std::pair<unsigned, BasicBlock *>> PNCanonNums;
+
   // We have to use the extracted function since we have merged this region into
   // the overall function yet.  We make sure to reassign the argument numbering
   // since it is possible that the argument ordering is different between the
   // functions.
   findCanonNumsForPHI(&PN, Region, OutputMappings, PNCanonNums,
                       /* ReplacedWithOutlinedCall = */ false);
 
   OutlinableRegion *FirstRegion = Group.Regions[0];
-  DenseSet<unsigned> CurrentCanonNums;
+
+  // A list of the canonical numbering assigned to each incoming value, paired
+  // with the incoming block for the PHINode that we are currently comparing
+  // the passed PHINode to.
+  SmallVector<std::pair<unsigned, BasicBlock *>> CurrentCanonNums;
+
   // Find the Canonical Numbering for each PHINode, if it matches, we replace
   // the uses of the PHINode we are searching for, with the found PHINode.
   for (PHINode &CurrPN : OverallPhiBlock->phis()) {
     // If this PHINode has already been matched to another PHINode to be merged,
     // we skip it.
     if (UsedPHIs.find(&CurrPN) != UsedPHIs.end())
       continue;
 
     CurrentCanonNums.clear();
     findCanonNumsForPHI(&CurrPN, *FirstRegion, OutputMappings, CurrentCanonNums,
                         /* ReplacedWithOutlinedCall = */ true);
-    if (all_of(PNCanonNums, [&CurrentCanonNums](unsigned CanonNum) {
-          return CurrentCanonNums.contains(CanonNum);
-        })) {
+
+    // If the list of incoming values is not the same length, then they cannot
+    // match since there is not an analogue for each incoming value.
+    if (PNCanonNums.size() != CurrentCanonNums.size())
+      continue;
+
+    bool FoundMatch = true;
+
+    // We compare the canonical value for each incoming value in the passed
+    // in PHINode to one already present in the outlined region.  If the
+    // incoming values do not match, then the PHINodes do not match.
+
+    // We also check to make sure that the incoming block matches as well by
+    // finding the corresponding incoming block in the combined outlined region
+    // for the current outlined region.
+    for (unsigned Idx = 0, Edx = PNCanonNums.size(); Idx < Edx; ++Idx) {
+      std::pair<unsigned, BasicBlock *> ToCompareTo = CurrentCanonNums[Idx];
+      std::pair<unsigned, BasicBlock *> ToAdd = PNCanonNums[Idx];
+      if (ToCompareTo.first != ToAdd.first) {
+        FoundMatch = false;
+        break;
+      }
+
+      BasicBlock *CorrespondingBlock =
+          Region.findCorrespondingBlockIn(*FirstRegion, ToAdd.second);
+      assert(CorrespondingBlock && "Found block is nullptr");
+      if (CorrespondingBlock != ToCompareTo.second) {
+        FoundMatch = false;
+        break;
+      }
+    }
+
+    // If all incoming values and branches matched, then we can merge
+    // into the found PHINode.
+    if (FoundMatch) {
       UsedPHIs.insert(&CurrPN);
       return &CurrPN;
     }
   }
 
   // If we've made it here, it means we weren't able to replace the PHINode, so
   // we must insert it ourselves.
   PHINode *NewPN = cast<PHINode>(PN.clone());
   NewPN->insertBefore(&*OverallPhiBlock->begin());
   for (unsigned Idx = 0, Edx = NewPN->getNumIncomingValues(); Idx < Edx;
        Idx++) {
     Value *IncomingVal = NewPN->getIncomingValue(Idx);
     BasicBlock *IncomingBlock = NewPN->getIncomingBlock(Idx);
 
     // Find corresponding basic block in the overall function for the incoming
     // block.
-    Instruction *FirstNonPHI = IncomingBlock->getFirstNonPHI();
-    assert(FirstNonPHI && "Incoming block is empty?");
-    Value *CorrespondingVal =
-        Region.findCorrespondingValueIn(*FirstRegion, FirstNonPHI);
-    assert(CorrespondingVal && "Value is nullptr?");
-    BasicBlock *BlockToUse = cast<Instruction>(CorrespondingVal)->getParent();
+    BasicBlock *BlockToUse =
+        Region.findCorrespondingBlockIn(*FirstRegion, IncomingBlock);
     NewPN->setIncomingBlock(Idx, BlockToUse);
 
     // If we have an argument we make sure we replace using the argument from
     // the correct function.
     if (Argument *A = dyn_cast<Argument>(IncomingVal)) {
       Value *Val = Group.OutlinedFunction->getArg(A->getArgNo());
       NewPN->setIncomingValue(Idx, Val);
       continue;

llvm/test/Transforms/IROutliner/different-order-phi-merges.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
+
+; Check that differently ordered phi nodes are not matched when merged, instead
+; generating two output paths.
+
+define void @f1() {
+bb1:
+  %0 = add i32 1, 2
+  %1 = add i32 3, 4
+  %2 = add i32 5, 6
+  %3 = add i32 7, 8
+  br i1 true, label %bb2, label %bb5
+bb2:
+  %4 = mul i32 5, 4
+  br label %bb5
+
+placeholder:
+  %a = sub i32 5, 4
+  ret void
+
+bb5:
+  %phinode = phi i32 [%3, %bb1], [%2, %bb2]
+  ret void
+}
+
+define void @f2() {
+bb1:
+  %0 = add i32 1, 2
+  %1 = add i32 3, 4
+  %2 = add i32 5, 6
+  %3 = add i32 7, 8
+  br i1 true, label %bb2, label %bb5
+bb2:
+  %4 = mul i32 5, 4
+  br label %bb5
+
+placeholder:
+  %a = sub i32 5, 4
+  ret void
+
+bb5:
+  %phinode = phi i32 [%2, %bb1], [%3, %bb2]
+  ret void
+}
+; CHECK-LABEL: @f1(
+; CHECK-NEXT:  bb1:
+; CHECK-NEXT:    [[PHINODE_CE_LOC:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[LT_CAST:%.*]] = bitcast i32* [[PHINODE_CE_LOC]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    [[TMP0:%.*]] = call i1 @outlined_ir_func_0(i32* [[PHINODE_CE_LOC]], i32 0)
+; CHECK-NEXT:    [[PHINODE_CE_RELOAD:%.*]] = load i32, i32* [[PHINODE_CE_LOC]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    br i1 [[TMP0]], label [[BB5:%.*]], label [[BB1_AFTER_OUTLINE:%.*]]
+; CHECK:       bb1_after_outline:
+; CHECK-NEXT:    ret void
+; CHECK:       bb5:
+; CHECK-NEXT:    [[PHINODE:%.*]] = phi i32 [ [[PHINODE_CE_RELOAD]], [[BB1:%.*]] ]
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: @f2(
+; CHECK-NEXT:  bb1:
+; CHECK-NEXT:    [[PHINODE_CE_LOC:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[LT_CAST:%.*]] = bitcast i32* [[PHINODE_CE_LOC]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    [[TMP0:%.*]] = call i1 @outlined_ir_func_0(i32* [[PHINODE_CE_LOC]], i32 1)
+; CHECK-NEXT:    [[PHINODE_CE_RELOAD:%.*]] = load i32, i32* [[PHINODE_CE_LOC]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 -1, i8* [[LT_CAST]])
+; CHECK-NEXT:    br i1 [[TMP0]], label [[BB5:%.*]], label [[BB1_AFTER_OUTLINE:%.*]]
+; CHECK:       bb1_after_outline:
+; CHECK-NEXT:    ret void
+; CHECK:       bb5:
+; CHECK-NEXT:    [[PHINODE:%.*]] = phi i32 [ [[PHINODE_CE_RELOAD]], [[BB1:%.*]] ]
+; CHECK-NEXT:    ret void
+;
+;
+; CHECK-LABEL: define internal i1 @outlined_ir_func_0(
+; CHECK-NEXT:  newFuncRoot:
+; CHECK-NEXT:    br label [[BB1_TO_OUTLINE:%.*]]
+; CHECK:       bb1_to_outline:
+; CHECK-NEXT:    [[TMP2:%.*]] = add i32 1, 2
+; CHECK-NEXT:    [[TMP3:%.*]] = add i32 3, 4
+; CHECK-NEXT:    [[TMP4:%.*]] = add i32 5, 6
+; CHECK-NEXT:    [[TMP5:%.*]] = add i32 7, 8
+; CHECK-NEXT:    br i1 true, label [[BB2:%.*]], label [[BB5_SPLIT:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[TMP6:%.*]] = mul i32 5, 4
+; CHECK-NEXT:    br label [[BB5_SPLIT]]
+; CHECK:       placeholder:
+; CHECK-NEXT:    [[A:%.*]] = sub i32 5, 4
+; CHECK-NEXT:    br label [[BB1_AFTER_OUTLINE_EXITSTUB:%.*]]
+; CHECK:       bb5.split:
+; CHECK-NEXT:    [[TMP7:%.*]] = phi i32 [ [[TMP4]], [[BB1_TO_OUTLINE]] ], [ [[TMP5]], [[BB2]] ]
+; CHECK-NEXT:    [[PHINODE_CE:%.*]] = phi i32 [ [[TMP5]], [[BB1_TO_OUTLINE]] ], [ [[TMP4]], [[BB2]] ]
+; CHECK-NEXT:    br label [[BB5_EXITSTUB:%.*]]
+; CHECK:       bb5.exitStub:
+; CHECK-NEXT:    switch i32 [[TMP1:%.*]], label [[FINAL_BLOCK_1:%.*]] [
+; CHECK-NEXT:    i32 0, label [[OUTPUT_BLOCK_0_1:%.*]]
+; CHECK-NEXT:    i32 1, label [[OUTPUT_BLOCK_1_1:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       bb1_after_outline.exitStub:
+; CHECK-NEXT:    switch i32 [[TMP1]], label [[FINAL_BLOCK_0:%.*]] [
+; CHECK-NEXT:    ]
+; CHECK:       output_block_0_1:
+; CHECK-NEXT:    store i32 [[PHINODE_CE]], i32* [[TMP0:%.*]], align 4
+; CHECK-NEXT:    br label [[FINAL_BLOCK_1]]
+; CHECK:       output_block_1_1:
+; CHECK-NEXT:    store i32 [[TMP7]], i32* [[TMP0]], align 4
+; CHECK-NEXT:    br label [[FINAL_BLOCK_1]]
+; CHECK:       final_block_0:
+; CHECK-NEXT:    ret i1 false
+; CHECK:       final_block_1:
+; CHECK-NEXT:    ret i1 true
+;