[LoopFlatten] Updating Phi nodes after IV widening

[SjoerdMeijer]

Summary

In rG6a076fa9539e, a problem with updating the old/narrow phi nodes after IV widening was introduced. If after widening of the IV the transformation is *not* applied, the narrow phi node was incorrectly modified, which should only happen if flattening happens. This can be seen in the added test widen-iv2.ll, which incorrectly had 1 incoming value, but should have its original 2 incoming values, which is now restored.

[xbolva00]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

725

Deleted is now unused?

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

725

Thanks for taking a look at this.

Function CreateWideIV takes and sets a bool if the phi node was trivially deleted. Above, on line 720, we pass FI.IsTriviallyDeletedPhi. We want to record this state for the inner induction phi, because it's the inner loop that gets removed. But here, for the outer loop, we don't really care about this. So it's not really used, and we pass in a dummy Deleted boolean.

[SjoerdMeijer]

Just to be complete: this passes a stage2 build and tests (and benchmarks I ran).

[dmgreen]

Can we add the inner phi that was not deleted to InnerPHIsToTransform? That would allow us to hold less state, hopefully simplifying things a little.

[SjoerdMeijer]

Thanks, now (re)using InnerPHIsToTransform , which indeed simplify things a bit.

[dmgreen]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

104

Other places add to InnerInductionPHI. Is it always safe (now and in the future) to assume they are all narrow induction phis?

[SjoerdMeijer]

Reverting to previous revision: don't conflate narrow and phis that will be transformed in InnerInductionPHI.

[SjoerdMeijer]

@dmgreen : I propose to keep this revision. The extra bit of state for NarrowInnerInductionPHI and NarrowOuterInductionPHI is needed anyway, to checks on the phi nodes (and skip these ones). That's why e.g. just this is not enough:

diff --git a/llvm/lib/Transforms/Scalar/LoopFlatten.cpp b/llvm/lib/Transforms/Scalar/LoopFlatten.cpp
index 931bcc2aada6..f81eb07b5425 100644                                                             
--- a/llvm/lib/Transforms/Scalar/LoopFlatten.cpp                                                    
+++ b/llvm/lib/Transforms/Scalar/LoopFlatten.cpp                                                    
@@ -711,7 +711,7 @@ static bool CanWidenIV(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,        
   // If the inner Phi node cannot be trivially deleted, we need to at least                        
   // bring it in a consistent state.                                                               
   if (!Deleted)                                                                                    
-    FI.InnerInductionPHI->removeIncomingValue(FI.InnerLoop->getLoopLatch());                       
+    FI.InnerPHIsToTransform.insert(FI.InnerInductionPHI);                                          
   if (!CreateWideIV({FI.OuterInductionPHI, MaxLegalType, false }, Deleted))                        
     return false;

And I think it is clearer to keep these narrow phi separate from InnerPHIsToTransform. Let me know what you think.

[SjoerdMeijer]

Ok, think I got what you meant. This gets rid of the extra state IsTriviallyDeletedPhi and adds to InnerPHIsToTransform so that the phi will be adjusted (and we don't need to add extra code for that).

[dmgreen]

OK Thanks. LGTM

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

struct FlattenInfo {
   SmallPtrSet<PHINode *, 4> InnerPHIsToTransform;
 
   // Whether this holds the flatten info before or after widening.
   bool Widened = false;
 
   // Holds the old/narrow induction phis, i.e. the Phis before IV widening has
   // been applied. This bookkeeping is used so we can skip some checks on these
   // phi nodes.
-  SmallPtrSet<PHINode *, 2> OldInductionPHIs;
+  PHINode *NarrowInnerInductionPHI = nullptr;
+  PHINode *NarrowOuterInductionPHI = nullptr;
 
   FlattenInfo(Loop *OL, Loop *IL) : OuterLoop(OL), InnerLoop(IL) {};
+
+  bool isNarrowInductionPhi(PHINode *Phi) {
+    // This can't be the narrow phi if we haven't widened the IV first.
+    if (!Widened)
+      return false;
+    return NarrowInnerInductionPHI == Phi || NarrowOuterInductionPHI == Phi;
+  }
 };
 
 static bool
 setLoopComponents(Value *&TC, Value *&TripCount, BinaryOperator *&Increment,
                   SmallPtrSetImpl<Instruction *> &IterationInstructions) {
   TripCount = TC;
   IterationInstructions.insert(Increment);
   LLVM_DEBUG(dbgs() << "Found Increment: "; Increment->dump());
static bool checkPHIs(FlattenInfo &FI, const TargetTransformInfo *TTI) {
 
   // Check that all PHI nodes in the inner loop header match one of the valid
   // patterns.
   for (PHINode &InnerPHI : FI.InnerLoop->getHeader()->phis()) {
     // The induction PHIs break these rules, and that's OK because we treat
     // them specially when doing the transformation.
     if (&InnerPHI == FI.InnerInductionPHI)
       continue;
-    if (FI.Widened && FI.OldInductionPHIs.count(&InnerPHI))
+    if (FI.isNarrowInductionPhi(&InnerPHI))
       continue;
 
     // Each inner loop PHI node must have two incoming values/blocks - one
     // from the pre-header, and one from the latch.
     assert(InnerPHI.getNumIncomingValues() == 2);
     Value *PreHeaderValue =
         InnerPHI.getIncomingValueForBlock(FI.InnerLoop->getLoopPreheader());
     Value *LatchValue =
for (PHINode &InnerPHI : FI.InnerLoop->getHeader()->phis()) {
     LLVM_DEBUG(dbgs() << "PHI pair is safe:\n");
     LLVM_DEBUG(dbgs() << "  Inner: "; InnerPHI.dump());
     LLVM_DEBUG(dbgs() << "  Outer: "; OuterPHI->dump());
     SafeOuterPHIs.insert(OuterPHI);
     FI.InnerPHIsToTransform.insert(&InnerPHI);
   }
 
   for (PHINode &OuterPHI : FI.OuterLoop->getHeader()->phis()) {
-    if (FI.Widened && FI.OldInductionPHIs.count(&OuterPHI))
+    if (FI.isNarrowInductionPhi(&OuterPHI))
       continue;
     if (!SafeOuterPHIs.count(&OuterPHI)) {
       LLVM_DEBUG(dbgs() << "found unsafe PHI in outer loop: "; OuterPHI.dump());
       return false;
     }
   }
 
   LLVM_DEBUG(dbgs() << "checkPHIs: OK\n");
auto CreateWideIV = [&] (WideIVInfo WideIV, bool &Deleted) -> bool {
     LLVM_DEBUG(dbgs() << "Deleting old phi: "; WideIV.NarrowIV->dump());
     Deleted = RecursivelyDeleteDeadPHINode(WideIV.NarrowIV);
     return true;
   };
 
   bool Deleted;
   if (!CreateWideIV({FI.InnerInductionPHI, MaxLegalType, false }, Deleted))
     return false;
-  // If the inner Phi node cannot be trivially deleted, we need to at least
-  // bring it in a consistent state.
+  // Add the narrow phi to list, so that it will be adjusted later when the
+  // the transformation is performed.
   if (!Deleted)
-    FI.InnerInductionPHI->removeIncomingValue(FI.InnerLoop->getLoopLatch());
+    FI.InnerPHIsToTransform.insert(FI.InnerInductionPHI);
+
   if (!CreateWideIV({FI.OuterInductionPHI, MaxLegalType, false }, Deleted))
     return false;
 
   assert(Widened && "Widened IV expected");
   FI.Widened = true;
 
   // Save the old/narrow induction phis, which we need to ignore in CheckPHIs.
-  FI.OldInductionPHIs.insert(FI.InnerInductionPHI);
-  FI.OldInductionPHIs.insert(FI.OuterInductionPHI);
+  FI.NarrowInnerInductionPHI = FI.InnerInductionPHI;
+  FI.NarrowOuterInductionPHI = FI.OuterInductionPHI;
 
   // After widening, rediscover all the loop components.
   return CanFlattenLoopPair(FI, DT, LI, SE, AC, TTI);
 }
 
 static bool FlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                             ScalarEvolution *SE, AssumptionCache *AC,
                             const TargetTransformInfo *TTI) {

llvm/test/Transforms/LoopFlatten/widen-iv2.ll

 ; CHECK-NEXT:    [[I_016_US:%.*]] = phi i32 [ [[INC5_US:%.*]], [[FOR_COND1_FOR_INC4_CRIT_EDGE_US]] ], [ 0, [[FOR_COND1_PREHEADER_US_PREHEADER]] ]
 ; CHECK-NEXT:    [[TMP4:%.*]] = load i32*, i32** @c, align 8
 ; CHECK-NEXT:    [[TMP5:%.*]] = mul nsw i64 [[INDVAR2]], [[TMP2]]
 ; CHECK-NEXT:    [[MUL_US:%.*]] = mul nsw i32 [[I_016_US]], [[TMP0]]
 ; CHECK-NEXT:    [[TMP6:%.*]] = sext i32 [[MUL_US]] to i64
 ; CHECK-NEXT:    br label [[FOR_BODY3_US:%.*]]
 ; CHECK:       for.body3.us:
 ; CHECK-NEXT:    [[INDVAR:%.*]] = phi i64 [ [[INDVAR_NEXT:%.*]], [[FOR_BODY3_US]] ], [ 0, [[FOR_COND1_PREHEADER_US]] ]
-; CHECK-NEXT:    [[J_014_US:%.*]] = phi i32 [ 0, [[FOR_COND1_PREHEADER_US]] ]
+; CHECK-NEXT:    [[J_014_US:%.*]] = phi i32 [ 0, [[FOR_COND1_PREHEADER_US]] ], [ [[INC_US:%.*]], [[FOR_BODY3_US]] ]
 ; CHECK-NEXT:    [[TMP7:%.*]] = add nsw i64 [[INDVAR]], [[TMP5]]
 ; CHECK-NEXT:    [[TMP8:%.*]] = sext i32 [[J_014_US]] to i64
 ; CHECK-NEXT:    [[TMP9:%.*]] = add nsw i64 [[TMP8]], [[TMP5]]
 ; CHECK-NEXT:    [[ADD_US:%.*]] = add nsw i32 [[J_014_US]], [[MUL_US]]
 ; CHECK-NEXT:    [[IDXPROM_US:%.*]] = sext i32 [[ADD_US]] to i64
 ; CHECK-NEXT:    [[ARRAYIDX_US:%.*]] = getelementptr inbounds i32, i32* [[TMP4]], i64 [[TMP7]]
 ; CHECK-NEXT:    store i32 32, i32* [[ARRAYIDX_US]], align 4
 ; CHECK-NEXT:    [[INDVAR_NEXT]] = add i64 [[INDVAR]], 1
-; CHECK-NEXT:    [[INC_US:%.*]] = add nuw nsw i32 [[J_014_US]], 1
+; CHECK-NEXT:    [[INC_US]] = add nuw nsw i32 [[J_014_US]], 1
 ; CHECK-NEXT:    [[CMP2_US:%.*]] = icmp slt i64 [[INDVAR_NEXT]], [[TMP1]]
 ; CHECK-NEXT:    br i1 [[CMP2_US]], label [[FOR_BODY3_US]], label [[FOR_COND1_FOR_INC4_CRIT_EDGE_US]]
 ; CHECK:       for.cond1.for.inc4_crit_edge.us:
 ; CHECK-NEXT:    [[INDVAR_NEXT3]] = add i64 [[INDVAR2]], 1
 ; CHECK-NEXT:    [[INC5_US]] = add nuw nsw i32 [[I_016_US]], 1
 ; CHECK-NEXT:    [[CMP_US:%.*]] = icmp slt i64 [[INDVAR_NEXT3]], [[TMP3]]
 ; CHECK-NEXT:    br i1 [[CMP_US]], label [[FOR_COND1_PREHEADER_US]], label [[FOR_END6_LOOPEXIT:%.*]]
 ; CHECK:       for.end6.loopexit: