[LoopFlatten] Widen the IV

[SjoerdMeijer]

Summary

Widen the IV to the widest available and legal integer type, which makes this transformations always safe so that we can skip overflow checks.

Motivation is to let this pass trigger on 64-bit targets too, and this is the last patch in a serie to achieve this: D90402 moves pass LoopFlatten to just before IndVarSimplify so that IVs are not already widened, D90421 factors out widening from IndVarSimplify into LoopUtils so that we can also use it in LoopFlatten, and D90408 is a refactoring in LoopFlatten to make it easier to reuse and rediscover loop components again after widening.

[dmgreen]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

339

Some formatting is off here.

351

I think this should be checking each use of the trunk, or checking the trunk has one use.

400

auto *V =

591

>= ?

I feel like something here needs to be checking that newbitwidth > InnerType->getScalarSizeInBits() * 2
(Which likely won't be very important for power-2 type sizes, but is still needed to make sure the new IV doesn't overflow?)

594

Can this call findLoopComponents directly, or is there something in CanFlattenloopPair that would be needed?

604

Do we need to be widening all the phi's, or just the induction?

611

What does DeadInstr contain after the calls to WidenIV? Should they be removed now, does that simplify things later on?

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

91

Can this test be cleaned up a bit (or perhaps a simpler test be added). This seems to be loop unswitched with this extra remainder loop left over.

[SjoerdMeijer]

Thanks for reviewing!
This is a rebase, will now start addressing the comments.

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

591

re: >= ?

I didn't see how the inner type can be larger than the legal max size, but since > certainly doesn't hurt I will add that.

594

Yeah, I looked at that, but there's a bit more going on there, e.g. it checks and sets the IV users. So, it's best and easiest just to run CanFlattenLoopPair again.

604

I borrowed this logic that drives createWideIV from IndVarSimplify, but that might indeed be interested in promoting all of them. But I guess you're right that we are interested in just the induction phi, so will look into that.

611

It's empty, because they are not "trivially" dead. I had a go at trying to clean it up, but it wasn't trivial, and needed a lot of code, for which we have other passes. You'll see some of the old phis in test cases, but they will cleaned up later by other passes.

[SjoerdMeijer]

Comments addressed

[SjoerdMeijer]

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

91

I have kept the test because this is my motivating example, and the IR input to loop flatten when it is generated from source. Skipping loop unswitch caused this to no longer trigger, and the test isn't that big, so I would prefer to keep it. If you insist, I will have another look though.

[dmgreen]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

353

Why continue, not returning false?

591

As far as I understand, "legal" here just means "can fit into a single register", so 64bits on aarch64. There may already be IV's that are larger.

611

Ah OK. Is it possible to call DeleteDeadPhi's on the loops? It might simplify some of the "OrigPHIs" logic if they have already been removed.

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

91

Yeah, I was mostly meaning to remove the blocks like for.cond1.preheader and for.cond1.preheader.preheader which is not part of the flattened loop. I think it should keep working on what remains if you manually edit it. If not and it doesn't keep working, but a big deal.

[SjoerdMeijer]

Comments addressed:

• now using RecursivelyDeleteDeadPHINode at the end to clean up dead phis. Doing this earlier proved to be a bit tricky with a few things in flight, but at least a nice clean up at the end.
• thanks for the suggestion: the test case has now been cleaned up.

[dmgreen]

now using RecursivelyDeleteDeadPHINode at the end to clean up dead phis. Doing this earlier proved to be a bit tricky with a few things in flight, but at least a nice clean up at the end.

What goes wrong when these are done earlier? I was hoping that it would help remove the need for OrigPHIs and Wide2OrigPHIs, cleaning up a bit of the awkwardness they are needed for.

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

7

dso_local and local_unnamed_addr #0 can be removed.

[SjoerdMeijer]

Was doing the RecursiveDelete in the wrong place. With that fixed we indeed no longer need the extra bookkeeping of OrigPhis and Orig2Wide, so that has been cleaned up.

[dmgreen]

Thanks. LGTM.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
+#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
+#include "llvm/Transforms/Utils/SimplifyIndVar.h"
 
 #define DEBUG_TYPE "loop-flatten"
 
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
 static cl::opt<unsigned> RepeatedInstructionThreshold(
     "loop-flatten-cost-threshold", cl::Hidden, cl::init(2),
struct FlattenInfo {
   Value *InnerLimit = nullptr;
   Value *OuterLimit = nullptr;
   BinaryOperator *InnerIncrement = nullptr;
   BinaryOperator *OuterIncrement = nullptr;
   BranchInst *InnerBranch = nullptr;
   BranchInst *OuterBranch = nullptr;
   SmallPtrSet<Value *, 4> LinearIVUses;
   SmallPtrSet<PHINode *, 4> InnerPHIsToTransform;
+  std::map<PHINode *, PHINode *> Wide2OrigPHIs;
+  SmallPtrSet<PHINode *, 4> OrigPHIs;
 
   FlattenInfo(Loop *OL, Loop *IL) : OuterLoop(OL), InnerLoop(IL) {};
 };
 
 // Finds the induction variable, increment and limit for a simple loop that we
 // can flatten.
 static bool findLoopComponents(
     Loop *L, SmallPtrSetImpl<Instruction *> &IterationInstructions,
static bool checkPHIs(struct FlattenInfo &FI,
 
   // 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;
+    // Ignore the original phi, i.e. the phis that were widened.
+    if (FI.OrigPHIs.find(&InnerPHI) != FI.OrigPHIs.end() )
+      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()) {
+    // Again, ignore the original phis.
+    if (FI.OrigPHIs.find(&OuterPHI) != FI.OrigPHIs.end())
+      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;
 }
 
 static bool
 checkOuterLoopInsts(struct FlattenInfo &FI,
                     SmallPtrSetImpl<Instruction *> &IterationInstructions,
                     const TargetTransformInfo *TTI) {
   // Check for instructions in the outer but not inner loop. If any of these
for (auto &I : *B) {
       RepeatedInstrCost += Cost;
     }
   }
 
   LLVM_DEBUG(dbgs() << "Cost of instructions that will be repeated: "
                     << RepeatedInstrCost << "\n");
   // Bail out if flattening the loops would cause instructions in the outer
   // loop but not in the inner loop to be executed extra times.
-  if (RepeatedInstrCost > RepeatedInstructionThreshold)
+  if (RepeatedInstrCost > RepeatedInstructionThreshold) {
+    LLVM_DEBUG(dbgs() << "checkOuterLoopInsts: not profitable, bailing.\n");
     return false;
+  }
 
+  LLVM_DEBUG(dbgs() << "checkOuterLoopInsts: OK\n");
   return true;
 }
 
 static bool checkIVUsers(struct FlattenInfo &FI) {
   // We require all uses of both induction variables to match this pattern:
   //
   //   (OuterPHI * InnerLimit) + InnerPHI
   //
   // Any uses of the induction variables not matching that pattern would
   // require a div/mod to reconstruct in the flattened loop, so the
   // transformation wouldn't be profitable.
 
+  Value *InnerLimit = FI.InnerLimit;
+  if (auto *I = dyn_cast<SExtInst>(InnerLimit))
+    InnerLimit = I->getOperand(0);
+
   // Check that all uses of the inner loop's induction variable match the
   // expected pattern, recording the uses of the outer IV.
   SmallPtrSet<Value *, 4> ValidOuterPHIUses;
   for (User *U : FI.InnerInductionPHI->users()) {
     if (U == FI.InnerIncrement)
       continue;
 
+    // After widening the IVs, a trunc instruction might have been introduced, so
+    // look through truncs.
+    if (dyn_cast<TruncInst>(U) ) {
+      if (!U->hasOneUse())
+        continue;
+      U = *U->user_begin();
+    }
+
     LLVM_DEBUG(dbgs() << "Found use of inner induction variable: "; U->dump());
 
-    Value *MatchedMul, *MatchedItCount;
-    if (match(U, m_c_Add(m_Specific(FI.InnerInductionPHI),
+    Value *MatchedMul;
+    Value *MatchedItCount;
+    bool IsAdd = match(U, m_c_Add(m_Specific(FI.InnerInductionPHI),
+                                  m_Value(MatchedMul))) &&
+                 match(MatchedMul, m_c_Mul(m_Specific(FI.OuterInductionPHI),
+                                           m_Value(MatchedItCount)));
+
+    // 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_Specific(FI.OuterInductionPHI),
-                      m_Value(MatchedItCount))) &&
-        MatchedItCount == FI.InnerLimit) {
+                            m_c_Mul(m_Trunc(m_Specific(FI.OuterInductionPHI)),
+                            m_Value(MatchedItCount)));
+
+    if ((IsAdd || IsAddTrunc) && MatchedItCount == InnerLimit) {
       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;
     }
   }
 
   // Check that there are no uses of the outer IV other than the ones found
   // as part of the pattern above.
   for (User *U : FI.OuterInductionPHI->users()) {
     if (U == FI.OuterIncrement)
       continue;
 
+    auto IsValidOuterPHIUses = [&] (User *U) -> bool {
       LLVM_DEBUG(dbgs() << "Found use of outer induction variable: "; U->dump());
-
       if (!ValidOuterPHIUses.count(U)) {
         LLVM_DEBUG(dbgs() << "Did not match expected pattern, bailing\n");
         return false;
-    } else {
+      }
       LLVM_DEBUG(dbgs() << "Use is optimisable\n");
+      return true;
+    };
+
+    if (auto *V = dyn_cast<TruncInst>(U)) {
+      for (auto *K : V->users()) {
+        if (!IsValidOuterPHIUses(K))
+          return false;
       }
+      continue;
     }
 
-  LLVM_DEBUG(dbgs() << "Found " << FI.LinearIVUses.size()
+    if (!IsValidOuterPHIUses(U))
+      return false;
+  }
+
+  LLVM_DEBUG(dbgs() << "checkIVUsers: OK\n";
+             dbgs() << "Found " << FI.LinearIVUses.size()
                     << " value(s) that can be replaced:\n";
              for (Value *V : FI.LinearIVUses) {
                dbgs() << "  ";
                V->dump();
              });
-
   return true;
 }
 
 // Return an OverflowResult dependant on if overflow of the multiplication of
 // InnerLimit and OuterLimit can be assumed not to happen.
 static OverflowResult checkOverflow(struct FlattenInfo &FI,
                                     DominatorTree *DT, AssumptionCache *AC) {
   Function *F = FI.OuterLoop->getHeader()->getParent();
for (Value *U : V->users()) {
         }
       }
     }
   }
 
   return OverflowResult::MayOverflow;
 }
 
-static bool FlattenLoopPair(struct FlattenInfo &FI, DominatorTree *DT,
+static bool CanFlattenLoopPair(struct FlattenInfo &FI, DominatorTree *DT,
                                LoopInfo *LI, ScalarEvolution *SE,
-                            AssumptionCache *AC, TargetTransformInfo *TTI) {
-  Function *F = FI.OuterLoop->getHeader()->getParent();
-  LLVM_DEBUG(dbgs() << "Loop flattening running on outer loop "
-                    << FI.OuterLoop->getHeader()->getName() << " and inner loop "
-                    << FI.InnerLoop->getHeader()->getName() << " in "
-                    << F->getName() << "\n");
-
+                               AssumptionCache *AC, const TargetTransformInfo *TTI) {
   SmallPtrSet<Instruction *, 8> IterationInstructions;
   if (!findLoopComponents(FI.InnerLoop, IterationInstructions, FI.InnerInductionPHI,
                           FI.InnerLimit, FI.InnerIncrement, FI.InnerBranch, SE))
     return false;
   if (!findLoopComponents(FI.OuterLoop, IterationInstructions, FI.OuterInductionPHI,
                           FI.OuterLimit, FI.OuterIncrement, FI.OuterBranch, SE))
     return false;
 
static bool CanFlattenLoopPair(struct FlattenInfo &FI, DominatorTree *DT,
   // Find the values in the loop that can be replaced with the linearized
   // induction variable, and check that there are no other uses of the inner
   // or outer induction variable. If there were, we could still do this
   // transformation, but we'd have to insert a div/mod to calculate the
   // original IVs, so it wouldn't be profitable.
   if (!checkIVUsers(FI))
     return false;
 
-  // Check if the new iteration variable might overflow. In this case, we
-  // need to version the loop, and select the original version at runtime if
-  // the iteration space is too large.
-  // TODO: We currently don't version the loop.
-  // TODO: it might be worth using a wider iteration variable rather than
-  // versioning the loop, if a wide enough type is legal.
-  bool MustVersionLoop = true;
-  OverflowResult OR = checkOverflow(FI, DT, AC);
-  if (OR == OverflowResult::AlwaysOverflowsHigh ||
-      OR == OverflowResult::AlwaysOverflowsLow) {
-    LLVM_DEBUG(dbgs() << "Multiply would always overflow, so not profitable\n");
-    return false;
-  } else if (OR == OverflowResult::MayOverflow) {
-    LLVM_DEBUG(dbgs() << "Multiply might overflow, not flattening\n");
-  } else {
-    LLVM_DEBUG(dbgs() << "Multiply cannot overflow, modifying loop in-place\n");
-    MustVersionLoop = false;
-  }
+  LLVM_DEBUG(dbgs() << "CanFlattenLoopPair: OK\n");
+  return true;
 
-  // We cannot safely flatten the loop. Exit now.
-  if (MustVersionLoop)
-    return false;
-
-  // Do the actual transformation.
+static bool DoFlattenLoopPair(struct FlattenInfo &FI, DominatorTree *DT,
+                              LoopInfo *LI, ScalarEvolution *SE,
+                              AssumptionCache *AC,
+                              const TargetTransformInfo *TTI) {
+  Function *F = FI.OuterLoop->getHeader()->getParent();
   LLVM_DEBUG(dbgs() << "Checks all passed, doing the transformation\n");
-
   {
     using namespace ore;
     OptimizationRemark Remark(DEBUG_TYPE, "Flattened", FI.InnerLoop->getStartLoc(),
                               FI.InnerLoop->getHeader());
     OptimizationRemarkEmitter ORE(F);
     Remark << "Flattened into outer loop";
     ORE.emit(Remark);
   }
 
   Value *NewTripCount =
       BinaryOperator::CreateMul(FI.InnerLimit, FI.OuterLimit, "flatten.tripcount",
                                 FI.OuterLoop->getLoopPreheader()->getTerminator());
   LLVM_DEBUG(dbgs() << "Created new trip count in preheader: ";
              NewTripCount->dump());
 
   // Fix up PHI nodes that take values from the inner loop back-edge, which
   // we are about to remove.
   FI.InnerInductionPHI->removeIncomingValue(FI.InnerLoop->getLoopLatch());
+  // The old Phi will be optimised away later, but for now we can't leave
+  // leave it in an invalid state, so are updating them too.
+  if (FI.Wide2OrigPHIs.find(FI.InnerInductionPHI) != FI.Wide2OrigPHIs.end()) {
+    auto *OrigPHI = FI.Wide2OrigPHIs.find(FI.InnerInductionPHI)->second;
+    OrigPHI->removeIncomingValue(FI.InnerLoop->getLoopLatch());
+  }
   for (PHINode *PHI : FI.InnerPHIsToTransform)
     PHI->removeIncomingValue(FI.InnerLoop->getLoopLatch());
 
   // Modify the trip count of the outer loop to be the product of the two
   // trip counts.
   cast<User>(FI.OuterBranch->getCondition())->setOperand(1, NewTripCount);
 
   // Replace the inner loop backedge with an unconditional branch to the exit.
   BasicBlock *InnerExitBlock = FI.InnerLoop->getExitBlock();
   BasicBlock *InnerExitingBlock = FI.InnerLoop->getExitingBlock();
   InnerExitingBlock->getTerminator()->eraseFromParent();
   BranchInst::Create(InnerExitBlock, InnerExitingBlock);
   DT->deleteEdge(InnerExitingBlock, FI.InnerLoop->getHeader());
 
+  auto HasSExtUser = [] (Value *V) -> Value * {
+    for (User *U : V->users() )
+      if (dyn_cast<SExtInst>(U))
+        return U;
+    return nullptr;
+  };
+
   // Replace all uses of the polynomial calculated from the two induction
   // variables with the one new one.
-  for (Value *V : FI.LinearIVUses)
+  for (Value *V : FI.LinearIVUses) {
+    // If the induction variable has been widened, look through the SExt.
+    if (Value *U = HasSExtUser(V))
+      V = U;
     V->replaceAllUsesWith(FI.OuterInductionPHI);
+  }
 
   // Tell LoopInfo, SCEV and the pass manager that the inner loop has been
   // deleted, and any information that have about the outer loop invalidated.
   SE->forgetLoop(FI.OuterLoop);
   SE->forgetLoop(FI.InnerLoop);
   LI->erase(FI.InnerLoop);
   return true;
 }
 
+static bool CanWidenIV(struct FlattenInfo &FI, DominatorTree *DT,
+                       LoopInfo *LI, ScalarEvolution *SE,
+                       AssumptionCache *AC, const TargetTransformInfo *TTI) {
+  LLVM_DEBUG(dbgs() << "Try widening the IVs\n");
+  Module *M = FI.InnerLoop->getHeader()->getParent()->getParent();
+  auto &DL = M->getDataLayout();
+  auto *InnerType = FI.InnerInductionPHI->getType();
+  auto *OuterType = FI.OuterInductionPHI->getType();
+  unsigned MaxLegalSize = DL.getLargestLegalIntTypeSizeInBits();
+  auto *MaxLegalType = DL.getLargestLegalIntType(M->getContext());
+
+  // If both induction types are less than the maximum legal integer width,
+  // promote both to the widest type available so we know calculating
+  // (OuterLimit * InnerLimit) as the new trip count is safe.
+  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 });
+  FI.OrigPHIs.insert(FI.InnerInductionPHI);
+  FI.OrigPHIs.insert(FI.OuterInductionPHI);
+  unsigned ElimExt;
+  unsigned Widened;
+
+  for (unsigned i = 0; i < WideIVs.size(); i++) {
+    PHINode *WidePhi = createWideIV(WideIVs[i], LI, SE, Rewriter, DT, DeadInsts,
+                                    ElimExt, Widened, true /* HasGuards */,
+                                    true /* UsePostIncrementRanges */);
+    if (!WidePhi)
+      return false;
+    LLVM_DEBUG(dbgs() << "Created wide phi: "; WidePhi->dump());
+    FI.Wide2OrigPHIs.insert({ WidePhi, WideIVs[i].NarrowIV});
+  }
+  // After widening, rediscover all the loop components.
+  return CanFlattenLoopPair(FI, DT, LI, SE, AC, TTI);
+}
+
+static bool FlattenLoopPair(struct FlattenInfo &FI, DominatorTree *DT,
+                            LoopInfo *LI, ScalarEvolution *SE,
+                            AssumptionCache *AC,
+                            const TargetTransformInfo *TTI) {
+  Function *F = FI.OuterLoop->getHeader()->getParent();
+  LLVM_DEBUG(dbgs() << "Loop flattening running on outer loop "
+                    << FI.OuterLoop->getHeader()->getName() << " and inner loop "
+                    << FI.InnerLoop->getHeader()->getName() << " in "
+                    << F->getName() << "\n");
+
+  if (!CanFlattenLoopPair(FI, DT, LI, SE, AC, TTI))
+    return false;
+
+  // Check if we can widen the induction variables to avoid overflow checks.
+  if (CanWidenIV(FI, DT, LI, SE, AC, TTI))
+    return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI);
+
+  // Check if the new iteration variable might overflow. In this case, we
+  // need to version the loop, and select the original version at runtime if
+  // the iteration space is too large.
+  // TODO: We currently don't version the loop.
+  OverflowResult OR = checkOverflow(FI, DT, AC);
+  if (OR == OverflowResult::AlwaysOverflowsHigh ||
+      OR == OverflowResult::AlwaysOverflowsLow) {
+    LLVM_DEBUG(dbgs() << "Multiply would always overflow, so not profitable\n");
+    return false;
+  } else if (OR == OverflowResult::MayOverflow) {
+    LLVM_DEBUG(dbgs() << "Multiply might overflow, not flattening\n");
+    return false;
+  }
+
+  LLVM_DEBUG(dbgs() << "Multiply cannot overflow, modifying loop in-place\n");
+  return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI);
+}
+
 bool Flatten(DominatorTree *DT, LoopInfo *LI, ScalarEvolution *SE,
              AssumptionCache *AC, TargetTransformInfo *TTI) {
   bool Changed = false;
   for (auto *InnerLoop : LI->getLoopsInPreorder()) {
     auto *OuterLoop = InnerLoop->getParentLoop();
     if (!OuterLoop)
       continue;
     struct FlattenInfo FI(OuterLoop, InnerLoop);

llvm/test/Transforms/LoopFlatten/loop-flatten-negative.ll

 ; RUN: opt < %s -S -loop-flatten -debug-only=loop-flatten 2>&1 | FileCheck %s
 ; REQUIRES: asserts
 
+target datalayout = "e-m:e-p:32:32-i64:64-v128:64:128-a:0:32-n32-S64"
+
 ; Every function in this file has a reason that it can't be transformed.
 
 ; CHECK-NOT: Checks all passed, doing the transformation
 
 ; Outer loop does not start at zero
 define void @test_1(i32 %N, i32* nocapture %C, i32* nocapture readonly %A, i32 %scale) {
 entry:
   %cmp25 = icmp sgt i32 %N, 0

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

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -S -loop-flatten -verify-loop-info -verify-dom-info -verify-scev -verify | FileCheck %s
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+
+; Function Attrs: nounwind
+define dso_local void @foo(i32* %A, i32 %N, i32 %M) local_unnamed_addr #0 {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP17:%.*]] = icmp sgt i32 [[N:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP17]], label [[FOR_COND1_PREHEADER_LR_PH:%.*]], label [[FOR_COND_CLEANUP:%.*]]
+; CHECK:       for.cond1.preheader.lr.ph:
+; 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 [[N]] to i64
+; CHECK-NEXT:    [[FLATTEN_TRIPCOUNT:%.*]] = mul i64 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    br label [[FOR_COND1_PREHEADER_US:%.*]]
+; CHECK:       for.cond1.preheader.us:
+; CHECK-NEXT:    [[INDVAR1:%.*]] = phi i64 [ [[INDVAR_NEXT2:%.*]], [[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:    [[TMP2:%.*]] = trunc i64 [[INDVAR1]] to i32
+; CHECK-NEXT:    [[MUL_US:%.*]] = mul nsw i32 [[TMP2]], [[M]]
+; CHECK-NEXT:    br label [[FOR_BODY4_US:%.*]]
+; CHECK:       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:    [[TMP3:%.*]] = trunc i64 [[INDVAR]] to i32
+; CHECK-NEXT:    [[ADD_US:%.*]] = add nsw i32 [[TMP3]], [[MUL_US]]
+; CHECK-NEXT:    [[IDXPROM_US:%.*]] = sext i32 [[ADD_US]] to i64
+; CHECK-NEXT:    [[ARRAYIDX_US:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 [[INDVAR1]]
+; CHECK-NEXT:    tail call void @f(i32* [[ARRAYIDX_US]])
+; 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 label [[FOR_COND1_FOR_COND_CLEANUP3_CRIT_EDGE_US]]
+; CHECK:       for.cond1.for.cond.cleanup3_crit_edge.us:
+; CHECK-NEXT:    [[INDVAR_NEXT2]] = add i64 [[INDVAR1]], 1
+; CHECK-NEXT:    [[INC6_US]] = add nuw nsw i32 [[I_018_US]], 1
+; CHECK-NEXT:    [[CMP_US:%.*]] = icmp slt i64 [[INDVAR_NEXT2]], [[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_LOOPEXIT20:%.*]]
+; CHECK:       for.cond.cleanup.loopexit:
+; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
+; CHECK:       for.cond.cleanup.loopexit20:
+; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp17 = icmp sgt i32 %N, 0
+  br i1 %cmp17, label %for.cond1.preheader.lr.ph, label %for.cond.cleanup
+
+for.cond1.preheader.lr.ph:                        ; preds = %entry
+  %cmp215 = icmp sgt i32 %M, 0
+  br i1 %cmp215, label %for.cond1.preheader.us.preheader, label %for.cond1.preheader.preheader
+
+for.cond1.preheader.preheader:                    ; preds = %for.cond1.preheader.lr.ph
+  br label %for.cond1.preheader
+
+for.cond1.preheader.us.preheader:                 ; preds = %for.cond1.preheader.lr.ph
+  br label %for.cond1.preheader.us
+
+for.cond1.preheader.us:                           ; preds = %for.cond1.preheader.us.preheader, %for.cond1.for.cond.cleanup3_crit_edge.us
+  %i.018.us = phi i32 [ %inc6.us, %for.cond1.for.cond.cleanup3_crit_edge.us ], [ 0, %for.cond1.preheader.us.preheader ]
+  %mul.us = mul nsw i32 %i.018.us, %M
+  br label %for.body4.us
+
+for.body4.us:                                     ; preds = %for.cond1.preheader.us, %for.body4.us
+  %j.016.us = phi i32 [ 0, %for.cond1.preheader.us ], [ %inc.us, %for.body4.us ]
+  %add.us = add nsw i32 %j.016.us, %mul.us
+  %idxprom.us = sext i32 %add.us to i64
+  %arrayidx.us = getelementptr inbounds i32, i32* %A, i64 %idxprom.us
+  tail call void @f(i32* %arrayidx.us) #2
+  %inc.us = add nuw nsw i32 %j.016.us, 1
+  %cmp2.us = icmp slt i32 %inc.us, %M
+  br i1 %cmp2.us, label %for.body4.us, label %for.cond1.for.cond.cleanup3_crit_edge.us
+
+for.cond1.for.cond.cleanup3_crit_edge.us:         ; preds = %for.body4.us
+  %inc6.us = add nuw nsw i32 %i.018.us, 1
+  %cmp.us = icmp slt i32 %inc6.us, %N
+  br i1 %cmp.us, label %for.cond1.preheader.us, label %for.cond.cleanup.loopexit
+
+for.cond1.preheader:                              ; preds = %for.cond1.preheader.preheader, %for.cond1.preheader
+  %i.018 = phi i32 [ %inc6, %for.cond1.preheader ], [ 0, %for.cond1.preheader.preheader ]
+  %inc6 = add nuw nsw i32 %i.018, 1
+  %cmp = icmp slt i32 %inc6, %N
+  br i1 %cmp, label %for.cond1.preheader, label %for.cond.cleanup.loopexit20
+
+for.cond.cleanup.loopexit:                        ; preds = %for.cond1.for.cond.cleanup3_crit_edge.us
+  br label %for.cond.cleanup
+
+for.cond.cleanup.loopexit20:                      ; preds = %for.cond1.preheader
+  br label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.cond.cleanup.loopexit20, %for.cond.cleanup.loopexit, %entry
+  ret void
+}
+
+declare dso_local void @f(i32* %0) local_unnamed_addr #1