[profcheck] Option to inject distinct small weights

llvm/lib/Transforms/Utils/ProfileVerify.cpp

@@ -28,6 +28,10 @@ static cl::opt<bool>
     AnnotateSelect("profcheck-annotate-select", cl::init(true),
                    cl::desc("Also inject (if missing) and verify MD_prof for "
                             "`select` instructions"));
+static cl::opt<bool>
+    WeightsForTest("profcheck-weights-for-test", cl::init(false),
+                   cl::desc("Generate weights with small values for tests."));
+
 static cl::opt<uint32_t> SelectTrueWeight(
     "profcheck-default-select-true-weight", cl::init(2U),
     cl::desc("When annotating `select` instructions, this value will be used "
@@ -91,6 +95,10 @@ bool ProfileInjector::inject() {
   if (F.getEntryCount(/*AllowSynthetic=*/true)->getCount() == 0)
     return false;
   bool Changed = false;
+  // Cycle through the weights list. If we didn't, tests with more than (say)
+  // one conditional branch would have the same !prof metadata on all of them,
+  // and numerically that may make for a poor unit test.
+  uint32_t WeightsForTestOffset = 0;
   for (auto &BB : F) {
     if (AnnotateSelect) {
       for (auto &I : BB) {
@@ -103,38 +111,48 @@ bool ProfileInjector::inject() {
     if (!Term || Term->getMetadata(LLVMContext::MD_prof))
       continue;
     SmallVector<BranchProbability> Probs;
-    Probs.reserve(Term->getNumSuccessors());
-    for (auto I = 0U, E = Term->getNumSuccessors(); I < E; ++I)
-      Probs.emplace_back(BPI.getEdgeProbability(&BB, Term->getSuccessor(I)));
 
-    assert(llvm::find_if(Probs,
-                         [](const BranchProbability &P) {
-                           return P.isUnknown();
-                         }) == Probs.end() &&
-           "All branch probabilities should be valid");
-    const auto *FirstZeroDenominator =
-        find_if(Probs, [](const BranchProbability &P) {
-          return P.getDenominator() == 0;
-        });
-    (void)FirstZeroDenominator;
-    assert(FirstZeroDenominator == Probs.end());
-    const auto *FirstNonZeroNumerator =
-        find_if(Probs, [](const BranchProbability &P) { return !P.isZero(); });
-    assert(FirstNonZeroNumerator != Probs.end());
-    DynamicAPInt LCM(Probs[0].getDenominator());
-    DynamicAPInt GCD(FirstNonZeroNumerator->getNumerator());
-    for (const auto &Prob : drop_begin(Probs)) {
-      if (!Prob.getNumerator())
-        continue;
-      LCM = llvm::lcm(LCM, DynamicAPInt(Prob.getDenominator()));
-      GCD = llvm::gcd(GCD, DynamicAPInt(Prob.getNumerator()));
-    }
     SmallVector<uint32_t> Weights;
     Weights.reserve(Term->getNumSuccessors());
-    for (const auto &Prob : Probs) {
-      DynamicAPInt W =
-          (Prob.getNumerator() * LCM / GCD) / Prob.getDenominator();
-      Weights.emplace_back(static_cast<uint32_t>((int64_t)W));
+    if (WeightsForTest) {
+      static const std::array Primes{3,  5,  7,  11, 13, 17, 19, 23, 29, 31,
+                                     37, 41, 43, 47, 53, 59, 61, 67, 71};
+      for (uint32_t I = 0, E = Term->getNumSuccessors(); I < E; ++I)
+        Weights.emplace_back(
+            Primes[(WeightsForTestOffset + I) % Primes.size()]);
+      ++WeightsForTestOffset;
+    } else {
+      Probs.reserve(Term->getNumSuccessors());
+      for (auto I = 0U, E = Term->getNumSuccessors(); I < E; ++I)
+        Probs.emplace_back(BPI.getEdgeProbability(&BB, Term->getSuccessor(I)));
+
+      assert(llvm::find_if(Probs,
+                           [](const BranchProbability &P) {
+                             return P.isUnknown();
+                           }) == Probs.end() &&
+             "All branch probabilities should be valid");
+      const auto *FirstZeroDenominator =
+          find_if(Probs, [](const BranchProbability &P) {
+            return P.getDenominator() == 0;
+          });
+      (void)FirstZeroDenominator;
+      assert(FirstZeroDenominator == Probs.end());
+      const auto *FirstNonZeroNumerator = find_if(
+          Probs, [](const BranchProbability &P) { return !P.isZero(); });
+      assert(FirstNonZeroNumerator != Probs.end());
+      DynamicAPInt LCM(Probs[0].getDenominator());
+      DynamicAPInt GCD(FirstNonZeroNumerator->getNumerator());
+      for (const auto &Prob : drop_begin(Probs)) {
+        if (!Prob.getNumerator())
+          continue;
+        LCM = llvm::lcm(LCM, DynamicAPInt(Prob.getDenominator()));
+        GCD = llvm::gcd(GCD, DynamicAPInt(Prob.getNumerator()));
+      }
+      for (const auto &Prob : Probs) {
+        DynamicAPInt W =
+            (Prob.getNumerator() * LCM / GCD) / Prob.getDenominator();
+        Weights.emplace_back(static_cast<uint32_t>((int64_t)W));
+      }
     }
     setBranchWeights(*Term, Weights, /*IsExpected=*/false);
     Changed = true;

llvm/test/Transforms/PGOProfile/profcheck-synthetic.ll

@@ -0,0 +1,73 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 6
+; RUN: opt -passes=prof-inject -profcheck-weights-for-test %s -S -o - | FileCheck %s --check-prefixes=TEST,CHECK
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s --check-prefixes=NORMAL,CHECK
+
+define void @foo(i32 %cond) {
+; TEST-LABEL: define void @foo(
+; TEST-SAME: i32 [[COND:%.*]]) !prof [[PROF0:![0-9]+]] {
+; TEST-NEXT:    [[I:%.*]] = icmp eq i32 [[COND]], 0
+; TEST-NEXT:    br i1 [[I]], label %[[A:.*]], label %[[B:.*]], !prof [[PROF1:![0-9]+]]
+; TEST:       [[A]]:
+; TEST-NEXT:    switch i32 [[COND]], label %[[DEFAULT:.*]] [
+; TEST-NEXT:      i32 10, label %[[C:.*]]
+; TEST-NEXT:      i32 20, label %[[D:.*]]
+; TEST-NEXT:    ], !prof [[PROF2:![0-9]+]]
+; TEST:       [[BB1:.*:]]
+; TEST-NEXT:    br label %[[B]]
+; TEST:       [[B]]:
+; TEST-NEXT:    ret void
+; TEST:       [[DEFAULT]]:
+; TEST-NEXT:    ret void
+; TEST:       [[C]]:
+; TEST-NEXT:    ret void
+; TEST:       [[D]]:
+; TEST-NEXT:    ret void
+;
+; NORMAL-LABEL: define void @foo(
+; NORMAL-SAME: i32 [[COND:%.*]]) !prof [[PROF0:![0-9]+]] {
+; NORMAL-NEXT:    [[I:%.*]] = icmp eq i32 [[COND]], 0
+; NORMAL-NEXT:    br i1 [[I]], label %[[A:.*]], label %[[B:.*]], !prof [[PROF1:![0-9]+]]
+; NORMAL:       [[A]]:
+; NORMAL-NEXT:    switch i32 [[COND]], label %[[DEFAULT:.*]] [
+; NORMAL-NEXT:      i32 10, label %[[C:.*]]
+; NORMAL-NEXT:      i32 20, label %[[D:.*]]
+; NORMAL-NEXT:    ], !prof [[PROF2:![0-9]+]]
+; NORMAL:       [[BB1:.*:]]
+; NORMAL-NEXT:    br label %[[B]]
+; NORMAL:       [[B]]:
+; NORMAL-NEXT:    ret void
+; NORMAL:       [[DEFAULT]]:
+; NORMAL-NEXT:    ret void
+; NORMAL:       [[C]]:
+; NORMAL-NEXT:    ret void
+; NORMAL:       [[D]]:
+; NORMAL-NEXT:    ret void
+;
+  %i = icmp eq i32 %cond, 0
+  br i1 %i, label %a, label %b
+a:
+  switch i32 %cond, label %default [
+  i32 10, label %c
+  i32 20, label %d
+  ]
+  br label %b
+b:
+  ret void
+default:
+  ret void
+c:
+  ret void
+d:
+  ret void
+}
+;.
+; TEST: [[PROF0]] = !{!"function_entry_count", i64 1000}
+; TEST: [[PROF1]] = !{!"branch_weights", i32 3, i32 5}
+; TEST: [[PROF2]] = !{!"branch_weights", i32 5, i32 7, i32 11}
+;.
+; NORMAL: [[PROF0]] = !{!"function_entry_count", i64 1000}
+; NORMAL: [[PROF1]] = !{!"branch_weights", i32 3, i32 5}
+; NORMAL: [[PROF2]] = !{!"branch_weights", i32 1, i32 1, i32 1}
+;.
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; CHECK: {{.*}}