[LoopTermFold] Allow pointer IVs with dereferenceable accesses to fold terminating condition

llvm/lib/Transforms/Scalar/LoopTermFold.cpp

@@ -21,7 +21,9 @@
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
@@ -48,6 +50,47 @@ using namespace llvm;
 STATISTIC(NumTermFold,
           "Number of terminating condition fold recognized and performed");
 
+// An aligned pointer with a power-of-two stride must visit null before it
+// repeats. If it is dereferenced on every iteration, and the alignment is at
+// least the stride, those iterations cannot all execute in a defined program.
+static bool hasNoSelfWrapViaDereference(PHINode &PN,
+                                        const SCEVAddRecExpr *AddRec,
+                                        ScalarEvolution &SE, DominatorTree &DT,
+                                        Instruction *LatchTerm) {
+  auto *PtrTy = dyn_cast<PointerType>(PN.getType());
+  if (!PtrTy || AddRec->getLoop()->getHeader() != PN.getParent() ||
+      NullPointerIsDefined(PN.getFunction(), PtrTy->getAddressSpace()))
+    return false;
+
+  const auto *Step = dyn_cast<SCEVConstant>(AddRec->getStepRecurrence(SE));
+  if (!Step)
+    return false;
+  APInt Stride = Step->getAPInt().abs();
+  if (!Stride.isPowerOf2() || Stride.getActiveBits() > 64)
+    return false;
+
+  const DataLayout &DL = PN.getDataLayout();
+  // GEP arithmetic must cover the entire pointer representation for its cycle
+  // to include null.
+  if (DL.isNonIntegralPointerType(PtrTy) ||
+      DL.getIndexSizeInBits(PtrTy->getAddressSpace()) !=
+          DL.getPointerSizeInBits(PtrTy->getAddressSpace()))
+    return false;
+
+  for (User *U : PN.users()) {
+    if (getLoadStorePointerOperand(U) != &PN)
+      continue;
+    // Volatile accesses may legally use address zero (for example, MMIO).
+    if (cast<Instruction>(U)->isVolatile())
+      continue;
+    if (getLoadStoreAlignment(U).value() >= Stride.getZExtValue() &&
+        DL.getTypeStoreSize(getLoadStoreType(U)).getKnownMinValue() != 0 &&
+        DT.dominates(cast<Instruction>(U), LatchTerm))
+      return true;
+  }
+  return false;
+}
+
 static std::optional<std::tuple<PHINode *, PHINode *, const SCEV *, bool>>
 canFoldTermCondOfLoop(Loop *L, ScalarEvolution &SE, DominatorTree &DT,
                       const LoopInfo &LI, const TargetTransformInfo &TTI) {
@@ -152,7 +195,9 @@ canFoldTermCondOfLoop(Loop *L, ScalarEvolution &SE, DominatorTree &DT,
     // iteration. The simplest case to consider is a candidate IV which is
     // narrower than the trip count (and thus original IV), but this can
     // also happen due to non-unit strides on the candidate IVs.
-    if (!AddRec->hasNoSelfWrap() ||
+    if ((!AddRec->hasNoSelfWrap() &&
+         !hasNoSelfWrapViaDereference(PN, AddRec, SE, DT,
+                                      LoopLatch->getTerminator())) ||
         !SE.isKnownNonZero(AddRec->getStepRecurrence(SE)))
       continue;
 

llvm/test/Transforms/LoopStrengthReduce/lsr-term-fold.ll

@@ -230,20 +230,20 @@ for.end:                                          ; preds = %for.body
   ret void
 }
 
-; In this case, the integer IV has a larger bitwidth than the pointer IV.
-; This means that the smaller IV may wrap around multiple times before
-; the original loop exit is taken.
+; The integer IV is wider than the pointer IV. However, the aligned store
+; would dereference null before the pointer repeats, so it is safe to fold.
 define void @iv_size(ptr %a, i128 %N) {
 ; CHECK-LABEL: @iv_size(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = trunc i128 [[N:%.*]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i64 [[TMP0]], 2
+; CHECK-NEXT:    [[SCEVGEP:%.*]] = getelementptr i8, ptr [[A:%.*]], i64 [[TMP1]]
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
-; CHECK-NEXT:    [[LSR_IV1:%.*]] = phi ptr [ [[UGLYGEP2:%.*]], [[FOR_BODY]] ], [ [[A:%.*]], [[ENTRY:%.*]] ]
-; CHECK-NEXT:    [[LSR_IV:%.*]] = phi i128 [ [[LSR_IV_NEXT:%.*]], [[FOR_BODY]] ], [ [[N:%.*]], [[ENTRY]] ]
+; CHECK-NEXT:    [[LSR_IV1:%.*]] = phi ptr [ [[UGLYGEP2:%.*]], [[FOR_BODY]] ], [ [[A]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    store i32 1, ptr [[LSR_IV1]], align 4
-; CHECK-NEXT:    [[LSR_IV_NEXT]] = add nsw i128 [[LSR_IV]], -1
 ; CHECK-NEXT:    [[UGLYGEP2]] = getelementptr i8, ptr [[LSR_IV1]], i64 4
-; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i128 [[LSR_IV_NEXT]], 0
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq ptr [[UGLYGEP2]], [[SCEVGEP]]
 ; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_END:%.*]], label [[FOR_BODY]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void

llvm/test/Transforms/LoopStrengthReduce/term-fold-dereference-narrow-index.ll

@@ -0,0 +1,34 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=loop-term-fold -S < %s | FileCheck %s
+
+; The unchanged high pointer bits can prevent the recurrence from visiting null.
+target datalayout = "e-p:64:64:64:32"
+
+define void @narrow_pointer_index(ptr %a, i32 %n) {
+; CHECK-LABEL: @narrow_pointer_index(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[P:%.*]] = phi ptr [ [[A:%.*]], [[ENTRY:%.*]] ], [ [[P_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[COUNT:%.*]] = phi i32 [ [[N:%.*]], [[ENTRY]] ], [ [[COUNT_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    store i8 1, ptr [[P]], align 2
+; CHECK-NEXT:    [[P_NEXT]] = getelementptr i8, ptr [[P]], i32 2
+; CHECK-NEXT:    [[COUNT_NEXT]] = add i32 [[COUNT]], -1
+; CHECK-NEXT:    [[DONE:%.*]] = icmp eq i32 [[COUNT_NEXT]], 0
+; CHECK-NEXT:    br i1 [[DONE]], label [[EXIT:%.*]], label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+loop:
+  %p = phi ptr [ %a, %entry ], [ %p.next, %loop ]
+  %count = phi i32 [ %n, %entry ], [ %count.next, %loop ]
+  store i8 1, ptr %p, align 2
+  %p.next = getelementptr i8, ptr %p, i32 2
+  %count.next = add i32 %count, -1
+  %done = icmp eq i32 %count.next, 0
+  br i1 %done, label %exit, label %loop
+exit:
+  ret void
+}