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

[Davide Italiano via llvm-commits]

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

>From 92ca8e33db6cb06dc1f5f5f6a322e663cf5c91d3 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 | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp b/llvm/lib/Transforms/IPO/SampleProfileProbe.cpp
index 56d711c588b0755..06320720ed62c7a 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"
@@ -221,12 +222,26 @@ void SampleProfileProber::computeProbeIdForBlocks() {
 }
 
 void SampleProfileProber::computeProbeIdForCallsites() {
+  LLVMContext &Ctx = F->getContext();
+  Module *M = F->getParent();
+
   for (auto &BB : *F) {
     for (auto &I : BB) {
       if (!isa<CallBase>(I))
         continue;
       if (isa<IntrinsicInst>(&I))
         continue;
+
+      // 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 be ignored.
+      if (LastProbeId >= 0xFFFF) {
+        std::string Msg = "Pseudo instrumentation incomplete for " +
+                          std::string(F->getName()) + " because it's too large";
+        Ctx.diagnose(DiagnosticInfoSampleProfile(M->getName().data(), Msg));
+        break;
+      }
+
       CallProbeIds[&I] = ++LastProbeId;
     }
   }