[LoopFlatten] Fix overflow check

[RosieSumpter]

Summary

Amended the GEP check in checkOverflow to ensure that the GEP
dominates the loop latch. If the GEP doesn't dominate the latch
then it is not necessarily calculated on each loop iteration, so
the assumption that the GEP will overflow (and be UB) before the
IV no longer holds.

[dmgreen]

I think this needs to check that the GEP is used by a load/store that is executed on each iteration.

Can this use one of the isGuaranteedToExecute function too? There is a isGuaranteedToExecuteForEveryIteration, but it does not look very powerful. It may be OK for most loops though, where checking the inner loop header is enough.

[RosieSumpter]

Changed the check to see if the GEP is used by a load/store instruction
and, if so, isGuaranteedToExecuteForEveryIteration is used with the instruction and
inner loop to determine if the GEP dominates the latch.

[RosieSumpter]

In D108107#2946377, @dmgreen wrote:

I think this needs to check that the GEP is used by a load/store that is executed on each iteration.

Can this use one of the isGuaranteedToExecute function too? There is a isGuaranteedToExecuteForEveryIteration, but it does not look very powerful. It may be OK for most loops though, where checking the inner loop header is enough.

Hopefully this change is what you had in mind - I'm just not sure whether the

if (isa<LoadInst>(GEPUser) || isa<StoreInst>(GEPUser))

is necessary (can a GEP be used by something that isn't a store or a load?)

[dmgreen]

Yeah, sounds good to me. There is quite a bit of nesting, but I hope it will be rare to get into it, or quick when it does usually finding the UB on the first item it visits.

In D108107#2946858, @RosieSumpter wrote:

Hopefully this change is what you had in mind - I'm just not sure whether the

if (isa<LoadInst>(GEPUser) || isa<StoreInst>(GEPUser))

is necessary (can a GEP be used by something that isn't a store or a load?)

Yeah sounds good. Can you add one extra check that the Store is using the GEP for the address, not the stored value. So it is operand #1, not #0.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

501

It may be better to reverse the condition so that we reduce the indentation and don't require GEPDominatesLatch any more.

if (!isa<Load>(..) && !isa<Store>(..))
  continue;
if (!isGuaranteedToExecuteForEveryIteration(..))
  continue;

[RosieSumpter]

• Added check that GEP is used as address of store instruction
• Reversed logic in some places and removed GEPDominatesLatch

[dmgreen]

LGTM. Thanks

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

OverflowResult OR = computeOverflowForUnsignedMul(
       FI.InnerTripCount, FI.OuterTripCount, DL, AC,
       FI.OuterLoop->getLoopPreheader()->getTerminator(), DT);
   if (OR != OverflowResult::MayOverflow)
     return OR;
 
   for (Value *V : FI.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.
+        for (Value *GEPUser : U->users()) {
+          Instruction *GEPUserInst = dyn_cast<Instruction>(GEPUser);
+          if (!isa<LoadInst>(GEPUserInst) &&
+              !(isa<StoreInst>(GEPUserInst) &&
+                GEP == GEPUserInst->getOperand(1)))
+            continue;
+          if (!isGuaranteedToExecuteForEveryIteration(GEPUserInst,
+                                                      FI.InnerLoop))
+            continue;
+          // The IV is used as the operand of a GEP which dominates the loop
+          // latch, 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())) {
             LLVM_DEBUG(
                 dbgs() << "use of linear IV would be UB if overflow occurred: ";
                 GEP->dump());
             return OverflowResult::NeverOverflows;
           }
         }
       }
     }
+  }
 
   return OverflowResult::MayOverflow;
 }
 
 static bool CanFlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                                ScalarEvolution *SE, AssumptionCache *AC,
                                const TargetTransformInfo *TTI) {
   SmallPtrSet<Instruction *, 8> IterationInstructions;

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

for.refetch:
   %3 = icmp eq i32 %c, 0
   br i1 %3, label %for.empty.loopexit, label %for.loopbody.outer
 for.empty.loopexit:
   br label %for.empty
 for.empty:
   ret void
 } 
 
+; GEP doesn't dominate the loop latch so can't guarantee N*M won't overflow.
+@first = global i32 1, align 4
+@a = external global [0 x i8], align 1
+define void @overflow(i32 %lim, i8* %a) {
+entry:
+  %cmp17.not = icmp eq i32 %lim, 0
+  br i1 %cmp17.not, label %for.cond.cleanup, label %for.cond1.preheader.preheader
+
+for.cond1.preheader.preheader:
+  br label %for.cond1.preheader
+
+for.cond1.preheader:
+  %i.018 = phi i32 [ %inc6, %for.cond.cleanup3 ], [ 0, %for.cond1.preheader.preheader ]
+  %mul = mul i32 %i.018, 100000
+  br label %for.body4
+
+for.cond.cleanup.loopexit:
+  br label %for.cond.cleanup
+
+for.cond.cleanup:
+  ret void
+
+for.cond.cleanup3:
+  %inc6 = add i32 %i.018, 1
+  %cmp = icmp ult i32 %inc6, %lim
+  br i1 %cmp, label %for.cond1.preheader, label %for.cond.cleanup.loopexit
+
+for.body4:
+  %j.016 = phi i32 [ 0, %for.cond1.preheader ], [ %inc, %if.end ]
+  %add = add i32 %j.016, %mul
+  %0 = load i32, i32* @first, align 4
+  %tobool.not = icmp eq i32 %0, 0
+  br i1 %tobool.not, label %if.end, label %if.then
+
+if.then:
+  %arrayidx = getelementptr inbounds [0 x i8], [0 x i8]* @a, i32 0, i32 %add
+  %1 = load i8, i8* %arrayidx, align 1
+  tail call void asm sideeffect "", "r"(i8 %1)
+  store i32 0, i32* @first, align 4
+  br label %if.end
+
+if.end:
+  tail call void asm sideeffect "", "r"(i32 %add)
+  %inc = add nuw nsw i32 %j.016, 1
+  %cmp2 = icmp ult i32 %j.016, 99999
+  br i1 %cmp2, label %for.body4, label %for.cond.cleanup3
+}
+
 declare void @objc_enumerationMutation(i8*)
 declare dso_local void @f(i32*)
 declare dso_local void @g(...)