[LoopFlatten] Make the analysis more robust after IV widening

[SjoerdMeijer]

Summary

LoopFlatten wasn't triggering on this motivating case after IV widening:

void foo(int *A, int N, int M) {
  for (int i = 0; i < N; ++i)
    for (int j = 0; j < M; ++j)
      f(A[i*M+j]);
}

The reason was that the old induction phi nodes were getting in the way. These narrow and dead induction phis are not always trivially dead, and having both the narrow and wide IVs confused the analysis and caused it to bail. This adds some extra bookkeeping for these old phis, so we can filter them out when checks on phi nodes are performed. Other clean up passes will get rid of these old phis and increment instructions.

As this was one of the motivating examples from the beginning, it was surprising this wasn't triggering from C/C++ code. It looks like the IR and CFG is just slightly different.

[RosieSumpter]

LGTM, but of course @dmgreen should be the one to approve.

[dmgreen]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

443

Do we know this is always valid? For any sext? Or do we need to check the types match the new IV?

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

443

Yes, I believe this is always valid. This is after widening, when all loop components and checks have been performed (twice actually, one on the narrow IV, and again on the wide IV). Thus, when there is an extend in the way at this point, I believe it is always in the shape that we expect it. But I will add an assert for the matching types, and check that MatchedItCount isn't null.

[SjoerdMeijer]

Added an assert.

[dmgreen]

OK. LGTM then

[uabelho]

Hi,

I've hunted down a miscompile that seems to start happening with this commit.
Reproduce with

opt -passes='function(loop-flatten)' -S -o - bbi-60764.ll

The input has two nested loops, the outer for i [0, 3} and the inner for j [0, 15] and loop-flatten flattens it to one loop iterating for i' in [0, 63] which looks ok.

In the inner loop there is a load from v[16*i+j] but after flattening as far as I can see it loads from v[16*i'] so it quickly goes out-of-bounds of v and read random data instead.

So in the result there is

entry:
  %flatten.tripcount = mul i32 16, 4
  br label %for.cond1.preheader

for.cond1.preheader:                              ; preds = %for.cond.cleanup3, %entry
  %indvar2 = phi i32 [ %indvar.next3, %for.cond.cleanup3 ], [ 0, %entry ]
  %i.013 = phi i16 [ 0, %entry ], [ %inc7, %for.cond.cleanup3 ]
  %sum.012 = phi i16 [ 0, %entry ], [ %add5.lcssa, %for.cond.cleanup3 ]
  %0 = mul nsw i32 %indvar2, 16

[...]
for.body4:                                        ; preds = %for.cond1.preheader
  %indvar = phi i32 [ 0, %for.cond1.preheader ]
  %2 = add nuw nsw i32 %indvar, %0
  %3 = trunc i32 %2 to i16
  %arrayidx = getelementptr inbounds [64 x i16], [64 x i16]* @v, i16 0, i16 %3
  %4 = load i16, i16* %arrayidx, align 1

Note that before this patch, the above reproducer hit an error about a broken PHI, so there was certainly something strange going on then too.

bbi-60764.ll1 KBDownload

[uabelho]

I wrote a PR about it:
https://bugs.llvm.org/show_bug.cgi?id=51980

[SjoerdMeijer]

Many thanks for that and the (short) reproducer, will be looking at that soon.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

struct FlattenInfo {
   BranchInst *InnerBranch = nullptr;
   BranchInst *OuterBranch = nullptr;
   SmallPtrSet<Value *, 4> LinearIVUses;
   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;
+
   FlattenInfo(Loop *OL, Loop *IL) : OuterLoop(OL), InnerLoop(IL) {};
 };
 
 static bool
 setLoopComponents(Value *&TC, Value *&TripCount, BinaryOperator *&Increment,
                   SmallPtrSetImpl<Instruction *> &IterationInstructions) {
   TripCount = TC;
   IterationInstructions.insert(Increment);
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))
+      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 =
         InnerPHI.getIncomingValueForBlock(FI.InnerLoop->getLoopLatch());
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))
+      continue;
     if (!SafeOuterPHIs.count(&OuterPHI)) {
       LLVM_DEBUG(dbgs() << "found unsafe PHI in outer loop: "; OuterPHI.dump());
       return false;
     }
   }
 
   LLVM_DEBUG(dbgs() << "checkPHIs: OK\n");
   return true;
for (User *U : FI.InnerInductionPHI->users()) {
     // Matches the same pattern as above, except it also looks for truncs
     // on the phi, which can be the result of widening the induction variables.
     bool IsAddTrunc = match(U, m_c_Add(m_Trunc(m_Specific(FI.InnerInductionPHI)),
                                        m_Value(MatchedMul))) &&
                       match(MatchedMul,
                             m_c_Mul(m_Trunc(m_Specific(FI.OuterInductionPHI)),
                             m_Value(MatchedItCount)));
 
+    // Look through extends if the IV has been widened.
+    if (FI.Widened && isa<SExtInst>(MatchedItCount))
+      MatchedItCount = dyn_cast<SExtInst>(MatchedItCount)->getOperand(0);
+    else if (FI.Widened && isa<ZExtInst>(MatchedItCount))
+      MatchedItCount = dyn_cast<ZExtInst>(MatchedItCount)->getOperand(0);
+
     if ((IsAdd || IsAddTrunc) && MatchedItCount == InnerTripCount) {
       LLVM_DEBUG(dbgs() << "Use is optimisable\n");
       ValidOuterPHIUses.insert(MatchedMul);
       FI.LinearIVUses.insert(U);
     } else {
       LLVM_DEBUG(dbgs() << "Did not match expected pattern, bailing\n");
       return false;
     }
static bool CanWidenIV(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   if (InnerType != OuterType ||
       InnerType->getScalarSizeInBits() >= MaxLegalSize ||
       MaxLegalType->getScalarSizeInBits() < InnerType->getScalarSizeInBits() * 2) {
     LLVM_DEBUG(dbgs() << "Can't widen the IV\n");
     return false;
   }
 
   SCEVExpander Rewriter(*SE, DL, "loopflatten");
-  SmallVector<WideIVInfo, 2> WideIVs;
   SmallVector<WeakTrackingVH, 4> DeadInsts;
-  WideIVs.push_back( {FI.InnerInductionPHI, MaxLegalType, false });
-  WideIVs.push_back( {FI.OuterInductionPHI, MaxLegalType, false });
   unsigned ElimExt = 0;
   unsigned Widened = 0;
 
-  for (const auto &WideIV : WideIVs) {
+  auto CreateWideIV = [&] (WideIVInfo WideIV, bool &Deleted) -> bool {
     PHINode *WidePhi = createWideIV(WideIV, LI, SE, Rewriter, DT, DeadInsts,
                                     ElimExt, Widened, true /* HasGuards */,
                                     true /* UsePostIncrementRanges */);
     if (!WidePhi)
       return false;
     LLVM_DEBUG(dbgs() << "Created wide phi: "; WidePhi->dump());
     LLVM_DEBUG(dbgs() << "Deleting old phi: "; WideIV.NarrowIV->dump());
-    RecursivelyDeleteDeadPHINode(WideIV.NarrowIV);
-  }
-  // After widening, rediscover all the loop components.
+    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.
+  if (!Deleted)
+    FI.InnerInductionPHI->removeIncomingValue(FI.InnerLoop->getLoopLatch());
+  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);
+
+  // 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_DEBUG(
       dbgs() << "Loop flattening running on outer loop "

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

 ; CHECK-NEXT:    [[CMP215:%.*]] = icmp sgt i32 [[M:%.*]], 0
 ; CHECK-NEXT:    br i1 [[CMP215]], label [[FOR_COND1_PREHEADER_US_PREHEADER:%.*]], label [[FOR_COND1_PREHEADER_PREHEADER:%.*]]
 ; CHECK:       for.cond1.preheader.preheader:
 ; CHECK-NEXT:    br label [[FOR_COND1_PREHEADER:%.*]]
 ; CHECK:       for.cond1.preheader.us.preheader:
 ; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[M]] to i64
 ; CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[M]] to i64
 ; CHECK-NEXT:    [[TMP2:%.*]] = sext i32 [[N]] to i64
+; CHECK-NEXT:    [[FLATTEN_TRIPCOUNT:%.*]] = mul i64 [[TMP0]], [[TMP2]]
 ; CHECK-NEXT:    br label [[FOR_COND1_PREHEADER_US:%.*]]
 ; CHECK:       for.cond1.preheader.us:
 ; CHECK-NEXT:    [[INDVAR2:%.*]] = phi i64 [ [[INDVAR_NEXT3:%.*]], [[FOR_COND1_FOR_COND_CLEANUP3_CRIT_EDGE_US:%.*]] ], [ 0, [[FOR_COND1_PREHEADER_US_PREHEADER]] ]
 ; CHECK-NEXT:    [[I_018_US:%.*]] = phi i32 [ [[INC6_US:%.*]], [[FOR_COND1_FOR_COND_CLEANUP3_CRIT_EDGE_US]] ], [ 0, [[FOR_COND1_PREHEADER_US_PREHEADER]] ]
 ; CHECK-NEXT:    [[TMP3:%.*]] = mul nsw i64 [[INDVAR2]], [[TMP1]]
 ; CHECK-NEXT:    [[MUL_US:%.*]] = mul nsw i32 [[I_018_US]], [[M]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = sext i32 [[MUL_US]] to i64
+; CHECK-NEXT:    [[FLATTEN_TRUNCIV:%.*]] = trunc i64 [[INDVAR2]] to i32
 ; CHECK-NEXT:    br label [[FOR_BODY4_US:%.*]]
 ; CHECK:       for.body4.us:
-; CHECK-NEXT:    [[INDVAR:%.*]] = phi i64 [ [[INDVAR_NEXT:%.*]], [[FOR_BODY4_US]] ], [ 0, [[FOR_COND1_PREHEADER_US]] ]
-; CHECK-NEXT:    [[J_016_US:%.*]] = phi i32 [ 0, [[FOR_COND1_PREHEADER_US]] ], [ [[INC_US:%.*]], [[FOR_BODY4_US]] ]
+; CHECK-NEXT:    [[INDVAR:%.*]] = phi i64 [ 0, [[FOR_COND1_PREHEADER_US]] ]
+; CHECK-NEXT:    [[J_016_US:%.*]] = phi i32 [ 0, [[FOR_COND1_PREHEADER_US]] ]
 ; CHECK-NEXT:    [[TMP5:%.*]] = add nsw i64 [[INDVAR]], [[TMP3]]
 ; CHECK-NEXT:    [[TMP6:%.*]] = sext i32 [[J_016_US]] to i64
 ; CHECK-NEXT:    [[TMP7:%.*]] = add nsw i64 [[TMP6]], [[TMP3]]
 ; CHECK-NEXT:    [[ADD_US:%.*]] = add nsw i32 [[J_016_US]], [[MUL_US]]
-; CHECK-NEXT:    [[IDXPROM_US:%.*]] = sext i32 [[ADD_US]] to i64
-; CHECK-NEXT:    [[ARRAYIDX_US:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 [[TMP5]]
+; CHECK-NEXT:    [[IDXPROM_US:%.*]] = sext i32 [[FLATTEN_TRUNCIV]] to i64
+; CHECK-NEXT:    [[ARRAYIDX_US:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 [[INDVAR2]]
 ; CHECK-NEXT:    [[TMP8:%.*]] = load i32, i32* [[ARRAYIDX_US]], align 4
 ; CHECK-NEXT:    tail call void @g(i32 [[TMP8]])
-; CHECK-NEXT:    [[INDVAR_NEXT]] = add i64 [[INDVAR]], 1
-; CHECK-NEXT:    [[INC_US]] = add nuw nsw i32 [[J_016_US]], 1
+; CHECK-NEXT:    [[INDVAR_NEXT:%.*]] = add i64 [[INDVAR]], 1
+; CHECK-NEXT:    [[INC_US:%.*]] = add nuw nsw i32 [[J_016_US]], 1
 ; CHECK-NEXT:    [[CMP2_US:%.*]] = icmp slt i64 [[INDVAR_NEXT]], [[TMP0]]
-; CHECK-NEXT:    br i1 [[CMP2_US]], label [[FOR_BODY4_US]], label [[FOR_COND1_FOR_COND_CLEANUP3_CRIT_EDGE_US]]
+; CHECK-NEXT:    br label [[FOR_COND1_FOR_COND_CLEANUP3_CRIT_EDGE_US]]
 ; CHECK:       for.cond1.for.cond.cleanup3_crit_edge.us:
 ; CHECK-NEXT:    [[INDVAR_NEXT3]] = add i64 [[INDVAR2]], 1
 ; CHECK-NEXT:    [[INC6_US]] = add nuw nsw i32 [[I_018_US]], 1
-; CHECK-NEXT:    [[CMP_US:%.*]] = icmp slt i64 [[INDVAR_NEXT3]], [[TMP2]]
+; CHECK-NEXT:    [[CMP_US:%.*]] = icmp slt i64 [[INDVAR_NEXT3]], [[FLATTEN_TRIPCOUNT]]
 ; CHECK-NEXT:    br i1 [[CMP_US]], label [[FOR_COND1_PREHEADER_US]], label [[FOR_COND_CLEANUP_LOOPEXIT:%.*]]
 ; CHECK:       for.cond1.preheader:
 ; CHECK-NEXT:    [[I_018:%.*]] = phi i32 [ [[INC6:%.*]], [[FOR_COND1_PREHEADER]] ], [ 0, [[FOR_COND1_PREHEADER_PREHEADER]] ]
 ; CHECK-NEXT:    [[INC6]] = add nuw nsw i32 [[I_018]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[INC6]], [[N]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_COND1_PREHEADER]], label [[FOR_COND_CLEANUP_LOOPEXIT19:%.*]]
 ; CHECK:       for.cond.cleanup.loopexit:
 ; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]