[llvm] [PGO] Skip optimizing probes that don't fit. (PR #70875)

[Davide Italiano via llvm-commits]

https://github.com/dcci created https://github.com/llvm/llvm-project/pull/70875

The discriminator can only pack 16 bits, so anything exceeding that value will cause the packing code to crash. Emit a diagnostic and skip the optimization instead.

>From 296e79df36c71d984447ea2a4296a53f29016958 Mon Sep 17 00:00:00 2001
From: Davide Italiano <davidino at fb.com>
Date: Tue, 31 Oct 2023 16:20:51 -0700
Subject: [PATCH] [PGO] Skip optimizing probes that don't fit.

The discriminator can only pack 16 bits, so anything exceeding
that value will cause the packing code to crash. Emit a diagnostic
and skip the optimization instead.
---
 llvm/lib/Transforms/IPO/SampleProfileProbe.cpp | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp b/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp
index 56d711c588b0755..6092c4dd8922400 100644
--- a/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp
@@ -18,6 +18,7 @@
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/IntrinsicInst.h"
@@ -244,7 +245,8 @@ uint32_t SampleProfileProber::getCallsiteId(const Instruction *Call) const {
 
 void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
   Module *M = F.getParent();
-  MDBuilder MDB(F.getContext());
+  LLVMContext &Ctx = M->getContext();
+  MDBuilder MDB(Ctx);
   // Since the GUID from probe desc and inline stack are computed seperately, we
   // need to make sure their names are consistent, so here also use the name
   // from debug info.
@@ -324,6 +326,19 @@ void SampleProfileProber::instrumentOneFunc(Function &F, TargetMachine *TM) {
   for (auto &I : CallProbeIds) {
     auto *Call = I.first;
     uint32_t Index = I.second;
+
+    // There's a limit on the number of probe indices that can be optimized.
+    // The current implementation uses the lower 16 bits of the discriminator
+    // so anything larger than 0xFFFF will ignored.
+    if (Index > 0xFFFF) {
+      std::string Msg = "Pseudo instrumentation incomplete for "
+                        + std::string(F.getName())
+                        + " because it's too large";
+      Ctx.diagnose(DiagnosticInfoSampleProfile(M->getName().data(), Msg,
+                                               DS_Warning));
+      continue;
+    }
+
     uint32_t Type = cast<CallBase>(Call)->getCalledFunction()
                         ? (uint32_t)PseudoProbeType::DirectCall
                         : (uint32_t)PseudoProbeType::IndirectCall;