[llvm] [SampleProfile] Skip counting mismatched weak symbols during profile loading (PR #185514)

[via llvm-commits]

https://github.com/HighW4y2H3ll created https://github.com/llvm/llvm-project/pull/185514

Weak symbols may be overridden during linking, and this may cause profile mismatch when compiling the weak symbols, while the profile was created based on the overriding function. Skip counting the weak symbol while checking the mismatched function profiles to avoid false alarm on rejecting legit profiles.

>From 12a0c4a48cb4445bd1d374a0b9916ef3d43ddcc8 Mon Sep 17 00:00:00 2001
From: h2h <h2h at meta.com>
Date: Mon, 9 Mar 2026 13:41:39 -0700
Subject: [PATCH] [SampleProfile] Skip counting mismatched weak symbols during
 profile loading

---
 .../Transforms/Utils/SampleProfileLoaderBaseImpl.h   |  9 +++++++++
 llvm/lib/Transforms/IPO/SampleProfile.cpp            |  3 ++-
 .../Inputs/weak-symbol-profile-mismatch.prof         |  5 +++++
 .../SampleProfile/weak-symbol-profile-mismatch.ll    | 12 ++++++++++++
 4 files changed, 28 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/SampleProfile/Inputs/weak-symbol-profile-mismatch.prof
 create mode 100644 llvm/test/Transforms/SampleProfile/weak-symbol-profile-mismatch.ll

diff --git a/llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h b/llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h
index fc2b13356350d..d1af683edaf0d 100644
--- a/llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h
+++ b/llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h
@@ -87,6 +87,7 @@ template <> struct IRTraits<BasicBlock> {
 // SampleProfileProber.
 class PseudoProbeManager {
   DenseMap<uint64_t, PseudoProbeDescriptor> GUIDToProbeDescMap;
+  DenseSet<uint64_t> GUIDIsWeakSymbol;
 
 public:
   PseudoProbeManager(const Module &M) {
@@ -98,7 +99,11 @@ class PseudoProbeManager {
                         ->getZExtValue();
         auto Hash = mdconst::dyn_extract<ConstantInt>(MD->getOperand(1))
                         ->getZExtValue();
+        auto Name = cast<MDString>(MD->getOperand(2))->getString();
         GUIDToProbeDescMap.try_emplace(GUID, PseudoProbeDescriptor(GUID, Hash));
+        if (auto *Func = M.getFunction(Name))
+          if (Func->hasWeakLinkage() || Func->hasExternalWeakLinkage())
+            GUIDIsWeakSymbol.insert(GUID);
       }
     }
   }
@@ -117,6 +122,10 @@ class PseudoProbeManager {
         FunctionSamples::getCanonicalFnName(F)));
   }
 
+  bool probeFromWeakSymbol(uint64_t GUID) const {
+    return GUIDIsWeakSymbol.count(GUID);
+  }
+
   bool profileIsHashMismatched(const PseudoProbeDescriptor &FuncDesc,
                                const FunctionSamples &Samples) const {
     return FuncDesc.getFunctionHash() != Samples.getFunctionHash();
diff --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index 8e76b79cc8f87..b2abc735046d2 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -2115,7 +2115,8 @@ bool SampleProfileLoader::rejectHighStalenessProfile(
       continue;
 
     TotalHotFunc++;
-    if (ProbeManager->profileIsHashMismatched(*FuncDesc, FS))
+    if (ProbeManager->profileIsHashMismatched(*FuncDesc, FS) &&
+        !ProbeManager->probeFromWeakSymbol(FS.getGUID()))
       NumMismatchedFunc++;
   }
   // Make sure that the num of selected function is not too small to distinguish
diff --git a/llvm/test/Transforms/SampleProfile/Inputs/weak-symbol-profile-mismatch.prof b/llvm/test/Transforms/SampleProfile/Inputs/weak-symbol-profile-mismatch.prof
new file mode 100644
index 0000000000000..60bfa0c35f161
--- /dev/null
+++ b/llvm/test/Transforms/SampleProfile/Inputs/weak-symbol-profile-mismatch.prof
@@ -0,0 +1,5 @@
+main:66508:408
+ 1: 408
+ 2: 0
+ 3: 408
+ !CFGChecksum: 3735928559
diff --git a/llvm/test/Transforms/SampleProfile/weak-symbol-profile-mismatch.ll b/llvm/test/Transforms/SampleProfile/weak-symbol-profile-mismatch.ll
new file mode 100644
index 0000000000000..9346544dac63d
--- /dev/null
+++ b/llvm/test/Transforms/SampleProfile/weak-symbol-profile-mismatch.ll
@@ -0,0 +1,12 @@
+; RUN: opt %s -passes='sample-profile' -sample-profile-file=%S/Inputs/weak-symbol-profile-mismatch.prof -min-functions-for-staleness-error=1 2>&1 | FileCheck %s
+
+; CHECK-NOT: Pseudo-probe-based profile requires SampleProfileProbePass
+; Function Attrs: noinline nounwind optnone uwtable
+define weak dso_local void @main() #0 align 8 {
+  ret void
+}
+
+attributes #0 = { noinline nounwind optnone uwtable }
+
+!llvm.pseudo_probe_desc = !{!1}
+!1 = !{i64 -2624081020897602054, i64 123456, !"main"}