[misexpect] Re-implement MisExpect Diagnostics

Reimplements MisExpect diagnostics from D66324 to reconstruct its
original checking methodology only using MD_prof branch_weights
metadata.

New checks rely on 2 invariants:

1) For frontend instrumentation, MD_prof branch_weights will always be
   populated before llvm.expect intrinsics are lowered.

2) for IR and sample profiling, llvm.expect intrinsics will always be
   lowered before branch_weights are populated from the IR profiles.

These invariants allow the checking to assume how the existing branch
weights are populated depending on the profiling method used, and emit
the correct diagnostics. If these invariants are ever invalidated, the
MisExpect related checks would need to be updated, potentially by
re-introducing MD_misexpect metadata, and ensuring it always will be
transformed the same way as branch_weights in other optimization passes.

Frontend based profiling is now enabled without using LLVM Args, by
introducing a new CodeGen option, and checking if the -Wmisexpect flag
has been passed on the command line.

Reviewed By: tejohnson

Differential Revision: https://reviews.llvm.org/D115907

8 files changed, 280 insertions, 2 deletions

diff --git a/llvm/lib/IR/DiagnosticInfo.cpp b/llvm/lib/IR/DiagnosticInfo.cpp
index f46f0fd..50fe682 100644
--- a/llvm/lib/IR/DiagnosticInfo.cpp
+++ b/llvm/lib/IR/DiagnosticInfo.cpp
@@ -393,6 +393,17 @@ std::string DiagnosticInfoOptimizationBase::getMsg() const {
   return OS.str();
 }
 
+DiagnosticInfoMisExpect::DiagnosticInfoMisExpect(const Instruction *Inst,
+                                                 Twine &Msg)
+    : DiagnosticInfoWithLocationBase(DK_MisExpect, DS_Warning,
+                                     *Inst->getParent()->getParent(),
+                                     Inst->getDebugLoc()),
+      Msg(Msg) {}
+
+void DiagnosticInfoMisExpect::print(DiagnosticPrinter &DP) const {
+  DP << getLocationStr() << ": " << getMsg();
+}
+
 void OptimizationRemarkAnalysisFPCommute::anchor() {}
 void OptimizationRemarkAnalysisAliasing::anchor() {}
 
diff --git a/llvm/lib/IR/LLVMContext.cpp b/llvm/lib/IR/LLVMContext.cpp
index f4e917c..df4a007 100644
--- a/llvm/lib/IR/LLVMContext.cpp
+++ b/llvm/lib/IR/LLVMContext.cpp
@@ -138,10 +138,22 @@ bool LLVMContext::getDiagnosticsHotnessRequested() const {
 void LLVMContext::setDiagnosticsHotnessThreshold(Optional<uint64_t> Threshold) {
   pImpl->DiagnosticsHotnessThreshold = Threshold;
 }
-
+void LLVMContext::setMisExpectWarningRequested(bool Requested) {
+  pImpl->MisExpectWarningRequested = Requested;
+}
+bool LLVMContext::getMisExpectWarningRequested() const {
+  return pImpl->MisExpectWarningRequested;
+}
 uint64_t LLVMContext::getDiagnosticsHotnessThreshold() const {
   return pImpl->DiagnosticsHotnessThreshold.getValueOr(UINT64_MAX);
 }
+void LLVMContext::setDiagnosticsMisExpectTolerance(
+    Optional<uint64_t> Tolerance) {
+  pImpl->DiagnosticsMisExpectTolerance = Tolerance;
+}
+uint64_t LLVMContext::getDiagnosticsMisExpectTolerance() const {
+  return pImpl->DiagnosticsMisExpectTolerance.getValueOr(0);
+}
 
 bool LLVMContext::isDiagnosticsHotnessThresholdSetFromPSI() const {
   return !pImpl->DiagnosticsHotnessThreshold.hasValue();
diff --git a/llvm/lib/IR/LLVMContextImpl.h b/llvm/lib/IR/LLVMContextImpl.h
index b87d285..0b39029 100644
--- a/llvm/lib/IR/LLVMContextImpl.h
+++ b/llvm/lib/IR/LLVMContextImpl.h
@@ -1380,6 +1380,11 @@ public:
   /// If threshold option is not specified, it is disabled (0) by default.
   Optional<uint64_t> DiagnosticsHotnessThreshold = 0;
 
+  /// The percentage of difference between profiling branch weights and
+  // llvm.expect branch weights to tolerate when emiting MisExpect diagnostics
+  Optional<uint64_t> DiagnosticsMisExpectTolerance = 0;
+  bool MisExpectWarningRequested = false;
+
   /// The specialized remark streamer used by LLVM's OptimizationRemarkEmitter.
   std::unique_ptr<LLVMRemarkStreamer> LLVMRS;
 
diff --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index 8232406..4bd7900 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -74,6 +74,8 @@
 #include "llvm/Transforms/Instrumentation.h"
 #include "llvm/Transforms/Utils/CallPromotionUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/MisExpect.h"
+#include "llvm/Transforms/Utils/SampleProfileInference.h"
 #include "llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h"
 #include "llvm/Transforms/Utils/SampleProfileLoaderBaseUtil.h"
 #include <algorithm>
@@ -1739,6 +1741,8 @@ void SampleProfileLoader::generateMDProfMetadata(Function &F) {
       }
     }
 
+    misexpect::checkExpectAnnotations(*TI, Weights, /*IsFrontend=*/false);
+
     uint64_t TempWeight;
     // Only set weights if there is at least one non-zero weight.
     // In any other case, let the analyzer set weights.
diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index d421894..35ea646 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -110,6 +110,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Instrumentation.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/MisExpect.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <algorithm>
 #include <cassert>
@@ -2117,6 +2118,8 @@ void llvm::setProfMetadata(Module *M, Instruction *TI,
     dbgs() << W << " ";
   } dbgs() << "\n";);
 
+  misexpect::checkExpectAnnotations(*TI, Weights, /*IsFrontend=*/false);
+
   TI->setMetadata(LLVMContext::MD_prof, MDB.createBranchWeights(Weights));
   if (EmitBranchProbability) {
     std::string BrCondStr = getBranchCondString(TI);
diff --git a/llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp b/llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp
index 97c8f33..88fad98 100644
--- a/llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp
+++ b/llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp
@@ -25,6 +25,7 @@
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Utils/MisExpect.h"
 
 using namespace llvm;
 
@@ -99,6 +100,8 @@ static bool handleSwitchExpect(SwitchInst &SI) {
   uint64_t Index = (Case == *SI.case_default()) ? 0 : Case.getCaseIndex() + 1;
   Weights[Index] = LikelyBranchWeightVal;
 
+  misexpect::checkExpectAnnotations(SI, Weights, /*IsFrontend=*/true);
+
   SI.setCondition(ArgValue);
 
   SI.setMetadata(LLVMContext::MD_prof,
@@ -313,13 +316,16 @@ template <class BrSelInst> static bool handleBrSelExpect(BrSelInst &BSI) {
   std::tie(LikelyBranchWeightVal, UnlikelyBranchWeightVal) =
       getBranchWeight(Fn->getIntrinsicID(), CI, 2);
 
+  SmallVector<uint32_t, 4> ExpectedWeights;
   if ((ExpectedValue->getZExtValue() == ValueComparedTo) ==
       (Predicate == CmpInst::ICMP_EQ)) {
     Node =
         MDB.createBranchWeights(LikelyBranchWeightVal, UnlikelyBranchWeightVal);
+    ExpectedWeights = {LikelyBranchWeightVal, UnlikelyBranchWeightVal};
   } else {
     Node =
         MDB.createBranchWeights(UnlikelyBranchWeightVal, LikelyBranchWeightVal);
+    ExpectedWeights = {UnlikelyBranchWeightVal, LikelyBranchWeightVal};
   }
 
   if (CmpI)
@@ -327,6 +333,8 @@ template <class BrSelInst> static bool handleBrSelExpect(BrSelInst &BSI) {
   else
     BSI.setCondition(ArgValue);
 
+  misexpect::checkFrontendInstrumentation(BSI, ExpectedWeights);
+
   BSI.setMetadata(LLVMContext::MD_prof, Node);
 
   return true;
@@ -407,7 +415,7 @@ public:
 
   bool runOnFunction(Function &F) override { return lowerExpectIntrinsic(F); }
 };
-}
+} // namespace
 
 char LowerExpectIntrinsic::ID = 0;
 INITIALIZE_PASS(LowerExpectIntrinsic, "lower-expect",
diff --git a/llvm/lib/Transforms/Utils/CMakeLists.txt b/llvm/lib/Transforms/Utils/CMakeLists.txt
index 58eddb7..3462f89 100644
--- a/llvm/lib/Transforms/Utils/CMakeLists.txt
+++ b/llvm/lib/Transforms/Utils/CMakeLists.txt
@@ -53,6 +53,7 @@ add_llvm_component_library(LLVMTransformUtils
   MemoryTaggingSupport.cpp
   Mem2Reg.cpp
   MetaRenamer.cpp
+  MisExpect.cpp
   ModuleUtils.cpp
   NameAnonGlobals.cpp
   PredicateInfo.cpp
diff --git a/llvm/lib/Transforms/Utils/MisExpect.cpp b/llvm/lib/Transforms/Utils/MisExpect.cpp
new file mode 100644
index 0000000..044a2ca
--- /dev/null
+++ b/llvm/lib/Transforms/Utils/MisExpect.cpp
@@ -0,0 +1,234 @@
+//===--- MisExpect.cpp - Check the use of llvm.expect with PGO data -------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This contains code to emit warnings for potentially incorrect usage of the
+// llvm.expect intrinsic. This utility extracts the threshold values from
+// metadata associated with the instrumented Branch or Switch instruction. The
+// threshold values are then used to determine if a warning should be emmited.
+//
+// MisExpect's implementation relies on two assumptions about how branch weights
+// are managed in LLVM.
+//
+// 1) Frontend profiling weights are always in place before llvm.expect is
+// lowered in LowerExpectIntrinsic.cpp. Frontend based instrumentation therefore
+// needs to extract the branch weights and then compare them to the weights
+// being added by the llvm.expect intrinsic lowering.
+//
+// 2) Sampling and IR based profiles will *only* have branch weight metadata
+// before profiling data is consulted if they are from a lowered llvm.expect
+// intrinsic. These profiles thus always extract the expected weights and then
+// compare them to the weights collected during profiling to determine if a
+// diagnostic message is warranted.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Utils/MisExpect.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Analysis/OptimizationRemarkEmitter.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/BranchProbability.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/FormatVariadic.h"
+#include <cstdint>
+#include <functional>
+#include <numeric>
+
+#define DEBUG_TYPE "misexpect"
+
+using namespace llvm;
+using namespace misexpect;
+
+namespace llvm {
+
+// Command line option to enable/disable the warning when profile data suggests
+// a mismatch with the use of the llvm.expect intrinsic
+static cl::opt<bool> PGOWarnMisExpect(
+    "pgo-warn-misexpect", cl::init(false), cl::Hidden,
+    cl::desc("Use this option to turn on/off "
+             "warnings about incorrect usage of llvm.expect intrinsics."));
+
+static cl::opt<unsigned> MisExpectTolerance(
+    "misexpect-tolerance", cl::init(0),
+    cl::desc("Prevents emiting diagnostics when profile counts are "
+             "within N% of the threshold.."));
+
+} // namespace llvm
+
+namespace {
+
+bool isMisExpectDiagEnabled(LLVMContext &Ctx) {
+  return PGOWarnMisExpect || Ctx.getMisExpectWarningRequested();
+}
+
+uint64_t getMisExpectTolerance(LLVMContext &Ctx) {
+  return std::max(static_cast<uint64_t>(MisExpectTolerance),
+                  Ctx.getDiagnosticsMisExpectTolerance());
+}
+
+Instruction *getInstCondition(Instruction *I) {
+  assert(I != nullptr && "MisExpect target Instruction cannot be nullptr");
+  Instruction *Ret = nullptr;
+  if (auto *B = dyn_cast<BranchInst>(I)) {
+    Ret = dyn_cast<Instruction>(B->getCondition());
+  }
+  // TODO: Find a way to resolve condition location for switches
+  // Using the condition of the switch seems to often resolve to an earlier
+  // point in the program, i.e. the calculation of the switch condition, rather
+  // than the switch's location in the source code. Thus, we should use the
+  // instruction to get source code locations rather than the condition to
+  // improve diagnostic output, such as the caret. If the same problem exists
+  // for branch instructions, then we should remove this function and directly
+  // use the instruction
+  //
+  else if (auto *S = dyn_cast<SwitchInst>(I)) {
+    Ret = dyn_cast<Instruction>(S->getCondition());
+  }
+  return Ret ? Ret : I;
+}
+
+void emitMisexpectDiagnostic(Instruction *I, LLVMContext &Ctx,
+                             uint64_t ProfCount, uint64_t TotalCount) {
+  double PercentageCorrect = (double)ProfCount / TotalCount;
+  auto PerString =
+      formatv("{0:P} ({1} / {2})", PercentageCorrect, ProfCount, TotalCount);
+  auto RemStr = formatv(
+      "Potential performance regression from use of the llvm.expect intrinsic: "
+      "Annotation was correct on {0} of profiled executions.",
+      PerString);
+  Twine Msg(PerString);
+  Instruction *Cond = getInstCondition(I);
+  if (isMisExpectDiagEnabled(Ctx))
+    Ctx.diagnose(DiagnosticInfoMisExpect(Cond, Msg));
+  OptimizationRemarkEmitter ORE(I->getParent()->getParent());
+  ORE.emit(OptimizationRemark(DEBUG_TYPE, "misexpect", Cond) << RemStr.str());
+}
+
+} // namespace
+
+namespace llvm {
+namespace misexpect {
+
+// Helper function to extract branch weights into a vector
+Optional<SmallVector<uint32_t, 4>> extractWeights(Instruction *I,
+                                                  LLVMContext &Ctx) {
+  assert(I && "MisExpect::extractWeights given invalid pointer");
+
+  auto *ProfileData = I->getMetadata(LLVMContext::MD_prof);
+  if (!ProfileData)
+    return None;
+
+  unsigned NOps = ProfileData->getNumOperands();
+  if (NOps < 3)
+    return None;
+
+  auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0));
+  if (!ProfDataName || !ProfDataName->getString().equals("branch_weights"))
+    return None;
+
+  SmallVector<uint32_t, 4> Weights(NOps - 1);
+  for (unsigned Idx = 1; Idx < NOps; Idx++) {
+    ConstantInt *Value =
+        mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(Idx));
+    uint32_t V = Value->getZExtValue();
+    Weights[Idx - 1] = V;
+  }
+
+  return Weights;
+}
+
+void verifyMisExpect(Instruction &I, ArrayRef<uint32_t> RealWeights,
+                     ArrayRef<uint32_t> ExpectedWeights) {
+  // To determine if we emit a diagnostic, we need to compare the branch weights
+  // from the profile to those added by the llvm.expect intrinsic.
+  // So first, we extract the "likely" and "unlikely" weights from
+  // ExpectedWeights And determine the correct weight in the profile to compare
+  // against.
+  uint64_t LikelyBranchWeight = 0,
+           UnlikelyBranchWeight = std::numeric_limits<uint32_t>::max();
+  size_t MaxIndex = 0;
+  for (size_t Idx = 0, End = ExpectedWeights.size(); Idx < End; Idx++) {
+    uint32_t V = ExpectedWeights[Idx];
+    if (LikelyBranchWeight < V) {
+      LikelyBranchWeight = V;
+      MaxIndex = Idx;
+    }
+    if (UnlikelyBranchWeight > V) {
+      UnlikelyBranchWeight = V;
+    }
+  }
+
+  const uint64_t ProfiledWeight = RealWeights[MaxIndex];
+  const uint64_t RealWeightsTotal =
+      std::accumulate(RealWeights.begin(), RealWeights.end(), (uint64_t)0,
+                      std::plus<uint64_t>());
+  const uint64_t NumUnlikelyTargets = RealWeights.size() - 1;
+
+  const uint64_t TotalBranchWeight =
+      LikelyBranchWeight + (UnlikelyBranchWeight * NumUnlikelyTargets);
+
+  // To determine our threshold value we need to obtain the branch probability
+  // for the weights added by llvm.expect and use that proportion to calculate
+  // our threshold based on the collected profile data.
+  const llvm::BranchProbability LikelyProbablilty(LikelyBranchWeight,
+                                                  TotalBranchWeight);
+  uint64_t ScaledThreshold = LikelyProbablilty.scale(RealWeightsTotal);
+
+  // clamp tolerance range to [0, 100)
+  auto Tolerance = getMisExpectTolerance(I.getContext());
+  if (Tolerance < 0)
+    Tolerance = 0;
+  else if (MisExpectTolerance > 99)
+    Tolerance = 99;
+
+  // Allow users to relax checking by N%  i.e., if they use a 5% tolerance,
+  // then we check against 0.95*ScaledThreshold
+  if (Tolerance > 0)
+    ScaledThreshold *= (1.0 - Tolerance / 100.0);
+
+  // When the profile weight is below the threshold, we emit the diagnostic
+  if (ProfiledWeight < ScaledThreshold)
+    emitMisexpectDiagnostic(&I, I.getContext(), ProfiledWeight,
+                            RealWeightsTotal);
+}
+
+void checkBackendInstrumentation(Instruction &I,
+                                 const ArrayRef<uint32_t> RealWeights) {
+  auto ExpectedWeightsOpt = extractWeights(&I, I.getContext());
+  if (!ExpectedWeightsOpt.hasValue())
+    return;
+  auto ExpectedWeights = ExpectedWeightsOpt.getValue();
+  verifyMisExpect(I, RealWeights, ExpectedWeights);
+}
+
+void checkFrontendInstrumentation(Instruction &I,
+                                  const ArrayRef<uint32_t> ExpectedWeights) {
+  auto RealWeightsOpt = extractWeights(&I, I.getContext());
+  if (!RealWeightsOpt.hasValue())
+    return;
+  auto RealWeights = RealWeightsOpt.getValue();
+  verifyMisExpect(I, RealWeights, ExpectedWeights);
+}
+
+void checkExpectAnnotations(Instruction &I,
+                            const ArrayRef<uint32_t> ExistingWeights,
+                            bool IsFrontendInstr) {
+  if (IsFrontendInstr) {
+    checkFrontendInstrumentation(I, ExistingWeights);
+  } else {
+    checkBackendInstrumentation(I, ExistingWeights);
+  }
+}
+
+} // namespace misexpect
+} // namespace llvm
+#undef DEBUG_TYPE