[llvm] ecb4171 - [LV] Handle zero cost loops in selectInterleaveCount.

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2022-03-29T22:52:43+01:00
New Revision: ecb4171dcbf1b433c9963fd605a74898303e850d

URL: https://github.com/llvm/llvm-project/commit/ecb4171dcbf1b433c9963fd605a74898303e850d
DIFF: https://github.com/llvm/llvm-project/commit/ecb4171dcbf1b433c9963fd605a74898303e850d.diff

LOG: [LV] Handle zero cost loops in selectInterleaveCount.

In some case, like in the added test case, we can reach
selectInterleaveCount with loops that actually have a cost of 0.

Unfortunately a loop cost of 0 is also used to communicate that the cost
has not been computed yet. To resolve the crash, bail out if the cost
remains zero after computing it.

This seems like the best option, as there are multiple code paths that
return a cost of 0 to force a computation in selectInterleaveCount.
Computing the cost at multiple places up front there would unnecessarily
complicate the logic.

Fixes #54413.

Added: 
    llvm/test/Transforms/LoopVectorize/X86/pr54413-select-interleave-count-loop-with-cost-zero.ll

Modified: 
    llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index efebf8efdec38..911d36e5c028e 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -5695,6 +5695,18 @@ unsigned LoopVectorizationCostModel::selectInterleaveCount(ElementCount VF,
       !(InterleaveSmallLoopScalarReduction && HasReductions && VF.isScalar()))
     return 1;
 
+  // If we did not calculate the cost for VF (because the user selected the VF)
+  // then we calculate the cost of VF here.
+  if (LoopCost == 0) {
+    InstructionCost C = expectedCost(VF).first;
+    assert(C.isValid() && "Expected to have chosen a VF with valid cost");
+    LoopCost = *C.getValue();
+
+    // Loop body is free and there is no need for interleaving.
+    if (LoopCost == 0)
+      return 1;
+  }
+
   RegisterUsage R = calculateRegisterUsage({VF})[0];
   // We divide by these constants so assume that we have at least one
   // instruction that uses at least one register.
@@ -5786,16 +5798,6 @@ unsigned LoopVectorizationCostModel::selectInterleaveCount(ElementCount VF,
 
   assert(IC > 0 && "Interleave count must be greater than 0.");
 
-  // If we did not calculate the cost for VF (because the user selected the VF)
-  // then we calculate the cost of VF here.
-  if (LoopCost == 0) {
-    InstructionCost C = expectedCost(VF).first;
-    assert(C.isValid() && "Expected to have chosen a VF with valid cost");
-    LoopCost = *C.getValue();
-  }
-
-  assert(LoopCost && "Non-zero loop cost expected");
-
   // Interleave if we vectorized this loop and there is a reduction that could
   // benefit from interleaving.
   if (VF.isVector() && HasReductions) {

diff  --git a/llvm/test/Transforms/LoopVectorize/X86/pr54413-select-interleave-count-loop-with-cost-zero.ll b/llvm/test/Transforms/LoopVectorize/X86/pr54413-select-interleave-count-loop-with-cost-zero.ll
new file mode 100644
index 0000000000000..3b7fa43eea8e7
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/X86/pr54413-select-interleave-count-loop-with-cost-zero.ll
@@ -0,0 +1,50 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes='loop(indvars),loop-vectorize' -S %s | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+
+; After indvars, the backedge taken count for %loop2 becomes 1, but SCEV
+; retains the cached original BTC, as the loop is in dead code. Make sure
+; LV does not crash when trying to select an interleave count for a loop with zero cost.
+define void @pr54413(i64* %ptr.base) {
+; CHECK-LABEL: @pr54413(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP1:%.*]]
+; CHECK:       loop1:
+; CHECK-NEXT:    br i1 true, label [[LOOP1_LATCH:%.*]], label [[LOOP2_PREHEADER:%.*]]
+; CHECK:       loop2.preheader:
+; CHECK-NEXT:    br label [[LOOP2:%.*]]
+; CHECK:       loop2:
+; CHECK-NEXT:    [[PTR:%.*]] = phi i64* [ [[PTR_BASE:%.*]], [[LOOP2_PREHEADER]] ], [ [[PTR_NEXT:%.*]], [[LOOP2]] ]
+; CHECK-NEXT:    [[PTR_NEXT]] = getelementptr inbounds i64, i64* [[PTR]], i64 1
+; CHECK-NEXT:    br i1 true, label [[LOOP2_EXIT:%.*]], label [[LOOP2]]
+; CHECK:       loop2.exit:
+; CHECK-NEXT:    [[PTR_NEXT_LCSSA:%.*]] = phi i64* [ [[PTR_NEXT]], [[LOOP2]] ]
+; CHECK-NEXT:    br label [[LOOP1_LATCH]]
+; CHECK:       loop1.latch:
+; CHECK-NEXT:    br label [[LOOP1]]
+;
+entry:
+  br label %loop1
+
+loop1:
+  br i1 true, label %loop1.latch, label %loop2.preheader
+
+loop2.preheader:
+  br label %loop2
+
+loop2:
+  %iv = phi i64 [ 0, %loop2.preheader ], [ %iv.next, %loop2 ]
+  %ptr = phi i64* [ %ptr.base, %loop2.preheader ], [ %ptr.next, %loop2 ]
+  %iv.next = add nuw nsw i64 %iv, 1
+  %ptr.next = getelementptr inbounds i64, i64* %ptr, i64 1
+  %cmp = icmp eq i64 %iv, 1024
+  br i1 %cmp, label %loop2.exit, label %loop2
+
+loop2.exit:
+  %ptr.next.lcssa = phi i64* [ %ptr.next, %loop2 ]
+  br label %loop1.latch
+
+loop1.latch:
+  br label %loop1
+}