[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

730

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 ↗

(On Diff #235561)

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

724–725

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

724–725

Completely agree.

[ebrevnov]

Fixed raised issue + Rebase.

[ebrevnov]

ping @Ayal

[Ayal]

llvm/lib/Transforms/Utils/LoopUtils.cpp

729

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

729

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

727–730

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

732

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

728

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

llvm/lib/Transforms/Utils/LoopUtils.cpp

Optional<unsigned> llvm::getLoopEstimatedTripCount(Loop *L) {
   // we exited the loop.
   uint64_t BackedgeTakenWeight, LatchExitWeight;
   if (!LatchBR->extractProfMetadata(BackedgeTakenWeight, LatchExitWeight))
     return None;
 
   if (LatchBR->getSuccessor(0) != L->getHeader())
     std::swap(BackedgeTakenWeight, LatchExitWeight);
 
-  if (!BackedgeTakenWeight || !LatchExitWeight)
-    return 0;
+  if (!LatchExitWeight)
+    return None;
 
-  // Divide the count of the backedge by the count of the edge exiting the loop,
-  // rounding to nearest.
-  return llvm::divideNearest(BackedgeTakenWeight, LatchExitWeight);
+  // Trip count is a ratio of loop header block execution weight to loop exit
+  // weight.
+  uint64_t HeaderBlockWeight = BackedgeTakenWeight + 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
   BasicBlock *InnerLoopLatch = InnerLoop->getLoopLatch();

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/LoopUtilsTest.cpp

 //===- LoopUtilsTest.cpp - Unit tests for LoopUtils -----------------------===//
 //
 // 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;
 
-static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
-  SMDiagnostic Err;
-  std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
-  if (!Mod)
-    Err.print("LoopUtilsTests", errs());
-  return Mod;
-}
-
-static void run(Module &M, StringRef FuncName,
-                function_ref<void(Function &F, DominatorTree &DT,
-                                  ScalarEvolution &SE, LoopInfo &LI)>
-                    Test) {
-  Function *F = M.getFunction(FuncName);
-  DominatorTree DT(*F);
-  TargetLibraryInfoImpl TLII;
-  TargetLibraryInfo TLI(TLII);
-  AssumptionCache AC(*F);
-  LoopInfo LI(DT);
-  ScalarEvolution SE(*F, TLI, AC, DT, LI);
-  Test(*F, DT, SE, LI);
-}
-
-TEST(LoopUtils, DeleteDeadLoopNest) {
-  LLVMContext C;
-  std::unique_ptr<Module> M =
-      parseIR(C, "define void @foo() {\n"
-                 "entry:\n"
-                 "  br label %for.i\n"
-                 "for.i:\n"
-                 "  %i = phi i64 [ 0, %entry ], [ %inc.i, %for.i.latch ]\n"
-                 "  br label %for.j\n"
-                 "for.j:\n"
-                 "  %j = phi i64 [ 0, %for.i ], [ %inc.j, %for.j ]\n"
-                 "  %inc.j = add nsw i64 %j, 1\n"
-                 "  %cmp.j = icmp slt i64 %inc.j, 100\n"
-                 "  br i1 %cmp.j, label %for.j, label %for.k.preheader\n"
-                 "for.k.preheader:\n"
-                 "  br label %for.k\n"
-                 "for.k:\n"
-                 "  %k = phi i64 [ %inc.k, %for.k ], [ 0, %for.k.preheader ]\n"
-                 "  %inc.k = add nsw i64 %k, 1\n"
-                 "  %cmp.k = icmp slt i64 %inc.k, 100\n"
-                 "  br i1 %cmp.k, label %for.k, label %for.i.latch\n"
-                 "for.i.latch:\n"
-                 "  %inc.i = add nsw i64 %i, 1\n"
-                 "  %cmp.i = icmp slt i64 %inc.i, 100\n"
-                 "  br i1 %cmp.i, label %for.i, label %for.end\n"
-                 "for.end:\n"
-                 "  ret void\n"
-                 "}\n");
+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, DeleteDeadLoopNest) {
+  ParseAssembly("define void @foo() {\n"
 
-  run(*M, "foo",
-      [&](Function &F, DominatorTree &DT, ScalarEvolution &SE, LoopInfo &LI) {
-        assert(LI.begin() != LI.end() && "Expecting loops in function F");
-        Loop *L = *LI.begin();
-        assert(L && L->getName() == "for.i" && "Expecting loop for.i");
+
+  LoopPassManager LPM;
+  LPM.addPass(LoopUtilsPass([=](Loop &L, LoopAnalysisManager &AM,
+                                LoopStandardAnalysisResults &AR,
+                                LPMUpdater &U) -> PreservedAnalyses {
+    auto &F = *L.getHeader()->getParent();
+    auto &SE = AR.SE;
+    auto &LI = AR.LI;
+    auto &DT = AR.DT;
+
+    std::string LoopName = L.getName();
+    if (LoopName == "for.i") {
 
-        deleteDeadLoop(L, &DT, &SE, &LI);
+      deleteDeadLoop(&L, &DT, &SE, &LI);
 
-        assert(DT.verify(DominatorTree::VerificationLevel::Fast) &&
-               "Expecting valid dominator tree");
-        LI.verify(DT);
-        assert(LI.begin() == LI.end() &&
-               "Expecting no loops left in function F");
-        SE.verify();
+      assert(LI.begin() == LI.end() && "Expecting no loops left in function F");
 
-        Function::iterator FI = F.begin();
-        BasicBlock *Entry = &*(FI++);
-        assert(Entry->getName() == "entry" && "Expecting BasicBlock entry");
-        const BranchInst *BI = dyn_cast<BranchInst>(Entry->getTerminator());
-        assert(BI && "Expecting valid branch instruction");
-        EXPECT_EQ(BI->getNumSuccessors(), (unsigned)1);
-        EXPECT_EQ(BI->getSuccessor(0)->getName(), "for.end");
-      });
+
+      U.markLoopAsDeleted(L, LoopName);
+      return PreservedAnalyses::none();
+    }
+    return PreservedAnalyses::all();
+  }));
+
+  FunctionPassManager FPM;
+  FPM.addPass(createFunctionToLoopPassAdaptor(std::move(LPM)));
+
+  ModulePassManager MPM;
+  MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+
+  MPM.run(*M, MAM);
+}
+
+TEST_F(LoopUtilsTest, getLoopEstimatedTripCountTest) {
+  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) -> PreservedAnalyses {
+    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);
+}
 }