[LoopUtils] Better accuracy for getLoopEstimatedTripCount.

[ebrevnov]

Summary

Current implementation of getLoopEstimatedTripCount returns 1 iteration less than it should. The reason is that in bottom tested loop first iteration is executed before first back branch is taken. For example for loop with !{!"branch_weights", i32 1 taken, i32 1 exit} metadata getLoopEstimatedTripCount gives 1 while actual number of iterations is 2.

[Ayal]

llvm/lib/Transforms/Utils/LoopUtils.cpp

719

Would be better to separate the functional (+1) change from the variable name changes and divideNearest(). The latter is an NFC which is independent of and hides the former.

If the functional change goes first, it can e.g. simply do

+  uint64_t EstimatedBackedgeTakenCount;
  if (LatchBR->getSuccessor(0) == L->getHeader())
-    return (TrueVal + (FalseVal / 2)) / FalseVal;
+    EstimatedBackedgeTakenCount = (TrueVal + (FalseVal / 2)) / FalseVal;
  else
-    return (FalseVal + (TrueVal / 2)) / TrueVal;
+    EstimatedBackedgeTakenCount = (FalseVal + (TrueVal / 2)) / TrueVal;
+  return EstimatedBackedgeTakenCount + 1;

Can alternatively renamed original function to getEstimatedBackedgeTakenCount(), and introduced a new getLoopEstimatedTripCount() which returns the former plus one, if it's useful(?)

llvm/unittests/Transforms/Utils/CMakeLists.txt

18

already there?

llvm/unittests/Transforms/Utils/LoopUtilsTest.cpp

1

"This file was added." - this file already exists, right?
What is this trying to check on top of the above fixed lit tests?

[ebrevnov]

Rebase.

[Ayal]

Just pointing out for the record that this is extracted from D67905

[ebrevnov]

Factor out NFC stuff + rebase

[Ayal]

llvm/lib/Transforms/Utils/LoopUtils.cpp

714

If !LatchCycleWeight (if !BackedgeTakenWeight) it should return 1 instead of 0. I.e., should always be equal to "getEstimateBackedgeTakenCount()" plus 1.

If !LatchExitWeight, perhaps better return None instead of 0? This implies the loop was never entered when gathering the profile, or loop was entered once and never exited(?); in either case estimating the trip count to be zero seems wrong.

[ebrevnov]

llvm/lib/Transforms/Utils/LoopUtils.cpp

714

Completely agree.

[ebrevnov]

Fixed raised issue + Rebase.

[ebrevnov]

ping @Ayal

[Ayal]

llvm/lib/Transforms/Utils/LoopUtils.cpp

717

Maybe clearer to emphasize the +1 effect, noting that Trip count is one plus the estimated backedge taken count, and doing

uint64_t BackedgeTakenCount = llvm::divideNearest(BackedgeTakenWeight, LatchExitWeight);
return BackedgeTakenCount + 1;

(cf. InnerLoopVectorizer::getOrCreateTripCount(Loop *L))

llvm/test/Transforms/LoopUnroll/peel-loop-conditions-pgo-1.ll

14

It is indeed confusing at first sight, to see that the total number of estimated trip count is 2+5, given branch_weights of 6 and 1 (and thus estimated backedge taken count of 6/1). Perhaps worth clarifying in the above comment about "remained iterations". Should -unroll-peel-max-count=7 be updated?

Note that loop peeling already updates the estimated trip count (from 7 to 5 above), by calling Transforms/Utils/LoopUnrollPeel.cpp's fixupBranchWeights(); potential reuse with D67905's setLoopEstimatedTripCount().

llvm/unittests/Transforms/Utils/LoopUtilsTest.cpp

1

What is this trying to check on top of the above fixed lit tests?

This change seems mostly unrelated to the +1 accuracy fix, right? Suggest to commit it first, testing the original (backedge taken) value, and then modify it along with the +1 patch, similar to above fixes of lit tests.

18

These should all be in lex order.

[ebrevnov]

llvm/lib/Transforms/Utils/LoopUtils.cpp

717

Sure, works for me. In fact, this +1 approach was used in original patch :-)

llvm/test/Transforms/LoopUnroll/peel-loop-conditions-pgo-1.ll

14

Should -unroll-peel-max-count=7 be updated?

Not sure I understand why it should be updated. Loop has 7 estimated iterations and we aim at peeling 2 iterations. Thus -unroll-peel-max-count=7 seems reasonable to me.

14

It is indeed confusing at first sight, to see that the total number of estimated trip count is 2+5, given branch_weights of 6 and 1 (and thus estimated backedge taken count of 6/1).

Even if I clarify here it won't help a lot since there are bunch of other similar places. Probably it's better to add explanation to getLoopEstimatedTripCount?

llvm/unittests/Transforms/Utils/LoopUtilsTest.cpp

1

What is this trying to check on top of the above fixed lit tests?

I believe you requested to add such test at some time. Anyway it's better to have such test than not :-). Will extract this change to a separate review.

18

Will fix

[ebrevnov]

Moved regression test to a separate review.

[Ayal]

llvm/lib/Transforms/Utils/LoopUtils.cpp

715–719

"Trip count" >> "Estimated backedge taken count"

720

Suffice to say that Estimated trip count is one plus estimated backedge taken count.

llvm/test/Transforms/LoopUnroll/peel-loop-conditions-pgo-1.ll

14

OK; just noting that original test estimated the trip count to be 6, and used -unroll-peel-max-count=7, so seemed reasonable to bump this threshold (to 8) given that the estimation bumped (to 7).

14

OK; just suggesting to extend the existing "remained iterations" comment above, clarifying how many they are.

llvm/unittests/Transforms/Utils/LoopUtilsTest.cpp

1

Sorry, my original request was just to make sure the +1 change is covered by tests, as opposed to NFC patches; the above lit tests seem fine.

[ebrevnov]

Fixed comments.

[Ayal]

This look good to me, with a last minor nit, thanks!

llvm/lib/Transforms/Utils/LoopUtils.cpp

716

"the the edge exiting weight" >> "the weight of the edge exiting the loop"

llvm/include/llvm/Support/MathExtras.h

template <uint64_t Align> constexpr inline uint64_t alignTo(uint64_t Value) {
   return (Value + Align - 1) / Align * Align;
 }
 
 /// Returns the integer ceil(Numerator / Denominator).
 inline uint64_t divideCeil(uint64_t Numerator, uint64_t Denominator) {
   return alignTo(Numerator, Denominator) / Denominator;
 }
 
+/// Returns the integer nearest(Numerator / Denominator).
+inline uint64_t divideNearest(uint64_t Numerator, uint64_t Denominator) {
+  return (Numerator + (Denominator / 2)) / Denominator;
+}
+
 /// Returns the largest uint64_t less than or equal to \p Value and is
 /// \p Skew mod \p Align. \p Align must be non-zero
 inline uint64_t alignDown(uint64_t Value, uint64_t Align, uint64_t Skew = 0) {
   assert(Align != 0u && "Align can't be 0.");
   Skew %= Align;
   return (Value - Skew) / Align * Align + Skew;
 }
 

llvm/lib/Transforms/Utils/LoopUtils.cpp

Optional<unsigned> llvm::getLoopEstimatedTripCount(Loop *L) {
   if (any_of(ExitBlocks, [](const BasicBlock *EB) {
         return !EB->getTerminatingDeoptimizeCall();
       }))
     return None;
 
   // To estimate the number of times the loop body was executed, we want to
   // know the number of times the backedge was taken, vs. the number of times
   // we exited the loop.
-  uint64_t TrueVal, FalseVal;
-  if (!LatchBR->extractProfMetadata(TrueVal, FalseVal))
+  uint64_t LatchCycleWeight, LatchExitWeight;
+  if (!LatchBR->extractProfMetadata(LatchCycleWeight, LatchExitWeight))
     return None;
 
-  if (!TrueVal || !FalseVal)
+  if (LatchBR->getSuccessor(0) != L->getHeader())
+    std::swap(LatchCycleWeight, LatchExitWeight);
+
+  if (!LatchCycleWeight || !LatchExitWeight)
     return 0;
 
-  // Divide the count of the backedge by the count of the edge exiting the loop,
-  // rounding to nearest.
-  if (LatchBR->getSuccessor(0) == L->getHeader())
-    return (TrueVal + (FalseVal / 2)) / FalseVal;
-  else
-    return (FalseVal + (TrueVal / 2)) / TrueVal;
+  // Trip count is a ratio of loop header block execution weight to loop exit
+  // weight.
+  uint64_t HeaderBlockWeight = LatchCycleWeight + LatchExitWeight;
+  return llvm::divideNearest(HeaderBlockWeight, LatchExitWeight);
 }
 
 bool llvm::hasIterationCountInvariantInParent(Loop *InnerLoop,
                                               ScalarEvolution &SE) {
   Loop *OuterL = InnerLoop->getParentLoop();
   if (!OuterL)
     return true;
 
   // Get the backedge taken count for the inner loop

llvm/test/Transforms/LoopUnroll/peel-loop-conditions-pgo-1.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -S -loop-unroll -loop-unroll -verify-dom-info -debug-only=loop-unroll -unroll-peel-max-count=7 2>&1 | FileCheck %s
 ; REQUIRES: asserts
 
 declare void @f1()
 declare void @f2()
 
 ; Check that we can peel off iterations that make conditions true.
 ; The second invocation of loop-unroll will do profile based peeling of
 ; remained iterations.
 define void @test1(i32 %k) !prof !4 {
 ; CHECK: Loop Unroll: F[test1] Loop %for.body
 ; CHECK: PEELING loop %for.body with iteration count 2!
-; CHECK: PEELING loop %for.body with iteration count 4!
+; CHECK: PEELING loop %for.body with iteration count 5!
 ; CHECK: llvm.loop.unroll.disable
 for.body.lr.ph:
   br label %for.body
 
 for.body:
   %i.05 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.inc ]
   %cmp1 = icmp ult i32 %i.05, 2
   br i1 %cmp1, label %if.then, label %if.else

llvm/test/Transforms/LoopUnroll/peel-loop-pgo-deopt-idom-2.ll

 ; REQUIRES: asserts
 ; RUN: opt < %s -S -debug-only=loop-unroll -loop-unroll -unroll-runtime -unroll-peel-multi-deopt-exit 2>&1 | FileCheck %s
 ; RUN: opt < %s -S -debug-only=loop-unroll -unroll-peel-multi-deopt-exit -passes='require<profile-summary>,function(require<opt-remark-emit>,unroll)' 2>&1 | FileCheck %s
 
 ; Regression test for setting the correct idom for exit blocks.
 
 ; CHECK: Loop Unroll: F[basic]
-; CHECK: PEELING loop %for.body with iteration count 1!
+; CHECK: PEELING loop %for.body with iteration count 2!
 
 define i32 @basic(i32* %p, i32 %k, i1 %c1, i1 %c2) #0 !prof !3 {
 entry:
   br label %for.body
 
 for.body:
   %i.05 = phi i32 [ 0, %entry ], [ %inc, %latch ]
   %p.addr.04 = phi i32* [ %p, %entry ], [ %incdec.ptr, %latch ]

llvm/test/Transforms/LoopUnroll/peel-loop-pgo-deopt-idom.ll

 ; REQUIRES: asserts
 ; RUN: opt < %s -S -debug-only=loop-unroll -loop-unroll -unroll-runtime -unroll-peel-multi-deopt-exit 2>&1 | FileCheck %s
 ; RUN: opt < %s -S -debug-only=loop-unroll -unroll-peel-multi-deopt-exit -passes='require<profile-summary>,function(require<opt-remark-emit>,unroll)' 2>&1 | FileCheck %s
 
 ; Regression test for setting the correct idom for exit blocks.
 
 ; CHECK: Loop Unroll: F[basic]
-; CHECK: PEELING loop %for.body with iteration count 1!
+; CHECK: PEELING loop %for.body with iteration count 2!
 
 define i32 @basic(i32* %p, i32 %k, i1 %c1, i1 %c2) #0 !prof !3 {
 entry:
   %cmp3 = icmp slt i32 0, %k
   br i1 %cmp3, label %for.body.lr.ph, label %for.end
 
 for.body.lr.ph:                                   ; preds = %entry
   br label %for.body

llvm/test/Transforms/LoopUnroll/peel-loop-pgo-deopt.ll

 ; REQUIRES: asserts
 ; RUN: opt < %s -S -debug-only=loop-unroll -loop-unroll -unroll-runtime -unroll-peel-multi-deopt-exit 2>&1 | FileCheck %s
 ; RUN: opt < %s -S -debug-only=loop-unroll -unroll-peel-multi-deopt-exit -passes='require<profile-summary>,function(require<opt-remark-emit>,unroll)' 2>&1 | FileCheck %s
 ; RUN: opt < %s -S -debug-only=loop-unroll -unroll-peel-multi-deopt-exit -passes='require<profile-summary>,function(require<opt-remark-emit>,unroll<no-profile-peeling>)' 2>&1 | FileCheck %s --check-prefixes=CHECK-NO-PEEL
 
 ; Make sure we use the profile information correctly to peel-off 3 iterations
 ; from the loop, and update the branch weights for the peeled loop properly.
 ; All side exits to deopt does not change weigths.
 
 ; CHECK: Loop Unroll: F[basic]
-; CHECK: PEELING loop %for.body with iteration count 3!
+; CHECK: PEELING loop %for.body with iteration count 4!
 ; CHECK-NO-PEEL-NOT: PEELING loop %for.body
 ; CHECK-LABEL: @basic
 ; CHECK: br i1 %c, label %{{.*}}, label %side_exit, !prof !15
 ; CHECK: br i1 %{{.*}}, label %[[NEXT0:.*]], label %for.cond.for.end_crit_edge, !prof !16
 ; CHECK: [[NEXT0]]:
 ; CHECK: br i1 %c, label %{{.*}}, label %side_exit, !prof !15
 ; CHECK: br i1 %{{.*}}, label %[[NEXT1:.*]], label %for.cond.for.end_crit_edge, !prof !17
 ; CHECK: [[NEXT1]]:

llvm/test/Transforms/LoopUnroll/peel-loop-pgo.ll

 ; RUN: opt < %s -S -debug-only=loop-unroll -loop-unroll 2>&1 | FileCheck %s
 ; RUN: opt < %s -S -debug-only=loop-unroll -passes='require<profile-summary>,function(require<opt-remark-emit>,unroll)' 2>&1 | FileCheck %s
 ; Confirm that peeling is disabled if the number of counts required to reach
 ; the hot percentile is above the threshold.
 ; RUN: opt < %s -S -profile-summary-huge-working-set-size-threshold=9 -debug-only=loop-unroll -passes='require<profile-summary>,function(require<opt-remark-emit>,unroll)' 2>&1 | FileCheck %s --check-prefix=NOPEEL
 ; REQUIRES: asserts
 
 ; Make sure we use the profile information correctly to peel-off 3 iterations
 ; from the loop, and update the branch weights for the peeled loop properly.
 
 ; CHECK: Loop Unroll: F[basic]
-; CHECK: PEELING loop %for.body with iteration count 3!
+; CHECK: PEELING loop %for.body with iteration count 4!
 ; CHECK: Loop Unroll: F[optsize]
 ; CHECK-NOT: PEELING
 
 ; Confirm that no peeling occurs when we are performing full unrolling.
 ; RUN: opt < %s -S -debug-only=loop-unroll -passes='require<opt-remark-emit>,loop(unroll-full)' 2>&1 | FileCheck %s --check-prefix=NOPEEL
 ; NOPEEL-NOT: PEELING
 
 ; CHECK-LABEL: @basic

llvm/unittests/Transforms/Utils/CMakeLists.txt

add_llvm_unittest(UtilsTests
   ASanStackFrameLayoutTest.cpp
   BasicBlockUtilsTest.cpp
   CloningTest.cpp
   CodeExtractorTest.cpp
   CodeMoverUtilsTest.cpp
   FunctionComparatorTest.cpp
   IntegerDivisionTest.cpp
   LocalTest.cpp
+  LoopUtilsTest.cpp
   SizeOptsTest.cpp
   SSAUpdaterBulkTest.cpp
   UnrollLoopTest.cpp
   ValueMapperTest.cpp
   VFABIUtils.cpp
   )

llvm/unittests/Transforms/Utils/LoopUtilsTest.cpp

+//===--------- LoopUtilsTest.cpp - unit tests -----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/LoopUtils.h"
+#include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/LoopPass.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Transforms/Scalar/LoopPassManager.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+struct LoopUtilsPass : public PassInfoMixin<LoopUtilsPass> {
+  template <typename T> LoopUtilsPass(T &&Arg) : Func(std::forward<T>(Arg)) {}
+
+  PreservedAnalyses run(Loop &L, LoopAnalysisManager &AM,
+      LoopStandardAnalysisResults &AR, LPMUpdater &U) {
+    return Func(L, AM, AR, U);
+  }
+
+  std::function<PreservedAnalyses(Loop &, LoopAnalysisManager &,
+                                  LoopStandardAnalysisResults &, LPMUpdater &)>
+  Func;
+};
+
+struct LoopUtilsTest : public testing::Test {
+
+  PassBuilder PB;
+  LoopAnalysisManager LAM;
+  FunctionAnalysisManager FAM;
+  CGSCCAnalysisManager CGAM;
+  ModuleAnalysisManager MAM;
+  LLVMContext Context;
+  std::unique_ptr<Module> M;
+
+  LoopUtilsTest() {
+    // Register all the basic analyses with the managers.
+    PB.registerModuleAnalyses(MAM);
+    PB.registerCGSCCAnalyses(CGAM);
+    PB.registerFunctionAnalyses(FAM);
+    PB.registerLoopAnalyses(LAM);
+    PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+  }
+
+  void ParseAssembly(const char *IR) {
+    SMDiagnostic Error;
+    M = parseAssemblyString(IR, Error, Context);
+
+    std::string errMsg;
+    raw_string_ostream os(errMsg);
+    Error.print("", os);
+
+    // A failure here means that the test itself is buggy.
+    if (!M)
+      report_fatal_error(os.str().c_str());
+  }
+};
+
+TEST_F(LoopUtilsTest, getLoopEstimatedTripCountTest1) {
+  ParseAssembly(R"IR(
+      @a = dso_local global [1024 x i32] zeroinitializer, align 16
+      @b = dso_local global [1024 x i32] zeroinitializer, align 16
+      define dso_local void @_Z3foov() {
+      entry:
+        br label %for.body
+      for.cond.cleanup:                                 ; preds = %for.body
+        ret void
+      for.body:                                         ; preds = %for.body, %entry
+        %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+        %arrayidx = getelementptr inbounds [1024 x i32], [1024 x i32]* @b, i64 0, i64 %indvars.iv
+        %0 = load i32, i32* %arrayidx, align 4
+        %1 = trunc i64 %indvars.iv to i32
+        %mul = mul nsw i32 %0, %1
+        %arrayidx2 = getelementptr inbounds [1024 x i32], [1024 x i32]* @a, i64 0, i64 %indvars.iv
+        %2 = load i32, i32* %arrayidx2, align 4
+        %add = add nsw i32 %2, %mul
+        store i32 %add, i32* %arrayidx2, align 4
+        %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+        %exitcond = icmp eq i64 %indvars.iv.next, 1024
+        br i1 %exitcond, label %for.cond.cleanup, label %for.body, !prof !1}
+
+        !1 = !{!"branch_weights", i32 1, i32 15 })IR");
+
+
+  LoopPassManager LPM;
+  LPM.addPass(LoopUtilsPass([&](Loop &L, LoopAnalysisManager &AM,
+      LoopStandardAnalysisResults &AR, LPMUpdater &U) {
+    auto TC = getLoopEstimatedTripCount(&L);
+    EXPECT_TRUE(TC.hasValue());
+    EXPECT_EQ(*TC, (unsigned)16);
+    return PreservedAnalyses::all();
+  }));
+
+  FunctionPassManager FPM;
+  FPM.addPass(createFunctionToLoopPassAdaptor(std::move(LPM)));
+
+  ModulePassManager MPM;
+  MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+
+  MPM.run(*M, MAM);
+}
+}