[LoopFlatten] Handle A[i][j] uses of IV's

[dmgreen]

Summary

This allow A[i][j] uses of the IV's in loop flattening,
along with the existing A[i*N+j] uses.

lib/Transforms/Scalar/LoopFlatten.cpp

if (RepeatedInstrCost > RepeatedInstructionThreshold)
     return false;
 
   return true;
 }
 
 static bool checkIVUsers(PHINode *InnerPHI, PHINode *OuterPHI,
                          BinaryOperator *InnerIncrement,
                          BinaryOperator *OuterIncrement, Value *InnerLimit,
-                         SmallPtrSetImpl<Value *> &LinearIVUses) {
+                         SmallPtrSetImpl<Value *> &LinearIVUses,
+                         SmallPtrSetImpl<GetElementPtrInst *> &GepIVUses) {
   // 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.
 
for (User *U : InnerPHI->users()) {
     Value *MatchedMul, *MatchedItCount;
     if (match(U, m_c_Add(m_Specific(InnerPHI), m_Value(MatchedMul))) &&
         match(MatchedMul,
               m_c_Mul(m_Specific(OuterPHI), m_Value(MatchedItCount))) &&
         MatchedItCount == InnerLimit) {
       DEBUG(dbgs() << "Use is optimisable\n");
       ValidOuterPHIUses.insert(MatchedMul);
       LinearIVUses.insert(U);
+    } else if (dyn_cast<Constant>(InnerLimit) &&
+               dyn_cast<GetElementPtrInst>(U) && U->getNumOperands() == 3 &&
+               U->getOperand(1) == OuterPHI && U->getOperand(2) == InnerPHI &&
+               U->getOperand(0)->getType()->isPointerTy() &&
+               U->getOperand(0)->getType()->getPointerElementType()->isArrayTy() &&
+               U->getOperand(0)->getType()->getPointerElementType()->getArrayNumElements() ==
+                   cast<Constant>(InnerLimit)->getUniqueInteger()) {
+      DEBUG(dbgs() << "Use is gep optimisable\n");
+      ValidOuterPHIUses.insert(U);
+      GepIVUses.insert(cast<GetElementPtrInst>(U));
     } else {
       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 : OuterPHI->users()) {
     if (U == OuterIncrement)
       continue;
 
     DEBUG(dbgs() << "Found use of outer induction variable: "; U->dump());
 
     if (!ValidOuterPHIUses.count(U)) {
       DEBUG(dbgs() << "Did not match expected pattern, bailing\n");
       return false;
     } else {
       DEBUG(dbgs() << "Use is optimisable\n");
     }
   }
 
-  DEBUG(dbgs() << "Found " << LinearIVUses.size()
+  DEBUG(dbgs() << "Found " << LinearIVUses.size() + GepIVUses.size()
                << " value(s) that can be replaced:\n";
         for (Value *V : LinearIVUses) {
           dbgs() << "  "; V->dump();
+        }
+        for (Value *V : GepIVUses) {
+          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(Loop *OuterLoop, Value *InnerLimit,
-                                    Value *OuterLimit,
-                                    SmallPtrSetImpl<Value *> &LinearIVUses,
-                                    DominatorTree *DT, AssumptionCache *AC) {
+static OverflowResult
+checkOverflow(Loop *OuterLoop, Value *InnerLimit, Value *OuterLimit,
+              SmallPtrSetImpl<GetElementPtrInst *> &GepIVUses,
   Function *F = OuterLoop->getHeader()->getParent();
   const DataLayout &DL = F->getParent()->getDataLayout();
 
   // For debugging/testing.
   if (AssumeNoOverflow)
     return OverflowResult::NeverOverflows;
 
   // Check if the multiply could not overflow due to known ranges of the
   // input values.
   OverflowResult OR = computeOverflowForUnsignedMul(
       InnerLimit, OuterLimit, DL, AC,
       OuterLoop->getLoopPreheader()->getTerminator(), DT);
   if (OR != OverflowResult::MayOverflow)
     return OR;
 
-  for (Value *V : LinearIVUses) {
-    for (Value *U : V->users()) {
-      if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) {
-        // The IV is used as the operand of a GEP, and the IV is at least as
-        // wide as the address space of the GEP. In this case, the GEP would
-        // wrap around the address space before the IV increment wraps, which
-        // would be UB.
-        if (GEP->isInBounds() &&
-            V->getType()->getIntegerBitWidth() >=
-                DL.getPointerTypeSizeInBits(GEP->getType())) {
-          DEBUG(dbgs() << "use of linear IV would be UB if overflow occurred: ";
-                GEP->dump());
-          return OverflowResult::NeverOverflows;
-        }
-      }
-    }
+  auto CheckGEPForOverflow = [&](GetElementPtrInst *GEP) {
+        GEP->getOperand(1)->getType()->getIntegerBitWidth() >=
+      return true;
     }
+    return false;
+  };
+
+  for (Value *V : LinearIVUses)
+    for (Value *U : V->users())
+      if (auto *GEP = dyn_cast<GetElementPtrInst>(U))
+        if(CheckGEPForOverflow(GEP))
+          return OverflowResult::NeverOverflows;
+
+  for (GetElementPtrInst *V : GepIVUses)
+    if (CheckGEPForOverflow(V))
+      return OverflowResult::NeverOverflows;
 
   return OverflowResult::MayOverflow;
 }
 
 static bool FlattenLoopPair(Loop *OuterLoop, Loop *InnerLoop, DominatorTree *DT,
                             LoopInfo *LI, ScalarEvolution *SE,
                             const LoopAccessInfo *LAI, AssumptionCache *AC,
                             TargetTransformInfo *TTI,
                             std::function<void(Loop *)> markLoopAsDeleted) {
   Function *F = OuterLoop->getHeader()->getParent();
-
   DEBUG(dbgs() << "Running on outer loop " << OuterLoop->getHeader()->getName()
                << " and inner loop " << InnerLoop->getHeader()->getName()
                << " in " << F->getName() << "\n");
 
   SmallPtrSet<Instruction *, 8> IterationInstructions;
 
   PHINode *InnerInductionPHI, *OuterInductionPHI;
   Value *InnerLimit, *OuterLimit;
if (!checkOuterLoopInsts(OuterLoop, InnerLoop, IterationInstructions,
     return false;
 
   // 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.
   SmallPtrSet<Value *, 4> LinearIVUses;
+  SmallPtrSet<GetElementPtrInst *, 4> GepIVUses;
   if (!checkIVUsers(InnerInductionPHI, OuterInductionPHI, InnerIncrement,
-                    OuterIncrement, InnerLimit, LinearIVUses))
+                    OuterIncrement, InnerLimit, LinearIVUses, GepIVUses))
     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: 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(OuterLoop, InnerLimit, OuterLimit, LinearIVUses, DT, AC);
+  OverflowResult OR = checkOverflow(OuterLoop, InnerLimit, OuterLimit,
+                                    LinearIVUses, GepIVUses, DT, AC);
   if (OR == OverflowResult::AlwaysOverflows) {
     DEBUG(dbgs() << "Multiply would always overflow, so not profitable\n");
     return false;
   } else if (OR == OverflowResult::MayOverflow) {
     DEBUG(dbgs() << "Multiply might overflow, versioning loop\n");
   } else {
     DEBUG(dbgs() << "Multiply cannot overflow, modifying loop in-place\n");
     MustVersionLoop = false;
static bool FlattenLoopPair(Loop *OuterLoop, Loop *InnerLoop, DominatorTree *DT,
   InnerExitingBlock->getTerminator()->eraseFromParent();
   BranchInst::Create(InnerExitBlock, InnerExitingBlock);
   DT->deleteEdge(InnerExitingBlock, InnerLoop->getHeader());
 
   // Replace all uses of the polynomial calculated from the two induction
   // variables with the one new one.
   for (Value *V : LinearIVUses)
     V->replaceAllUsesWith(OuterInductionPHI);
+  for (GetElementPtrInst *V : GepIVUses) {
+    V->setOperand(1, Constant::getNullValue(V->getOperand(1)->getType()));
+    V->setOperand(2, OuterInductionPHI);
+  }
 
   // If we made a fallback copy of the loop, it will still be flattenable if
   // this pass is run again, but that wouldn't be profitable, so disable
   // flattening of that loop.
   if (FallbackLoop)
     FallbackLoop->setLoopAlreadyFlattened();
 
   // Tell LoopInfo, SCEV and the pass manager that the inner loop has been

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

for.inc6:                                         ; preds = %for.body3
   br i1 %exitcond19, label %for.body, label %for.end8
 
 for.end8:                                         ; preds = %for.inc6
   ret i32 10
 }
 
 
 ; Outer loop conditional phi
-define i32 @e() {
+define i32 @test_10() {
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.end16
   %f.033 = phi i32 [ 0, %entry ], [ %inc18, %for.end16 ]
   %g.032 = phi i32 [ undef, %entry ], [ %g.3.lcssa, %for.end16 ]
   %.pr = add i32 10, 10
   %tobool29 = icmp eq i32 %.pr, 0
for.end16:                                        ; preds = %lor.end
   %g.3.lcssa = phi i32 [ %g.3, %lor.end ]
   %inc18 = add nuw nsw i32 %f.033, 1
   %exitcond34 = icmp ne i32 %inc18, 7
   br i1 %exitcond34, label %for.body, label %for.end19
 
 for.end19:                                        ; preds = %for.end16
   ret i32 undef
 }
+
+define i32 @test11([8 x float]* nocapture %A) {
+entry:
+  br label %for.body
+
+for.body:
+  %i.016 = phi i32 [ 0, %entry ], [ %inc6, %for.inc5 ]
+  %mul = mul nuw nsw i32 %i.016, 10
+  br label %for.body3
+
+for.body3:
+  %j.015 = phi i32 [ 0, %for.body ], [ %inc, %for.body3 ]
+  %add = add nuw nsw i32 %j.015, %mul
+  %conv = sitofp i32 %add to float
+  %arrayidx4 = getelementptr inbounds [8 x float], [8 x float]* %A, i32 %i.016, i32 %j.015
+  store float %conv, float* %arrayidx4, align 4
+  %inc = add nuw nsw i32 %j.015, 1
+  %exitcond = icmp ne i32 %inc, 10
+  br i1 %exitcond, label %for.body3, label %for.inc5
+
+for.inc5:
+  %inc6 = add nuw nsw i32 %i.016, 1
+  %exitcond17 = icmp ne i32 %inc6, 10
+  br i1 %exitcond17, label %for.body, label %for.end7
+
+for.end7:
+  ret i32 0
+}

test/Transforms/LoopFlatten/loop-flatten.ll

 ; CHECK: br label %for.cond.cleanup
 
 for.cond.cleanup:                                 ; preds = %for.cond.cleanup.loopexit26, %for.cond.cleanup.loopexit, %entry
   ret void
 ; CHECK: for.cond.cleanup:
 ; CHECK: ret void
 }
 
-declare i32 @func(i32)
+; CHECK-LABEL: test11
+; A[i][j]
+define i32 @test11([10 x float]* nocapture %A) local_unnamed_addr #0 {
+entry:
+  br label %for.body
+; CHECK: entry:
+; CHECK:   %flatten.tripcount = mul i32 10, 10
+; CHECK:   br label %for.body
+
+for.body:
+  %i.016 = phi i32 [ 0, %entry ], [ %inc6, %for.inc5 ]
+  %mul = mul nuw nsw i32 %i.016, 10
+  br label %for.body3
+; CHECK: for.body:
+; CHECK:   %i.016 = phi i32 [ 0, %entry ], [ %inc6, %for.inc5 ]
+; CHECK:   %mul = mul nuw nsw i32 %i.016, 10
+; CHECK:   br label %for.body3
+
+for.body3:
+  %j.015 = phi i32 [ 0, %for.body ], [ %inc, %for.body3 ]
+  %add = add nuw nsw i32 %j.015, %mul
+  %conv = sitofp i32 %add to float
+  %arrayidx4 = getelementptr inbounds [10 x float], [10 x float]* %A, i32 %i.016, i32 %j.015
+  store float %conv, float* %arrayidx4, align 4
+  %inc = add nuw nsw i32 %j.015, 1
+  %exitcond = icmp ne i32 %inc, 10
+  br i1 %exitcond, label %for.body3, label %for.inc5
+; CHECK: for.body3:
+; CHECK:   %j.015 = phi i32 [ 0, %for.body ]
+; CHECK:   %add = add nuw nsw i32 %j.015, %mul
+; CHECK:   %conv = sitofp i32 %i.016 to float
+; CHECK:   %arrayidx4 = getelementptr inbounds [10 x float], [10 x float]* %A, i32 0, i32 %i.016
+; CHECK:   store float %conv, float* %arrayidx4, align 4
+; CHECK:   %inc = add nuw nsw i32 %j.015, 1
+; CHECK:   %exitcond = icmp ne i32 %inc, 10
+; CHECK:   br label %for.inc5
+
+for.inc5:
+  %inc6 = add nuw nsw i32 %i.016, 1
+  %exitcond17 = icmp ne i32 %inc6, 10
+  br i1 %exitcond17, label %for.body, label %for.end7
+; CHECK: for.inc5:
+; CHECK:   %inc6 = add nuw nsw i32 %i.016, 1
+; CHECK:   %exitcond17 = icmp ne i32 %inc6, %flatten.tripcount
+; CHECK:   br i1 %exitcond17, label %for.body, label %for.end7
+
+for.end7:
+  ret i32 0
+}
+
 
+declare i32 @func(i32)