[PGO] Add ProfileInjector and ProfileVerifier passes

llvm/include/llvm/Transforms/Utils/ProfileVerify.h

@@ -0,0 +1,36 @@
+//===- ProfileVerify.h - Verify profile info for testing ----0-----*-C++-*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Inject profile information, as part of tests, to verify passes don't
+// accidentally drop it.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_TRANSFORMS_UTILS_PROFILEVERIFY_H
+#define LLVM_TRANSFORMS_UTILS_PROFILEVERIFY_H
+
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+/// Inject MD_prof metadata where it's missing. Used for testing that passes
+/// don't accidentally drop this metadata.
+class ProfileInjectorPass : public PassInfoMixin<ProfileInjectorPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+/// Checks that MD_prof is present on every instruction that supports it. Used
+/// in conjunction with the ProfileInjectorPass. MD_prof "unknown" is considered
+/// valid (i.e. !{!"unknown"})
+class ProfileVerifierPass : public PassInfoMixin<ProfileVerifierPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+} // namespace llvm
+#endif

llvm/lib/Passes/PassBuilder.cpp

@@ -357,6 +357,7 @@
 #include "llvm/Transforms/Utils/MoveAutoInit.h"
 #include "llvm/Transforms/Utils/NameAnonGlobals.h"
 #include "llvm/Transforms/Utils/PredicateInfo.h"
+#include "llvm/Transforms/Utils/ProfileVerify.h"
 #include "llvm/Transforms/Utils/RelLookupTableConverter.h"
 #include "llvm/Transforms/Utils/StripGCRelocates.h"
 #include "llvm/Transforms/Utils/StripNonLineTableDebugInfo.h"

llvm/lib/Passes/PassRegistry.def

@@ -517,6 +517,8 @@ FUNCTION_PASS("print<regions>", RegionInfoPrinterPass(errs()))
 FUNCTION_PASS("print<scalar-evolution>", ScalarEvolutionPrinterPass(errs()))
 FUNCTION_PASS("print<stack-safety-local>", StackSafetyPrinterPass(errs()))
 FUNCTION_PASS("print<uniformity>", UniformityInfoPrinterPass(errs()))
+FUNCTION_PASS("prof-inject", ProfileInjectorPass())
+FUNCTION_PASS("prof-verify", ProfileVerifierPass())
 FUNCTION_PASS("reassociate", ReassociatePass())
 FUNCTION_PASS("redundant-dbg-inst-elim", RedundantDbgInstEliminationPass())
 FUNCTION_PASS("reg2mem", RegToMemPass())

llvm/lib/Transforms/Utils/CMakeLists.txt

@@ -67,6 +67,7 @@ add_llvm_component_library(LLVMTransformUtils
   MoveAutoInit.cpp
   NameAnonGlobals.cpp
   PredicateInfo.cpp
+  ProfileVerify.cpp
   PromoteMemoryToRegister.cpp
   RelLookupTableConverter.cpp
   ScalarEvolutionExpander.cpp

llvm/lib/Transforms/Utils/ProfileVerify.cpp

@@ -0,0 +1,113 @@
+//===- ProfileVerify.cpp - Verify profile info for testing ----------------===//
+//
+// 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/ProfileVerify.h"
+#include "llvm/ADT/DynamicAPInt.h"
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/ProfDataUtils.h"
+#include "llvm/Support/BranchProbability.h"
+
+using namespace llvm;
+namespace {
+class ProfileInjector {
+  Function &F;
+  FunctionAnalysisManager &FAM;
+
+public:
+  static bool supportsBranchWeights(const Instruction &I) {
+    return isa<BranchInst>(&I) ||
+
+           isa<SwitchInst>(&I) ||
+
+           isa<IndirectBrInst>(&I) || isa<SelectInst>(&I) ||
+           isa<CallBrInst>(&I);
+  }
+
+  ProfileInjector(Function &F, FunctionAnalysisManager &FAM) : F(F), FAM(FAM) {}
+  bool inject();
+};
+} // namespace
+
+bool ProfileInjector::inject() {
+  auto &BPI = FAM.getResult<BranchProbabilityAnalysis>(F);
+
+  for (auto &BB : F) {
+    if (succ_size(&BB) <= 1)
+      continue;
+    auto *Term = BB.getTerminator();
+    assert(Term);
+    if (Term->getMetadata(LLVMContext::MD_prof) ||
+        !supportsBranchWeights(*Term))
+      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)));
+
+    const auto *FirstZeroDenominator =
+        find_if(Probs, [](const BranchProbability &P) {
+          return P.getDenominator() == 0;
+        });
+    assert(FirstZeroDenominator == Probs.end());
+    const auto *FirstNonzeroNumerator =
+        find_if(Probs, [](const BranchProbability &P) {
+          return P.getNumerator() != 0;
+        });
+    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::lcm(GCD, DynamicAPInt(Prob.getNumerator()));
+    }
+    SmallVector<uint32_t> Weights;
+    Weights.reserve(Term->getNumSuccessors());
+    for (const auto &Prob : Probs) {
+      auto W = Prob.getNumerator() * LCM / GCD;
+      Weights.emplace_back(static_cast<int32_t>((int64_t)W));
+    }
+    setBranchWeights(*Term, Weights, false);
+  }
+  return true;
+}
+
+PreservedAnalyses ProfileInjectorPass::run(Function &F,
+                                           FunctionAnalysisManager &FAM) {
+  ProfileInjector PI(F, FAM);
+  if (!PI.inject())
+    return PreservedAnalyses::all();
+
+  return PreservedAnalyses::none();
+}
+
+PreservedAnalyses ProfileVerifierPass::run(Function &F,
+                                           FunctionAnalysisManager &FAM) {
+  bool Changed = false;
+  for (auto &BB : F)
+    if (succ_size(&BB) >= 2)
+      if (auto *Term = BB.getTerminator())
+        if (ProfileInjector::supportsBranchWeights(*Term)) {
+          if (!Term->getMetadata(LLVMContext::MD_prof)) {
+            F.getContext().emitError("Profile verification failed");
+          } else {
+            Changed = true;
+          }
+        }
+
+  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}

llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll

@@ -0,0 +1,20 @@
+; Test that prof-inject only injects missing metadata
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no, !prof !0
+yes:
+  br i1 %c, label %yes2, label %no
+yes2:
+  ret void
+no:
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: br i1 %c, label %yes, label %no, !prof !0
+; CHECK: br i1 %c, label %yes2, label %no, !prof !1
+; CHECK: !0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: !1 = !{!"branch_weights", i32 429496729, i32 715827882}

llvm/test/Transforms/PGOProfile/prof-verify-existing.ll

@@ -0,0 +1,21 @@
+; Test that prof-verify does not modify existing metadata (incl. "unknown")
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s
+; RUN: opt -passes=prof-verify %s -S --disable-output
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no, !prof !0
+yes:
+  br i1 %c, label %yes2, label %no, !prof !1
+yes2:
+  ret void
+no:
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 1, i32 2}
+!1 = !{!"unknown"}
+; CHECK: br i1 %c, label %yes, label %no, !prof !0
+; CHECK: !0 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: !1 = !{!"unknown"}

llvm/test/Transforms/PGOProfile/prof-verify.ll

@@ -0,0 +1,19 @@
+; Test prof-inject and prof-verify
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s --check-prefix=INJECT
+; RUN: not opt -passes=prof-verify %s -S -o - 2>&1 | FileCheck %s --check-prefix=VERIFY
+; RUN: opt -passes=prof-inject,prof-verify %s --disable-output
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no
+yes:
+  ret void
+no:
+  ret void
+}
+
+; INJECT: br i1 %c, label %yes, label %no, !prof !0
+; INJECT: !0 = !{!"branch_weights", i32 429496729, i32 715827882}
+
+; VERIFY: Profile verification failed