[LoopPeel] Peel if it turns invariant loads dereferenceable.

fhahn · fhahn · commit 405eae583837 · 2021-10-13T16:50:38.000+01:00
This patch adds a new cost heuristic that allows peeling a single iteration off read-only loops, if the loop contains a load that 1. is feeding an exit condition, 2. dominates the latch, 3. is not already known to be dereferenceable, 4. and has a loop invariant address. If all non-latch exits are terminated with unreachable, such loads in the loop are guaranteed to be dereferenceable after peeling, enabling hoisting/CSE'ing them. This enables vectorization of loops with certain runtime-checks, like multiple calls to `std::vector::at` if the vector is passed as pointer. Reviewed By: mkazantsev Differential Revision: https://reviews.llvm.org/D108114 (cherry-picked from cd0ba9dc58c5806f4e3cc9635ab1f64af6973a83)
diff --git a/llvm/include/llvm/Transforms/Utils/LoopPeel.h b/llvm/include/llvm/Transforms/Utils/LoopPeel.h
@@ -32,8 +32,8 @@ gatherPeelingPreferences(Loop *L, ScalarEvolution &SE,
 
 void computePeelCount(Loop *L, unsigned LoopSize,
                       TargetTransformInfo::PeelingPreferences &PP,
-                      unsigned &TripCount, ScalarEvolution &SE,
-                      unsigned Threshold = UINT_MAX);
+                      unsigned &TripCount, DominatorTree &DT,
+                      ScalarEvolution &SE, unsigned Threshold = UINT_MAX);
 
 } // end namespace llvm
 
diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -877,7 +877,7 @@ bool llvm::computeUnrollCount(
   }
 
   // 4th priority is loop peeling.
-  computePeelCount(L, LoopSize, PP, TripCount, SE, UP.Threshold);
+  computePeelCount(L, LoopSize, PP, TripCount, DT, SE, UP.Threshold);
   if (PP.PeelCount) {
     UP.Runtime = false;
     UP.Count = 1;
diff --git a/llvm/lib/Transforms/Utils/LoopPeel.cpp b/llvm/lib/Transforms/Utils/LoopPeel.cpp
@@ -14,6 +14,7 @@
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopIterator.h"
 #include "llvm/Analysis/ScalarEvolution.h"
@@ -165,6 +166,66 @@ static unsigned calculateIterationsToInvariance(
   return ToInvariance;
 }
 
+// Try to find any invariant memory reads that will become dereferenceable in
+// the remainder loop after peeling. The load must also be used (transitively)
+// by an exit condition. Returns the number of iterations to peel off (at the
+// moment either 0 or 1).
+static unsigned peelToTurnInvariantLoadsDerefencebale(Loop &L,
+                                                      DominatorTree &DT) {
+  // Skip loops with a single exiting block, because there should be no benefit
+  // for the heuristic below.
+  if (L.getExitingBlock())
+    return 0;
+
+  // All non-latch exit blocks must have an UnreachableInst terminator.
+  // Otherwise the heuristic below may not be profitable.
+  SmallVector<BasicBlock *, 4> Exits;
+  L.getUniqueNonLatchExitBlocks(Exits);
+  if (any_of(Exits, [](const BasicBlock *BB) {
+        return !isa<UnreachableInst>(BB->getTerminator());
+      }))
+    return 0;
+
+  // Now look for invariant loads that dominate the latch and are not known to
+  // be dereferenceable. If there are such loads and no writes, they will become
+  // dereferenceable in the loop if the first iteration is peeled off. Also
+  // collect the set of instructions controlled by such loads. Only peel if an
+  // exit condition uses (transitively) such a load.
+  BasicBlock *Header = L.getHeader();
+  BasicBlock *Latch = L.getLoopLatch();
+  SmallPtrSet<Value *, 8> LoadUsers;
+  const DataLayout &DL = L.getHeader()->getModule()->getDataLayout();
+  for (BasicBlock *BB : L.blocks()) {
+    for (Instruction &I : *BB) {
+      if (I.mayWriteToMemory())
+        return 0;
+
+      auto Iter = LoadUsers.find(&I);
+      if (Iter != LoadUsers.end()) {
+        for (Value *U : I.users())
+          LoadUsers.insert(U);
+      }
+      // Do not look for reads in the header; they can already be hoisted
+      // without peeling.
+      if (BB == Header)
+        continue;
+      if (auto *LI = dyn_cast<LoadInst>(&I)) {
+        Value *Ptr = LI->getPointerOperand();
+        if (DT.dominates(BB, Latch) && L.isLoopInvariant(Ptr) &&
+            !isDereferenceablePointer(Ptr, LI->getType(), DL, LI, &DT))
+          for (Value *U : I.users())
+            LoadUsers.insert(U);
+      }
+    }
+  }
+  SmallVector<BasicBlock *> ExitingBlocks;
+  L.getExitingBlocks(ExitingBlocks);
+  for (BasicBlock *Exiting : ExitingBlocks)
+    if (LoadUsers.find(Exiting->getTerminator()) != LoadUsers.end())
+      return 1;
+  return 0;
+}
+
 // Return the number of iterations to peel off that make conditions in the
 // body true/false. For example, if we peel 2 iterations off the loop below,
 // the condition i < 2 can be evaluated at compile time.
@@ -280,8 +341,8 @@ static unsigned countToEliminateCompares(Loop &L, unsigned MaxPeelCount,
 // Return the number of iterations we want to peel off.
 void llvm::computePeelCount(Loop *L, unsigned LoopSize,
                             TargetTransformInfo::PeelingPreferences &PP,
-                            unsigned &TripCount, ScalarEvolution &SE,
-                            unsigned Threshold) {
+                            unsigned &TripCount, DominatorTree &DT,
+                            ScalarEvolution &SE, unsigned Threshold) {
   assert(LoopSize > 0 && "Zero loop size is not allowed!");
   // Save the PP.PeelCount value set by the target in
   // TTI.getPeelingPreferences or by the flag -unroll-peel-count.
@@ -348,6 +409,9 @@ void llvm::computePeelCount(Loop *L, unsigned LoopSize,
     DesiredPeelCount = std::max(DesiredPeelCount,
                                 countToEliminateCompares(*L, MaxPeelCount, SE));
 
+    if (DesiredPeelCount == 0)
+      DesiredPeelCount = peelToTurnInvariantLoadsDerefencebale(*L, DT);
+
     if (DesiredPeelCount > 0) {
       DesiredPeelCount = std::min(DesiredPeelCount, MaxPeelCount);
       // Consider max peel count limitation.
diff --git a/llvm/test/Transforms/LoopUnroll/peel-to-turn-invariant-accesses-dereferenceable.ll b/llvm/test/Transforms/LoopUnroll/peel-to-turn-invariant-accesses-dereferenceable.ll
@@ -6,15 +6,99 @@ declare void @foo()
 define i32 @peel_readonly_to_make_loads_derefenceable(i32* %ptr, i32 %N, i32* %inv, i1 %c.1) {
 ; CHECK-LABEL: @peel_readonly_to_make_loads_derefenceable(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP_HEADER_PEEL_BEGIN:%.*]]
+; CHECK:       loop.header.peel.begin:
+; CHECK-NEXT:    br label [[LOOP_HEADER_PEEL:%.*]]
+; CHECK:       loop.header.peel:
+; CHECK-NEXT:    br i1 [[C_1:%.*]], label [[THEN_PEEL:%.*]], label [[UNREACHABLE_EXIT:%.*]]
+; CHECK:       then.peel:
+; CHECK-NEXT:    [[I_PEEL:%.*]] = load i32, i32* [[INV:%.*]], align 4
+; CHECK-NEXT:    [[C_2_PEEL:%.*]] = icmp ult i32 [[I_PEEL]], 2
+; CHECK-NEXT:    br i1 [[C_2_PEEL]], label [[LOOP_LATCH_PEEL:%.*]], label [[UNREACHABLE_EXIT]]
+; CHECK:       loop.latch.peel:
+; CHECK-NEXT:    [[GEP_PEEL:%.*]] = getelementptr i32, i32* [[PTR:%.*]], i32 1
+; CHECK-NEXT:    [[LV_PEEL:%.*]] = load i32, i32* [[GEP_PEEL]], align 4
+; CHECK-NEXT:    [[SUM_NEXT_PEEL:%.*]] = add i32 0, [[LV_PEEL]]
+; CHECK-NEXT:    [[IV_NEXT_PEEL:%.*]] = add nuw nsw i32 1, 1
+; CHECK-NEXT:    [[C_3_PEEL:%.*]] = icmp ult i32 1, 1000
+; CHECK-NEXT:    br i1 [[C_3_PEEL]], label [[LOOP_HEADER_PEEL_NEXT:%.*]], label [[EXIT:%.*]]
+; CHECK:       loop.header.peel.next:
+; CHECK-NEXT:    br label [[LOOP_HEADER_PEEL_NEXT1:%.*]]
+; CHECK:       loop.header.peel.next1:
+; CHECK-NEXT:    br label [[ENTRY_PEEL_NEWPH:%.*]]
+; CHECK:       entry.peel.newph:
+; CHECK-NEXT:    br label [[LOOP_HEADER:%.*]]
+; CHECK:       loop.header:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[IV_NEXT_PEEL]], [[ENTRY_PEEL_NEWPH]] ], [ [[IV_NEXT:%.*]], [[LOOP_LATCH:%.*]] ]
+; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ [[SUM_NEXT_PEEL]], [[ENTRY_PEEL_NEWPH]] ], [ [[SUM_NEXT:%.*]], [[LOOP_LATCH]] ]
+; CHECK-NEXT:    br i1 [[C_1]], label [[THEN:%.*]], label [[UNREACHABLE_EXIT_LOOPEXIT:%.*]]
+; CHECK:       then:
+; CHECK-NEXT:    [[I:%.*]] = load i32, i32* [[INV]], align 4
+; CHECK-NEXT:    [[C_2:%.*]] = icmp ult i32 [[I]], 2
+; CHECK-NEXT:    br i1 [[C_2]], label [[LOOP_LATCH]], label [[UNREACHABLE_EXIT_LOOPEXIT]]
+; CHECK:       loop.latch:
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, i32* [[PTR]], i32 [[IV]]
+; CHECK-NEXT:    [[LV:%.*]] = load i32, i32* [[GEP]], align 4
+; CHECK-NEXT:    [[SUM_NEXT]] = add i32 [[SUM]], [[LV]]
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 1
+; CHECK-NEXT:    [[C_3:%.*]] = icmp ult i32 [[IV]], 1000
+; CHECK-NEXT:    br i1 [[C_3]], label [[LOOP_HEADER]], label [[EXIT_LOOPEXIT:%.*]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK:       exit.loopexit:
+; CHECK-NEXT:    [[SUM_NEXT_LCSSA_PH:%.*]] = phi i32 [ [[SUM_NEXT]], [[LOOP_LATCH]] ]
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT_PEEL]], [[LOOP_LATCH_PEEL]] ], [ [[SUM_NEXT_LCSSA_PH]], [[EXIT_LOOPEXIT]] ]
+; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
+; CHECK:       unreachable.exit.loopexit:
+; CHECK-NEXT:    br label [[UNREACHABLE_EXIT]]
+; CHECK:       unreachable.exit:
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    unreachable
+;
+entry:
+  br label %loop.header
+
+loop.header:
+  %iv = phi i32 [ 1, %entry ], [ %iv.next, %loop.latch ]
+  %sum = phi i32 [ 0, %entry ], [ %sum.next, %loop.latch ]
+  br i1 %c.1, label %then, label %unreachable.exit
+
+then:
+  %i = load i32, i32* %inv
+  %c.2 = icmp ult i32 %i, 2
+  br i1 %c.2, label %loop.latch, label %unreachable.exit
+
+loop.latch:
+  %gep = getelementptr i32, i32* %ptr, i32 %iv
+  %lv = load i32, i32* %gep
+  %sum.next = add i32 %sum, %lv
+  %iv.next = add nuw nsw i32  %iv, 1
+  %c.3 = icmp ult i32 %iv, 1000
+  br i1 %c.3, label %loop.header, label %exit
+
+exit:
+  ret i32 %sum.next
+
+unreachable.exit:
+  call void @foo()
+  unreachable
+}
+
+define i32 @peel_readonly_to_make_loads_derefenceable_exits_lead_to_unreachable(i32* %ptr, i32 %N, i32* %inv, i1 %c.1) {
+; CHECK-LABEL: @peel_readonly_to_make_loads_derefenceable_exits_lead_to_unreachable(
+; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP_HEADER:%.*]]
 ; CHECK:       loop.header:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 1, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP_LATCH:%.*]] ]
 ; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[SUM_NEXT:%.*]], [[LOOP_LATCH]] ]
-; CHECK-NEXT:    br i1 [[C_1:%.*]], label [[THEN:%.*]], label [[UNREACHABLE_EXIT:%.*]]
+; CHECK-NEXT:    br i1 [[C_1:%.*]], label [[THEN:%.*]], label [[EXIT_2:%.*]]
 ; CHECK:       then:
 ; CHECK-NEXT:    [[I:%.*]] = load i32, i32* [[INV:%.*]], align 4
 ; CHECK-NEXT:    [[C_2:%.*]] = icmp ult i32 [[I]], 2
-; CHECK-NEXT:    br i1 [[C_2]], label [[LOOP_LATCH]], label [[UNREACHABLE_EXIT]]
+; CHECK-NEXT:    br i1 [[C_2]], label [[THEN_2:%.*]], label [[EXIT_2]]
+; CHECK:       then.2:
+; CHECK-NEXT:    [[C_4:%.*]] = icmp ult i32 [[I]], 4
+; CHECK-NEXT:    br i1 [[C_4]], label [[LOOP_LATCH]], label [[EXIT_3:%.*]]
 ; CHECK:       loop.latch:
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, i32* [[PTR:%.*]], i32 [[IV]]
 ; CHECK-NEXT:    [[LV:%.*]] = load i32, i32* [[GEP]], align 4
@@ -25,7 +109,11 @@ define i32 @peel_readonly_to_make_loads_derefenceable(i32* %ptr, i32 %N, i32* %i
 ; CHECK:       exit:
 ; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT]], [[LOOP_LATCH]] ]
 ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
-; CHECK:       unreachable.exit:
+; CHECK:       exit.2:
+; CHECK-NEXT:    br label [[UNREACHABLE_BB:%.*]]
+; CHECK:       exit.3:
+; CHECK-NEXT:    br label [[UNREACHABLE_BB]]
+; CHECK:       unreachable.bb:
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    unreachable
 ;
@@ -35,12 +123,16 @@ entry:
 loop.header:
   %iv = phi i32 [ 1, %entry ], [ %iv.next, %loop.latch ]
   %sum = phi i32 [ 0, %entry ], [ %sum.next, %loop.latch ]
-  br i1 %c.1, label %then, label %unreachable.exit
+  br i1 %c.1, label %then, label %exit.2
 
 then:
   %i = load i32, i32* %inv
   %c.2 = icmp ult i32 %i, 2
-  br i1 %c.2, label %loop.latch, label %unreachable.exit
+  br i1 %c.2, label %then.2, label %exit.2
+
+then.2:
+  %c.4 = icmp ult i32 %i, 4
+  br i1 %c.4, label %loop.latch, label %exit.3
 
 loop.latch:
   %gep = getelementptr i32, i32* %ptr, i32 %iv
@@ -53,7 +145,13 @@ loop.latch:
 exit:
   ret i32 %sum.next
 
-unreachable.exit:
+exit.2:
+  br label %unreachable.bb
+
+exit.3:
+  br label %unreachable.bb
+
+unreachable.bb:
   call void @foo()
   unreachable
 }
@@ -302,18 +400,18 @@ unreachable.exit:
 
 declare i32 @llvm.experimental.deoptimize.i32(...)
 
-define i32 @do_not_peel_with_deopt_exit(i32* %ptr, i32 %N, i32* %inv, i1 %c.1) {
-; CHECK-LABEL: @do_not_peel_with_deopt_exit(
+define i32 @peel_with_deopt_exit(i32* %ptr, i32 %N, i32* %inv, i1 %c.1) {
+; CHECK-LABEL: @peel_with_deopt_exit(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP_HEADER:%.*]]
 ; CHECK:       loop.header:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 1, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP_LATCH:%.*]] ]
 ; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[SUM_NEXT:%.*]], [[LOOP_LATCH]] ]
-; CHECK-NEXT:    br i1 [[C_1:%.*]], label [[THEN:%.*]], label [[UNREACHABLE_EXIT:%.*]]
+; CHECK-NEXT:    br i1 [[C_1:%.*]], label [[THEN:%.*]], label [[DEOPT_EXIT:%.*]]
 ; CHECK:       then:
 ; CHECK-NEXT:    [[I:%.*]] = load i32, i32* [[INV:%.*]], align 4
 ; CHECK-NEXT:    [[C_2:%.*]] = icmp ult i32 [[I]], 2
-; CHECK-NEXT:    br i1 [[C_2]], label [[LOOP_LATCH]], label [[UNREACHABLE_EXIT]]
+; CHECK-NEXT:    br i1 [[C_2]], label [[LOOP_LATCH]], label [[DEOPT_EXIT]]
 ; CHECK:       loop.latch:
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, i32* [[PTR:%.*]], i32 [[IV]]
 ; CHECK-NEXT:    [[LV:%.*]] = load i32, i32* [[GEP]], align 4
@@ -324,7 +422,7 @@ define i32 @do_not_peel_with_deopt_exit(i32* %ptr, i32 %N, i32* %inv, i1 %c.1) {
 ; CHECK:       exit:
 ; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT]], [[LOOP_LATCH]] ]
 ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
-; CHECK:       unreachable.exit:
+; CHECK:       deopt.exit:
 ; CHECK-NEXT:    [[SUM_LCSSA:%.*]] = phi i32 [ [[SUM]], [[THEN]] ], [ [[SUM]], [[LOOP_HEADER]] ]
 ; CHECK-NEXT:    [[RVAL:%.*]] = call i32 (...) @llvm.experimental.deoptimize.i32() [ "deopt"(i32 [[SUM_LCSSA]]) ]
 ; CHECK-NEXT:    ret i32 [[RVAL]]
@@ -335,12 +433,12 @@ entry:
 loop.header:
   %iv = phi i32 [ 1, %entry ], [ %iv.next, %loop.latch ]
   %sum = phi i32 [ 0, %entry ], [ %sum.next, %loop.latch ]
-  br i1 %c.1, label %then, label %unreachable.exit
+  br i1 %c.1, label %then, label %deopt.exit
 
 then:
   %i = load i32, i32* %inv
   %c.2 = icmp ult i32 %i, 2
-  br i1 %c.2, label %loop.latch, label %unreachable.exit
+  br i1 %c.2, label %loop.latch, label %deopt.exit
 
 loop.latch:
   %gep = getelementptr i32, i32* %ptr, i32 %iv
@@ -353,7 +451,7 @@ loop.latch:
 exit:
   ret i32 %sum.next
 
-unreachable.exit:
+deopt.exit:
   %rval = call i32(...) @llvm.experimental.deoptimize.i32() [ "deopt"(i32 %sum) ]
   ret i32 %rval
 }
diff --git a/llvm/test/Transforms/PhaseOrdering/AArch64/peel-multiple-unreachable-exits-for-vectorization.ll b/llvm/test/Transforms/PhaseOrdering/AArch64/peel-multiple-unreachable-exits-for-vectorization.ll

Original file line number	Diff line number	Diff line change
`@@ -877,7 +877,7 @@ bool llvm::computeUnrollCount(`
`877`	`877`	`}`
`878`	`878`
`879`	`879`	`// 4th priority is loop peeling.`
`880`		`- computePeelCount(L, LoopSize, PP, TripCount, SE, UP.Threshold);`
	`880`	`+ computePeelCount(L, LoopSize, PP, TripCount, DT, SE, UP.Threshold);`
`881`	`881`	`if (PP.PeelCount) {`
`882`	`882`	`UP.Runtime = false;`
`883`	`883`	`UP.Count = 1;`