[lld] [llvm] [llvm][lld] Respect temporal profile weights in balanced partitioning (PR #213837)
Karim Alweheshy via llvm-commits
llvm-commits at lists.llvm.org
Tue Aug 4 22:59:18 PDT 2026
https://github.com/karim-alweheshy updated https://github.com/llvm/llvm-project/pull/213837
>From 49d4ed229ba597d32f60bc005eaf15c872699a30 Mon Sep 17 00:00:00 2001
From: Karim Alweheshy <karim.alweheshy at gmail.com>
Date: Tue, 4 Aug 2026 06:56:44 +0200
Subject: [PATCH 1/3] Respect temporal trace weights in balanced partitioning
---
.../lld/Common/BPSectionOrdererBase.inc | 55 ++++++-
.../bp-section-orderer-weighted-profile.s | 149 ++++++++++++++++++
.../llvm/Support/BalancedPartitioning.h | 30 +++-
llvm/lib/ProfileData/InstrProf.cpp | 30 +++-
llvm/lib/Support/BalancedPartitioning.cpp | 36 ++++-
.../ProfileData/BPFunctionNodeTest.cpp | 38 +++--
.../Support/BalancedPartitioningTest.cpp | 39 +++++
7 files changed, 342 insertions(+), 35 deletions(-)
create mode 100644 lld/test/MachO/bp-section-orderer-weighted-profile.s
diff --git a/lld/include/lld/Common/BPSectionOrdererBase.inc b/lld/include/lld/Common/BPSectionOrdererBase.inc
index b2e87d6c240fc..78d1ab79a3bb9 100644
--- a/lld/include/lld/Common/BPSectionOrdererBase.inc
+++ b/lld/include/lld/Common/BPSectionOrdererBase.inc
@@ -34,6 +34,7 @@
#include "llvm/ADT/Twine.h"
#include "llvm/ProfileData/InstrProfReader.h"
#include "llvm/Support/BalancedPartitioning.h"
+#include "llvm/Support/MathExtras.h"
#include "llvm/Support/TimeProfiler.h"
#include "llvm/Support/VirtualFileSystem.h"
#include <memory>
@@ -172,9 +173,13 @@ auto BPOrderer<D>::computeOrder(
BPFunctionNode::UtilityNodeT maxUN = 0;
MapVector<unsigned, UtilityNodes> startupSectionIdxUNs;
+ DenseMap<unsigned, SmallVector<BPFunctionNode::UtilityNodeWeightT>>
+ startupSectionIdxUNWeights;
// Used to define the initial order for startup functions.
DenseMap<unsigned, size_t> sectionIdxToTimestamp;
uint64_t numProfiledFunctionRefs = 0, numResolvedProfiledFunctionRefs = 0;
+ // Weighted verbose counts saturate instead of wrapping on malformed or
+ // artificially large profile weights.
DenseSet<CachedHashStringRef> uniqueProfiledFunctions;
DenseSet<CachedHashStringRef> uniqueResolvedProfiledFunctions;
std::unique_ptr<InstrProfReader> reader;
@@ -195,6 +200,10 @@ auto BPOrderer<D>::computeOrder(
MapVector<unsigned, BPFunctionNode::UtilityNodeT> sectionIdxToFirstUN;
for (size_t traceIdx = 0; traceIdx < traces.size(); traceIdx++) {
+ const uint64_t traceWeight = traces[traceIdx].Weight;
+ // A zero-weight trace is equivalent to no copies of the trace.
+ if (traceWeight == 0)
+ continue;
uint64_t currentSize = 0, cutoffSize = 1;
size_t cutoffTimestamp = 1;
auto &trace = traces[traceIdx].FunctionNameRefs;
@@ -203,7 +212,8 @@ auto BPOrderer<D>::computeOrder(
reader->getSymtab().getFuncOrVarName(trace[timestamp]));
parsedFuncName = lld::utils::getRootSymbol(parsedFuncName);
if (verbose) {
- ++numProfiledFunctionRefs;
+ numProfiledFunctionRefs =
+ SaturatingAdd(numProfiledFunctionRefs, traceWeight);
uniqueProfiledFunctions.insert(CachedHashStringRef(parsedFuncName));
}
@@ -212,7 +222,8 @@ auto BPOrderer<D>::computeOrder(
if (sectionIdxsIt == rootSymbolToSectionIdxs.end())
continue;
if (verbose) {
- ++numResolvedProfiledFunctionRefs;
+ numResolvedProfiledFunctionRefs =
+ SaturatingAdd(numResolvedProfiledFunctionRefs, traceWeight);
uniqueResolvedProfiledFunctions.insert(
CachedHashStringRef(parsedFuncName));
}
@@ -239,8 +250,10 @@ auto BPOrderer<D>::computeOrder(
sectionIdxToFirstUN.try_emplace(sectionIdx, maxUN);
}
for (auto &[sectionIdx, firstUN] : sectionIdxToFirstUN)
- for (auto un = firstUN; un <= maxUN; ++un)
+ for (auto un = firstUN; un <= maxUN; ++un) {
startupSectionIdxUNs[sectionIdx].push_back(un);
+ startupSectionIdxUNWeights[sectionIdx].push_back(traceWeight);
+ }
++maxUN;
sectionIdxToFirstUN.clear();
}
@@ -333,9 +346,27 @@ auto BPOrderer<D>::computeOrder(
for (auto &[sectionIdx, compressionUns] :
unsForStartupFunctionCompression) {
auto &uns = startupSectionIdxUNs[sectionIdx];
+ auto &weights = startupSectionIdxUNWeights[sectionIdx];
uns.append(compressionUns);
- llvm::sort(uns);
- uns.erase(llvm::unique(uns), uns.end());
+ weights.append(compressionUns.size(), 1);
+
+ SmallVector<std::pair<BPFunctionNode::UtilityNodeT,
+ BPFunctionNode::UtilityNodeWeightT>>
+ weightedUNs;
+ for (auto [un, weight] : llvm::zip(uns, weights))
+ weightedUNs.emplace_back(un, weight);
+ llvm::sort(weightedUNs);
+ weightedUNs.erase(llvm::unique(weightedUNs,
+ [](const auto &left, const auto &right) {
+ return left.first == right.first;
+ }),
+ weightedUNs.end());
+ uns.clear();
+ weights.clear();
+ for (auto [un, weight] : weightedUNs) {
+ uns.push_back(un);
+ weights.push_back(weight);
+ }
}
}
@@ -357,7 +388,8 @@ auto BPOrderer<D>::computeOrder(
std::vector<BPFunctionNode> nodesForStartup;
for (auto &[sectionIdx, uns] : startupSectionIdxUNs)
- nodesForStartup.emplace_back(sectionIdx, uns);
+ nodesForStartup.emplace_back(sectionIdx, uns,
+ startupSectionIdxUNWeights[sectionIdx]);
// Use the first timestamp to define the initial order for startup nodes.
llvm::sort(nodesForStartup, [§ionIdxToTimestamp](auto &L, auto &R) {
@@ -459,9 +491,14 @@ auto BPOrderer<D>::computeOrder(
// The area under the curve F where F(t) is the total number of page
// faults at step t.
- unsigned area = 0;
+ // Saturation makes an oversized weighted profile deterministic rather
+ // than allowing the diagnostic-only aggregate to wrap.
+ uint64_t area = 0;
for (auto &trace : reader->getTemporalProfTraces()) {
+ if (trace.Weight == 0)
+ continue;
SmallSet<uint64_t, 0> touchedPages;
+ uint64_t traceArea = 0;
for (unsigned step = 0; step < trace.FunctionNameRefs.size(); step++) {
auto traceId = trace.FunctionNameRefs[step];
auto [Filename, ParsedFuncName] =
@@ -473,8 +510,10 @@ auto BPOrderer<D>::computeOrder(
for (uint64_t i = firstPage; i <= lastPage; i++)
touchedPages.insert(i);
}
- area += touchedPages.size();
+ traceArea = SaturatingAdd(traceArea,
+ static_cast<uint64_t>(touchedPages.size()));
}
+ area = SaturatingMultiplyAdd(traceArea, trace.Weight, area);
}
dbgs() << "Total area under the page fault curve: " << (float)area
<< "\n";
diff --git a/lld/test/MachO/bp-section-orderer-weighted-profile.s b/lld/test/MachO/bp-section-orderer-weighted-profile.s
new file mode 100644
index 0000000000000..1907a697c7fa7
--- /dev/null
+++ b/lld/test/MachO/bp-section-orderer-weighted-profile.s
@@ -0,0 +1,149 @@
+# REQUIRES: aarch64
+
+# RUN: rm -rf %t && split-file %s %t
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %t/input.s -o %t/input.o
+# RUN: llvm-profdata merge %t/weighted.proftext %t/functions.proftext -o %t/weighted.profdata
+# RUN: llvm-profdata merge %t/replicated.proftext %t/functions.proftext -o %t/replicated.profdata
+# RUN: llvm-profdata merge %t/unweighted.proftext %t/functions.proftext -o %t/unweighted.profdata
+
+# A weight of 10 is equivalent to ten copies of the same trace.
+# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/weighted.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - | FileCheck %s --check-prefix=WEIGHTED
+# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/replicated.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - | FileCheck %s --check-prefix=WEIGHTED
+# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/unweighted.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - | FileCheck %s --check-prefix=UNWEIGHTED
+
+# WEIGHTED: D
+# WEIGHTED-NEXT: A
+# WEIGHTED-NEXT: B
+# WEIGHTED-NEXT: C
+# WEIGHTED-NEXT: E
+# WEIGHTED-NEXT: F
+# UNWEIGHTED: A
+# UNWEIGHTED-NEXT: B
+# UNWEIGHTED-NEXT: D
+# UNWEIGHTED-NEXT: C
+# UNWEIGHTED-NEXT: E
+# UNWEIGHTED-NEXT: F
+
+# Verbose trace-reference counts and the page-fault area also respect weights.
+# RUN: %lld -arch arm64 -lSystem -e _main -o %t/weighted.out %t/input.o --irpgo-profile=%t/weighted.profdata --bp-startup-sort=function --verbose-bp-section-orderer 2>&1 | FileCheck %s --check-prefix=WEIGHTED-STATS
+# RUN: %lld -arch arm64 -lSystem -e _main -o %t/replicated.out %t/input.o --irpgo-profile=%t/replicated.profdata --bp-startup-sort=function --verbose-bp-section-orderer 2>&1 | FileCheck %s --check-prefix=WEIGHTED-STATS
+# WEIGHTED-STATS: Temporal profile function references: 66 / 66 resolved (6 / 6 unique)
+# WEIGHTED-STATS: Total area under the page fault curve: 6.600000e+01
+
+#--- input.s
+.text
+.globl _main, A, B, C, D, E, F
+_main:
+ ret
+A:
+ ret
+B:
+ ret
+C:
+ ret
+D:
+ ret
+E:
+ ret
+F:
+ ret
+.subsections_via_symbols
+
+#--- weighted.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
+2
+# Trace Stream Size
+11
+# Weight
+10
+A, B, C, D, E, F
+# Weight
+1
+A, D, B, C, E, F
+
+#--- replicated.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
+11
+# Trace Stream Size
+11
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, D, B, C, E, F
+
+#--- unweighted.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
+2
+# Trace Stream Size
+2
+# Weight
+1
+A, B, C, D, E, F
+# Weight
+1
+A, D, B, C, E, F
+
+#--- functions.proftext
+:ir
+A
+# Func Hash
+1
+# Num Counters
+1
+# Counter Values
+1
+B
+2
+1
+1
+C
+3
+1
+1
+D
+4
+1
+1
+E
+5
+1
+1
+F
+6
+1
+1
diff --git a/llvm/include/llvm/Support/BalancedPartitioning.h b/llvm/include/llvm/Support/BalancedPartitioning.h
index 05307d74c209c..355155f3d512c 100644
--- a/llvm/include/llvm/Support/BalancedPartitioning.h
+++ b/llvm/include/llvm/Support/BalancedPartitioning.h
@@ -41,9 +41,12 @@
#include "raw_ostream.h"
#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Compiler.h"
#include <atomic>
+#include <cassert>
#include <condition_variable>
#include <mutex>
#include <random>
@@ -60,10 +63,30 @@ class BPFunctionNode {
public:
using IDT = uint64_t;
using UtilityNodeT = uint32_t;
+ using UtilityNodeWeightT = uint64_t;
/// \param UtilityNodes the set of utility nodes (must be unique'd)
BPFunctionNode(IDT Id, ArrayRef<UtilityNodeT> UtilityNodes)
- : Id(Id), UtilityNodes(UtilityNodes) {}
+ : Id(Id), UtilityNodes(UtilityNodes),
+ UtilityNodeWeights(UtilityNodes.size(), 1) {}
+
+ /// \param UtilityNodes the set of utility nodes (must be unique'd)
+ /// \param UtilityNodeWeights the weight of each corresponding utility node.
+ /// All occurrences of a utility node must have the same weight. Zero-weight
+ /// utility nodes are ignored.
+ BPFunctionNode(IDT Id, ArrayRef<UtilityNodeT> UtilityNodes,
+ ArrayRef<UtilityNodeWeightT> UtilityNodeWeights)
+ : Id(Id) {
+ assert(UtilityNodes.size() == UtilityNodeWeights.size() &&
+ "each utility node must have a weight");
+ for (auto [UtilityNode, Weight] :
+ llvm::zip(UtilityNodes, UtilityNodeWeights)) {
+ if (Weight == 0)
+ continue;
+ this->UtilityNodes.push_back(UtilityNode);
+ this->UtilityNodeWeights.push_back(Weight);
+ }
+ }
/// The ID of this node
IDT Id;
@@ -73,6 +96,9 @@ class BPFunctionNode {
protected:
/// The list of utility nodes associated with this node
SmallVector<UtilityNodeT, 4> UtilityNodes;
+ /// The weight of each corresponding utility node. A utility node's cost is
+ /// multiplied by its weight without materializing duplicate utility nodes.
+ SmallVector<UtilityNodeWeightT, 4> UtilityNodeWeights;
/// The bucket assigned by balanced partitioning
std::optional<unsigned> Bucket;
/// The index of the input order of the FunctionNodes
@@ -189,6 +215,8 @@ class BalancedPartitioning {
float CachedGainRL;
/// Whether \p CachedGainLR and \p CachedGainRL are valid
bool CachedGainIsValid = false;
+ /// The number of times this utility contributes to the objective
+ BPFunctionNode::UtilityNodeWeightT Weight = 1;
};
protected:
diff --git a/llvm/lib/ProfileData/InstrProf.cpp b/llvm/lib/ProfileData/InstrProf.cpp
index 692786a146d46..49cec36ec1bb2 100644
--- a/llvm/lib/ProfileData/InstrProf.cpp
+++ b/llvm/lib/ProfileData/InstrProf.cpp
@@ -1153,9 +1153,11 @@ void TemporalProfTraceTy::createBPFunctionNodes(
DenseMap<IDT, size_t> IdToFirstTimestamp;
DenseMap<IDT, UtilityNodeT> IdToFirstUN;
DenseMap<IDT, SmallVector<UtilityNodeT>> IdToUNs;
- // TODO: We need to use the Trace.Weight field to give more weight to more
- // important utilities
+ DenseMap<IDT, SmallVector<BPFunctionNode::UtilityNodeWeightT>> IdToUNWeights;
for (auto &Trace : Traces) {
+ // A zero-weight trace is equivalent to no copies of the trace.
+ if (Trace.Weight == 0)
+ continue;
size_t CutoffTimestamp = 1;
for (size_t Timestamp = 0; Timestamp < Trace.FunctionNameRefs.size();
Timestamp++) {
@@ -1170,8 +1172,10 @@ void TemporalProfTraceTy::createBPFunctionNodes(
IdToFirstUN.try_emplace(Id, MaxUN);
}
for (auto &[Id, FirstUN] : IdToFirstUN)
- for (auto UN = FirstUN; UN <= MaxUN; ++UN)
+ for (auto UN = FirstUN; UN <= MaxUN; ++UN) {
IdToUNs[Id].push_back(UN);
+ IdToUNWeights[Id].push_back(Trace.Weight);
+ }
++MaxUN;
IdToFirstUN.clear();
}
@@ -1183,15 +1187,25 @@ void TemporalProfTraceTy::createBPFunctionNodes(
++UNFrequency[UN];
// Filter out utility nodes that are too infrequent or too prevalent to make
// BalancedPartitioning more effective.
- for (auto &[Id, UNs] : IdToUNs)
- llvm::erase_if(UNs, [&](auto &UN) {
+ for (auto &[Id, UNs] : IdToUNs) {
+ auto &Weights = IdToUNWeights[Id];
+ size_t WriteIndex = 0;
+ for (size_t ReadIndex = 0; ReadIndex < UNs.size(); ++ReadIndex) {
+ auto UN = UNs[ReadIndex];
unsigned Freq = UNFrequency[UN];
- return Freq <= 1 || 2 * Freq > IdToUNs.size();
- });
+ if (Freq <= 1 || 2 * Freq > IdToUNs.size())
+ continue;
+ UNs[WriteIndex] = UN;
+ Weights[WriteIndex] = Weights[ReadIndex];
+ ++WriteIndex;
+ }
+ UNs.resize(WriteIndex);
+ Weights.resize(WriteIndex);
+ }
}
for (auto &[Id, UNs] : IdToUNs)
- Nodes.emplace_back(Id, UNs);
+ Nodes.emplace_back(Id, UNs, IdToUNWeights[Id]);
// Since BalancedPartitioning is sensitive to the initial order, we explicitly
// order nodes by their earliest timestamp.
diff --git a/llvm/lib/Support/BalancedPartitioning.cpp b/llvm/lib/Support/BalancedPartitioning.cpp
index 2ae20e96845a2..c18ec77e2a3a3 100644
--- a/llvm/lib/Support/BalancedPartitioning.cpp
+++ b/llvm/lib/Support/BalancedPartitioning.cpp
@@ -175,11 +175,20 @@ void BalancedPartitioning::runIterations(const FunctionNodeRange Nodes,
++UtilityNodeIndex[UN];
// Remove utility nodes if they have just one edge or are connected to all
// functions
- for (auto &N : Nodes)
- llvm::erase_if(N.UtilityNodes, [&](auto &UN) {
+ for (auto &N : Nodes) {
+ size_t WriteIndex = 0;
+ for (size_t ReadIndex = 0; ReadIndex < N.UtilityNodes.size(); ++ReadIndex) {
+ auto UN = N.UtilityNodes[ReadIndex];
unsigned UNI = UtilityNodeIndex[UN];
- return UNI == 1 || UNI == NumNodes;
- });
+ if (UNI == 1 || UNI == NumNodes)
+ continue;
+ N.UtilityNodes[WriteIndex] = UN;
+ N.UtilityNodeWeights[WriteIndex] = N.UtilityNodeWeights[ReadIndex];
+ ++WriteIndex;
+ }
+ N.UtilityNodes.resize(WriteIndex);
+ N.UtilityNodeWeights.resize(WriteIndex);
+ }
// Renumber utility nodes so they can be used to index into Signatures
UtilityNodeIndex.clear();
@@ -190,12 +199,17 @@ void BalancedPartitioning::runIterations(const FunctionNodeRange Nodes,
// Initialize signatures
SignaturesT Signatures(/*Size=*/UtilityNodeIndex.size());
for (auto &N : Nodes) {
- for (auto &UN : N.UtilityNodes) {
+ for (auto [UN, Weight] : llvm::zip(N.UtilityNodes, N.UtilityNodeWeights)) {
assert(UN < Signatures.size());
+ auto &Signature = Signatures[UN];
+ if (Signature.LeftCount != 0 || Signature.RightCount != 0)
+ assert(Signature.Weight == Weight &&
+ "a utility node must have a consistent weight");
+ Signature.Weight = Weight;
if (N.Bucket == LeftBucket) {
- Signatures[UN].LeftCount++;
+ Signature.LeftCount++;
} else {
- Signatures[UN].RightCount++;
+ Signature.RightCount++;
}
}
}
@@ -227,6 +241,14 @@ unsigned BalancedPartitioning::runIteration(const FunctionNodeRange Nodes,
Signature.CachedGainLR = Cost - logCost(L - 1, R + 1);
if (R > 0)
Signature.CachedGainRL = Cost - logCost(L + 1, R - 1);
+ // Scaling a utility's move gain is mathematically equivalent to adding
+ // Weight distinct utility nodes with the same signature, but avoids
+ // expanding the graph. Weight == 1 deliberately takes no FP operation so
+ // unweighted inputs preserve their previous behavior.
+ if (Signature.Weight != 1) {
+ Signature.CachedGainLR *= static_cast<float>(Signature.Weight);
+ Signature.CachedGainRL *= static_cast<float>(Signature.Weight);
+ }
Signature.CachedGainIsValid = true;
}
diff --git a/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp b/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
index 24586b5aa31af..058dd8d17e355 100644
--- a/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
+++ b/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
@@ -24,10 +24,14 @@ void PrintTo(const BPFunctionNode &Node, std::ostream *OS) {
TEST(BPFunctionNodeTest, Basic) {
auto NodeIs = [](BPFunctionNode::IDT Id,
- ArrayRef<BPFunctionNode::UtilityNodeT> UNs) {
+ ArrayRef<BPFunctionNode::UtilityNodeT> UNs,
+ ArrayRef<BPFunctionNode::UtilityNodeWeightT> Weights) {
return AllOf(Field("Id", &BPFunctionNode::Id, Id),
Field("UtilityNodes", &BPFunctionNode::UtilityNodes,
- UnorderedElementsAreArray(UNs)));
+ UnorderedElementsAreArray(UNs)),
+ Field("UtilityNodeWeights",
+ &BPFunctionNode::UtilityNodeWeights,
+ UnorderedElementsAreArray(Weights)));
};
std::vector<BPFunctionNode> Nodes;
@@ -35,27 +39,39 @@ TEST(BPFunctionNodeTest, Basic) {
{TemporalProfTraceTy({0, 1, 2, 3})}, Nodes, /*RemoveOutlierUNs=*/false);
// Utility nodes that are too infrequent or too prevalent are filtered out.
EXPECT_THAT(Nodes,
- UnorderedElementsAre(NodeIs(0, {0, 1, 2}), NodeIs(1, {1, 2}),
- NodeIs(2, {2}), NodeIs(3, {2})));
+ UnorderedElementsAre(NodeIs(0, {0, 1, 2}, {1, 1, 1}),
+ NodeIs(1, {1, 2}, {1, 1}),
+ NodeIs(2, {2}, {1}), NodeIs(3, {2}, {1})));
Nodes.clear();
TemporalProfTraceTy::createBPFunctionNodes(
{TemporalProfTraceTy({0, 1, 2, 3, 4}), TemporalProfTraceTy({4, 2})},
Nodes, /*RemoveOutlierUNs=*/false);
- EXPECT_THAT(Nodes,
- UnorderedElementsAre(NodeIs(0, {0, 1, 2, 3}),
- NodeIs(1, {1, 2, 3}), NodeIs(2, {2, 3, 5}),
- NodeIs(3, {2, 3}), NodeIs(4, {3, 4, 5})));
+ EXPECT_THAT(Nodes, UnorderedElementsAre(NodeIs(0, {0, 1, 2, 3}, {1, 1, 1, 1}),
+ NodeIs(1, {1, 2, 3}, {1, 1, 1}),
+ NodeIs(2, {2, 3, 5}, {1, 1, 1}),
+ NodeIs(3, {2, 3}, {1, 1}),
+ NodeIs(4, {3, 4, 5}, {1, 1, 1})));
Nodes.clear();
TemporalProfTraceTy::createBPFunctionNodes(
{TemporalProfTraceTy({0, 1, 2, 3, 4}), TemporalProfTraceTy({4, 2})},
Nodes, /*RemoveOutlierUNs=*/true);
- EXPECT_THAT(Nodes, UnorderedElementsAre(NodeIs(0, {1}), NodeIs(1, {1}),
- NodeIs(2, {5}), NodeIs(3, {}),
- NodeIs(4, {5})));
+ EXPECT_THAT(Nodes,
+ UnorderedElementsAre(NodeIs(0, {1}, {1}), NodeIs(1, {1}, {1}),
+ NodeIs(2, {5}, {1}), NodeIs(3, {}, {}),
+ NodeIs(4, {5}, {1})));
+
+ Nodes.clear();
+ TemporalProfTraceTy::createBPFunctionNodes(
+ {TemporalProfTraceTy({0, 1, 2, 3}, 10), TemporalProfTraceTy({4, 5}, 0)},
+ Nodes, /*RemoveOutlierUNs=*/false);
+ EXPECT_THAT(Nodes,
+ UnorderedElementsAre(NodeIs(0, {0, 1, 2}, {10, 10, 10}),
+ NodeIs(1, {1, 2}, {10, 10}),
+ NodeIs(2, {2}, {10}), NodeIs(3, {2}, {10})));
}
} // end namespace llvm
diff --git a/llvm/unittests/Support/BalancedPartitioningTest.cpp b/llvm/unittests/Support/BalancedPartitioningTest.cpp
index ebe518a8e89ca..a781d7b5d405e 100644
--- a/llvm/unittests/Support/BalancedPartitioningTest.cpp
+++ b/llvm/unittests/Support/BalancedPartitioningTest.cpp
@@ -97,4 +97,43 @@ TEST_F(BalancedPartitioningTest, MoveGain) {
30.f);
}
+TEST_F(BalancedPartitioningTest, WeightedUtilitiesMatchReplication) {
+ Config.SkipProbability = 0;
+ Config.TaskSplitDepth = 0;
+
+ std::vector<BPFunctionNode> WeightedNodes = {
+ BPFunctionNode(0, {1}, {1}),
+ BPFunctionNode(1, {1, 2}, {1, 1}),
+ BPFunctionNode(2, {0, 1, 2}, {10, 1, 1}),
+ BPFunctionNode(3, {0, 2}, {10, 1}),
+ };
+ std::vector<BPFunctionNode> ReplicatedNodes = {
+ BPFunctionNode(0, {10}),
+ BPFunctionNode(1, {10, 11}),
+ BPFunctionNode(2, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}),
+ BPFunctionNode(3, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 11}),
+ };
+ std::vector<BPFunctionNode> UnweightedNodes = {
+ BPFunctionNode(0, {1}),
+ BPFunctionNode(1, {1, 2}),
+ BPFunctionNode(2, {0, 1, 2}),
+ BPFunctionNode(3, {0, 2}),
+ };
+ std::vector<BPFunctionNode> WeightOneNodes = {
+ BPFunctionNode(0, {1}, {1}),
+ BPFunctionNode(1, {1, 2}, {1, 1}),
+ BPFunctionNode(2, {0, 1, 2}, {1, 1, 1}),
+ BPFunctionNode(3, {0, 2}, {1, 1}),
+ };
+
+ Bp.run(WeightedNodes);
+ Bp.run(ReplicatedNodes);
+ Bp.run(UnweightedNodes);
+ Bp.run(WeightOneNodes);
+
+ EXPECT_EQ(getIds(WeightedNodes), getIds(ReplicatedNodes));
+ EXPECT_NE(getIds(WeightedNodes), getIds(UnweightedNodes));
+ EXPECT_EQ(getIds(WeightOneNodes), getIds(UnweightedNodes));
+}
+
} // end namespace llvm
>From a3ef1d7fec2596da98a022f260e399a3a39688d6 Mon Sep 17 00:00:00 2001
From: Karim Alweheshy <karim.alweheshy at gmail.com>
Date: Tue, 4 Aug 2026 15:30:18 +0200
Subject: [PATCH 2/3] [Support][ProfileData][lld] Harden weighted utility
handling
---
.../lld/Common/BPSectionOrdererBase.inc | 46 +++----
.../ELF/bp-section-orderer-weighted-profile.s | 130 ++++++++++++++++++
.../bp-section-orderer-weighted-profile.s | 82 +++++------
.../llvm/Support/BalancedPartitioning.h | 55 ++++----
llvm/lib/ProfileData/InstrProf.cpp | 53 +++----
llvm/lib/Support/BalancedPartitioning.cpp | 56 ++++----
.../llvm-profdata/order-weighted-traces.test | 87 ++++++++++++
.../show-order-weight-saturation.proftext | 31 +++++
.../tools/llvm-profdata/show-order.proftext | 4 +-
llvm/tools/llvm-profdata/llvm-profdata.cpp | 33 ++---
.../ProfileData/BPFunctionNodeTest.cpp | 43 +++---
.../Support/BalancedPartitioningTest.cpp | 27 ++--
12 files changed, 446 insertions(+), 201 deletions(-)
create mode 100644 lld/test/ELF/bp-section-orderer-weighted-profile.s
create mode 100644 llvm/test/tools/llvm-profdata/order-weighted-traces.test
create mode 100644 llvm/test/tools/llvm-profdata/show-order-weight-saturation.proftext
diff --git a/lld/include/lld/Common/BPSectionOrdererBase.inc b/lld/include/lld/Common/BPSectionOrdererBase.inc
index 78d1ab79a3bb9..845a0bebcd36f 100644
--- a/lld/include/lld/Common/BPSectionOrdererBase.inc
+++ b/lld/include/lld/Common/BPSectionOrdererBase.inc
@@ -83,6 +83,7 @@ template <class D> struct BPOrderer {
} // namespace lld
using UtilityNodes = SmallVector<BPFunctionNode::UtilityNodeT>;
+using WeightedUtilityNodes = SmallVector<BPFunctionNode::WeightedUtilityNode>;
template <class D>
static SmallVector<std::pair<unsigned, UtilityNodes>> getUnsForCompression(
@@ -172,9 +173,7 @@ auto BPOrderer<D>::computeOrder(
sectionToIdx.try_emplace(isec, i);
BPFunctionNode::UtilityNodeT maxUN = 0;
- MapVector<unsigned, UtilityNodes> startupSectionIdxUNs;
- DenseMap<unsigned, SmallVector<BPFunctionNode::UtilityNodeWeightT>>
- startupSectionIdxUNWeights;
+ MapVector<unsigned, WeightedUtilityNodes> startupSectionIdxUNs;
// Used to define the initial order for startup functions.
DenseMap<unsigned, size_t> sectionIdxToTimestamp;
uint64_t numProfiledFunctionRefs = 0, numResolvedProfiledFunctionRefs = 0;
@@ -250,10 +249,8 @@ auto BPOrderer<D>::computeOrder(
sectionIdxToFirstUN.try_emplace(sectionIdx, maxUN);
}
for (auto &[sectionIdx, firstUN] : sectionIdxToFirstUN)
- for (auto un = firstUN; un <= maxUN; ++un) {
- startupSectionIdxUNs[sectionIdx].push_back(un);
- startupSectionIdxUNWeights[sectionIdx].push_back(traceWeight);
- }
+ for (auto un = firstUN; un <= maxUN; ++un)
+ startupSectionIdxUNs[sectionIdx].emplace_back(un, traceWeight);
++maxUN;
sectionIdxToFirstUN.clear();
}
@@ -346,27 +343,16 @@ auto BPOrderer<D>::computeOrder(
for (auto &[sectionIdx, compressionUns] :
unsForStartupFunctionCompression) {
auto &uns = startupSectionIdxUNs[sectionIdx];
- auto &weights = startupSectionIdxUNWeights[sectionIdx];
- uns.append(compressionUns);
- weights.append(compressionUns.size(), 1);
-
- SmallVector<std::pair<BPFunctionNode::UtilityNodeT,
- BPFunctionNode::UtilityNodeWeightT>>
- weightedUNs;
- for (auto [un, weight] : llvm::zip(uns, weights))
- weightedUNs.emplace_back(un, weight);
- llvm::sort(weightedUNs);
- weightedUNs.erase(llvm::unique(weightedUNs,
- [](const auto &left, const auto &right) {
- return left.first == right.first;
- }),
- weightedUNs.end());
- uns.clear();
- weights.clear();
- for (auto [un, weight] : weightedUNs) {
- uns.push_back(un);
- weights.push_back(weight);
- }
+ for (auto un : compressionUns)
+ uns.emplace_back(un, 1);
+ llvm::sort(uns, [](const auto &left, const auto &right) {
+ return left.Id < right.Id;
+ });
+ uns.erase(llvm::unique(uns,
+ [](const auto &left, const auto &right) {
+ return left.Id == right.Id;
+ }),
+ uns.end());
}
}
@@ -388,8 +374,8 @@ auto BPOrderer<D>::computeOrder(
std::vector<BPFunctionNode> nodesForStartup;
for (auto &[sectionIdx, uns] : startupSectionIdxUNs)
- nodesForStartup.emplace_back(sectionIdx, uns,
- startupSectionIdxUNWeights[sectionIdx]);
+ nodesForStartup.push_back(
+ BPFunctionNode::createWithWeightedUtilities(sectionIdx, uns));
// Use the first timestamp to define the initial order for startup nodes.
llvm::sort(nodesForStartup, [§ionIdxToTimestamp](auto &L, auto &R) {
diff --git a/lld/test/ELF/bp-section-orderer-weighted-profile.s b/lld/test/ELF/bp-section-orderer-weighted-profile.s
new file mode 100644
index 0000000000000..1814090e0a21a
--- /dev/null
+++ b/lld/test/ELF/bp-section-orderer-weighted-profile.s
@@ -0,0 +1,130 @@
+# REQUIRES: aarch64
+
+# RUN: rm -rf %t && split-file %s %t
+# RUN: llvm-mc -filetype=obj -triple=aarch64 %t/input.s -o %t/input.o
+# RUN: llvm-profdata merge %t/weighted.proftext %t/functions.proftext -o %t/weighted.profdata
+# RUN: llvm-profdata merge %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext \
+# RUN: %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext \
+# RUN: %t/competing.proftext %t/functions.proftext -o %t/replicated.profdata
+# RUN: llvm-profdata merge %t/primary.proftext %t/competing.proftext %t/functions.proftext -o %t/unweighted.profdata
+
+# A weight of 10 is equivalent to ten copies of the same trace in the ELF
+# balanced-partitioning consumer.
+# RUN: ld.lld -e _start -o %t/weighted.out %t/input.o --irpgo-profile=%t/weighted.profdata --bp-startup-sort=function
+# RUN: ld.lld -e _start -o %t/replicated.out %t/input.o --irpgo-profile=%t/replicated.profdata --bp-startup-sort=function
+# RUN: ld.lld -e _start -o %t/unweighted.out %t/input.o --irpgo-profile=%t/unweighted.profdata --bp-startup-sort=function
+# RUN: llvm-nm -jn %t/weighted.out > %t/weighted.order
+# RUN: llvm-nm -jn %t/replicated.out > %t/replicated.order
+# RUN: llvm-nm -jn %t/unweighted.out > %t/unweighted.order
+# RUN: cmp %t/weighted.order %t/replicated.order
+# RUN: not cmp %t/weighted.order %t/unweighted.order
+# RUN: FileCheck %s --input-file=%t/weighted.order --check-prefix=FUNCTIONS
+
+# FUNCTIONS-DAG: A
+# FUNCTIONS-DAG: B
+# FUNCTIONS-DAG: C
+# FUNCTIONS-DAG: D
+# FUNCTIONS-DAG: E
+# FUNCTIONS-DAG: F
+
+#--- input.s
+.section .text._start,"ax", at progbits
+.globl _start
+_start:
+ ret
+
+.section .text.A,"ax", at progbits
+.globl A
+A:
+ ret
+
+.section .text.B,"ax", at progbits
+.globl B
+B:
+ ret
+
+.section .text.C,"ax", at progbits
+.globl C
+C:
+ ret
+
+.section .text.D,"ax", at progbits
+.globl D
+D:
+ ret
+
+.section .text.E,"ax", at progbits
+.globl E
+E:
+ ret
+
+.section .text.F,"ax", at progbits
+.globl F
+F:
+ ret
+
+#--- weighted.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
+2
+# Trace Stream Size
+11
+# Weight
+10
+A, B, C, D, E, F
+# Weight
+1
+A, D, B, C, E, F
+
+#--- primary.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
+1
+# Trace Stream Size
+1
+# Weight
+1
+A, B, C, D, E, F
+
+#--- competing.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
+1
+# Trace Stream Size
+1
+# Weight
+1
+A, D, B, C, E, F
+
+#--- functions.proftext
+:ir
+A
+# Func Hash
+1
+# Num Counters
+1
+# Counter Values
+1
+B
+2
+1
+1
+C
+3
+1
+1
+D
+4
+1
+1
+E
+5
+1
+1
+F
+6
+1
+1
diff --git a/lld/test/MachO/bp-section-orderer-weighted-profile.s b/lld/test/MachO/bp-section-orderer-weighted-profile.s
index 1907a697c7fa7..b2278265c6196 100644
--- a/lld/test/MachO/bp-section-orderer-weighted-profile.s
+++ b/lld/test/MachO/bp-section-orderer-weighted-profile.s
@@ -3,13 +3,20 @@
# RUN: rm -rf %t && split-file %s %t
# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %t/input.s -o %t/input.o
# RUN: llvm-profdata merge %t/weighted.proftext %t/functions.proftext -o %t/weighted.profdata
-# RUN: llvm-profdata merge %t/replicated.proftext %t/functions.proftext -o %t/replicated.profdata
-# RUN: llvm-profdata merge %t/unweighted.proftext %t/functions.proftext -o %t/unweighted.profdata
+# RUN: llvm-profdata merge %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext \
+# RUN: %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext \
+# RUN: %t/competing.proftext %t/functions.proftext -o %t/replicated.profdata
+# RUN: llvm-profdata merge %t/primary.proftext %t/competing.proftext %t/functions.proftext -o %t/unweighted.profdata
+# RUN: llvm-profdata merge %t/saturated.proftext %t/functions.proftext -o %t/saturated.profdata
# A weight of 10 is equivalent to ten copies of the same trace.
-# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/weighted.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - | FileCheck %s --check-prefix=WEIGHTED
-# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/replicated.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - | FileCheck %s --check-prefix=WEIGHTED
-# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/unweighted.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - | FileCheck %s --check-prefix=UNWEIGHTED
+# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/weighted.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - > %t/weighted.order
+# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/replicated.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - > %t/replicated.order
+# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/unweighted.profdata --bp-startup-sort=function | llvm-nm --numeric-sort --format=just-symbols - > %t/unweighted.order
+# RUN: cmp %t/weighted.order %t/replicated.order
+# RUN: not cmp %t/weighted.order %t/unweighted.order
+# RUN: FileCheck %s --input-file=%t/weighted.order --check-prefix=WEIGHTED
+# RUN: FileCheck %s --input-file=%t/unweighted.order --check-prefix=UNWEIGHTED
# WEIGHTED: D
# WEIGHTED-NEXT: A
@@ -24,28 +31,47 @@
# UNWEIGHTED-NEXT: E
# UNWEIGHTED-NEXT: F
+# Combining weight-one compression utilities with weighted temporal utilities
+# preserves the same weighted-versus-replicated equivalence.
+# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/weighted.profdata --bp-startup-sort=function --bp-compression-sort-startup-functions | llvm-nm --numeric-sort --format=just-symbols - > %t/weighted-compression.order
+# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/replicated.profdata --bp-startup-sort=function --bp-compression-sort-startup-functions | llvm-nm --numeric-sort --format=just-symbols - > %t/replicated-compression.order
+# RUN: %lld -arch arm64 -lSystem -e _main -o - %t/input.o --irpgo-profile=%t/unweighted.profdata --bp-startup-sort=function --bp-compression-sort-startup-functions | llvm-nm --numeric-sort --format=just-symbols - > %t/unweighted-compression.order
+# RUN: cmp %t/weighted-compression.order %t/replicated-compression.order
+# RUN: not cmp %t/weighted-compression.order %t/unweighted-compression.order
+
# Verbose trace-reference counts and the page-fault area also respect weights.
# RUN: %lld -arch arm64 -lSystem -e _main -o %t/weighted.out %t/input.o --irpgo-profile=%t/weighted.profdata --bp-startup-sort=function --verbose-bp-section-orderer 2>&1 | FileCheck %s --check-prefix=WEIGHTED-STATS
# RUN: %lld -arch arm64 -lSystem -e _main -o %t/replicated.out %t/input.o --irpgo-profile=%t/replicated.profdata --bp-startup-sort=function --verbose-bp-section-orderer 2>&1 | FileCheck %s --check-prefix=WEIGHTED-STATS
# WEIGHTED-STATS: Temporal profile function references: 66 / 66 resolved (6 / 6 unique)
# WEIGHTED-STATS: Total area under the page fault curve: 6.600000e+01
+# Oversized weights saturate verbose diagnostics instead of wrapping.
+# RUN: %lld -arch arm64 -lSystem -e _main -o %t/saturated.out %t/input.o --irpgo-profile=%t/saturated.profdata --bp-startup-sort=function --verbose-bp-section-orderer 2>&1 | FileCheck %s --check-prefix=SATURATED-STATS
+# SATURATED-STATS: Temporal profile function references: 18446744073709551615 / 18446744073709551615 resolved (2 / 2 unique)
+# SATURATED-STATS: Total area under the page fault curve: 1.844674e+19
+
#--- input.s
.text
.globl _main, A, B, C, D, E, F
_main:
ret
A:
+ add w0, w0, #1
ret
B:
+ add w0, w0, #1
ret
C:
+ add w0, w0, #2
ret
D:
+ add w0, w0, #2
ret
E:
+ add w0, w0, #3
ret
F:
+ add w0, w0, #3
ret
.subsections_via_symbols
@@ -63,60 +89,38 @@ A, B, C, D, E, F
1
A, D, B, C, E, F
-#--- replicated.proftext
+#--- primary.proftext
:ir
:temporal_prof_traces
# Num Traces
-11
-# Trace Stream Size
-11
-# Weight
-1
-A, B, C, D, E, F
-# Weight
-1
-A, B, C, D, E, F
-# Weight
-1
-A, B, C, D, E, F
-# Weight
-1
-A, B, C, D, E, F
-# Weight
-1
-A, B, C, D, E, F
-# Weight
1
-A, B, C, D, E, F
-# Weight
+# Trace Stream Size
1
-A, B, C, D, E, F
# Weight
1
A, B, C, D, E, F
-# Weight
+
+#--- competing.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
1
-A, B, C, D, E, F
-# Weight
+# Trace Stream Size
1
-A, B, C, D, E, F
# Weight
1
A, D, B, C, E, F
-#--- unweighted.proftext
+#--- saturated.proftext
:ir
:temporal_prof_traces
# Num Traces
-2
+1
# Trace Stream Size
-2
-# Weight
1
-A, B, C, D, E, F
# Weight
-1
-A, D, B, C, E, F
+18446744073709551615
+A, B
#--- functions.proftext
:ir
diff --git a/llvm/include/llvm/Support/BalancedPartitioning.h b/llvm/include/llvm/Support/BalancedPartitioning.h
index 355155f3d512c..d088b78591b28 100644
--- a/llvm/include/llvm/Support/BalancedPartitioning.h
+++ b/llvm/include/llvm/Support/BalancedPartitioning.h
@@ -41,12 +41,10 @@
#include "raw_ostream.h"
#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Compiler.h"
#include <atomic>
-#include <cassert>
#include <condition_variable>
#include <mutex>
#include <random>
@@ -65,27 +63,38 @@ class BPFunctionNode {
using UtilityNodeT = uint32_t;
using UtilityNodeWeightT = uint64_t;
- /// \param UtilityNodes the set of utility nodes (must be unique'd)
- BPFunctionNode(IDT Id, ArrayRef<UtilityNodeT> UtilityNodes)
- : Id(Id), UtilityNodes(UtilityNodes),
- UtilityNodeWeights(UtilityNodes.size(), 1) {}
+ struct WeightedUtilityNode {
+ UtilityNodeT Id;
+ UtilityNodeWeightT Weight;
+
+ WeightedUtilityNode(UtilityNodeT Id, UtilityNodeWeightT Weight)
+ : Id(Id), Weight(Weight) {}
+
+ bool operator==(const WeightedUtilityNode &Other) const {
+ return Id == Other.Id && Weight == Other.Weight;
+ }
+ };
/// \param UtilityNodes the set of utility nodes (must be unique'd)
- /// \param UtilityNodeWeights the weight of each corresponding utility node.
- /// All occurrences of a utility node must have the same weight. Zero-weight
- /// utility nodes are ignored.
- BPFunctionNode(IDT Id, ArrayRef<UtilityNodeT> UtilityNodes,
- ArrayRef<UtilityNodeWeightT> UtilityNodeWeights)
- : Id(Id) {
- assert(UtilityNodes.size() == UtilityNodeWeights.size() &&
- "each utility node must have a weight");
- for (auto [UtilityNode, Weight] :
- llvm::zip(UtilityNodes, UtilityNodeWeights)) {
- if (Weight == 0)
+ BPFunctionNode(IDT Id, ArrayRef<UtilityNodeT> UtilityNodes) : Id(Id) {
+ for (UtilityNodeT UtilityNode : UtilityNodes)
+ this->UtilityNodes.emplace_back(UtilityNode, 1);
+ }
+
+ /// Create a node with weighted utility nodes (which must be unique'd). All
+ /// occurrences of a utility node must have the same weight. A utility's cost
+ /// is multiplied by its weight without materializing duplicate utility
+ /// nodes. Zero-weight utility nodes are ignored.
+ static BPFunctionNode
+ createWithWeightedUtilities(IDT Id,
+ ArrayRef<WeightedUtilityNode> UtilityNodes) {
+ BPFunctionNode Node(Id);
+ for (WeightedUtilityNode UtilityNode : UtilityNodes) {
+ if (UtilityNode.Weight == 0)
continue;
- this->UtilityNodes.push_back(UtilityNode);
- this->UtilityNodeWeights.push_back(Weight);
+ Node.UtilityNodes.push_back(UtilityNode);
}
+ return Node;
}
/// The ID of this node
@@ -95,10 +104,7 @@ class BPFunctionNode {
protected:
/// The list of utility nodes associated with this node
- SmallVector<UtilityNodeT, 4> UtilityNodes;
- /// The weight of each corresponding utility node. A utility node's cost is
- /// multiplied by its weight without materializing duplicate utility nodes.
- SmallVector<UtilityNodeWeightT, 4> UtilityNodeWeights;
+ SmallVector<WeightedUtilityNode, 4> UtilityNodes;
/// The bucket assigned by balanced partitioning
std::optional<unsigned> Bucket;
/// The index of the input order of the FunctionNodes
@@ -107,6 +113,9 @@ class BPFunctionNode {
friend class BPFunctionNodeTest_Basic_Test;
friend class BalancedPartitioningTest_Basic_Test;
friend class BalancedPartitioningTest_Large_Test;
+
+private:
+ explicit BPFunctionNode(IDT Id) : Id(Id) {}
};
/// Algorithm parameters; default values are tuned on real-world binaries
diff --git a/llvm/lib/ProfileData/InstrProf.cpp b/llvm/lib/ProfileData/InstrProf.cpp
index 49cec36ec1bb2..6b8c60a39f399 100644
--- a/llvm/lib/ProfileData/InstrProf.cpp
+++ b/llvm/lib/ProfileData/InstrProf.cpp
@@ -718,9 +718,10 @@ Error collectGlobalObjectNameStrings(ArrayRef<std::string> NameStrs,
std::string UncompressedNameStrings =
join(NameStrs.begin(), NameStrs.end(), getInstrProfNameSeparator());
- assert(StringRef(UncompressedNameStrings)
- .count(getInstrProfNameSeparator()) == (NameStrs.size() - 1) &&
- "PGO name is invalid (contains separator token)");
+ assert(
+ StringRef(UncompressedNameStrings).count(getInstrProfNameSeparator()) ==
+ (NameStrs.size() - 1) &&
+ "PGO name is invalid (contains separator token)");
unsigned EncLen = encodeULEB128(UncompressedNameStrings.length(), P);
P += EncLen;
@@ -1152,8 +1153,7 @@ void TemporalProfTraceTy::createBPFunctionNodes(
UtilityNodeT MaxUN = 0;
DenseMap<IDT, size_t> IdToFirstTimestamp;
DenseMap<IDT, UtilityNodeT> IdToFirstUN;
- DenseMap<IDT, SmallVector<UtilityNodeT>> IdToUNs;
- DenseMap<IDT, SmallVector<BPFunctionNode::UtilityNodeWeightT>> IdToUNWeights;
+ DenseMap<IDT, SmallVector<BPFunctionNode::WeightedUtilityNode>> IdToUNs;
for (auto &Trace : Traces) {
// A zero-weight trace is equivalent to no copies of the trace.
if (Trace.Weight == 0)
@@ -1172,10 +1172,8 @@ void TemporalProfTraceTy::createBPFunctionNodes(
IdToFirstUN.try_emplace(Id, MaxUN);
}
for (auto &[Id, FirstUN] : IdToFirstUN)
- for (auto UN = FirstUN; UN <= MaxUN; ++UN) {
- IdToUNs[Id].push_back(UN);
- IdToUNWeights[Id].push_back(Trace.Weight);
- }
+ for (auto UN = FirstUN; UN <= MaxUN; ++UN)
+ IdToUNs[Id].emplace_back(UN, Trace.Weight);
++MaxUN;
IdToFirstUN.clear();
}
@@ -1184,28 +1182,18 @@ void TemporalProfTraceTy::createBPFunctionNodes(
DenseMap<UtilityNodeT, unsigned> UNFrequency;
for (auto &[Id, UNs] : IdToUNs)
for (auto &UN : UNs)
- ++UNFrequency[UN];
+ ++UNFrequency[UN.Id];
// Filter out utility nodes that are too infrequent or too prevalent to make
// BalancedPartitioning more effective.
- for (auto &[Id, UNs] : IdToUNs) {
- auto &Weights = IdToUNWeights[Id];
- size_t WriteIndex = 0;
- for (size_t ReadIndex = 0; ReadIndex < UNs.size(); ++ReadIndex) {
- auto UN = UNs[ReadIndex];
- unsigned Freq = UNFrequency[UN];
- if (Freq <= 1 || 2 * Freq > IdToUNs.size())
- continue;
- UNs[WriteIndex] = UN;
- Weights[WriteIndex] = Weights[ReadIndex];
- ++WriteIndex;
- }
- UNs.resize(WriteIndex);
- Weights.resize(WriteIndex);
- }
+ for (auto &[Id, UNs] : IdToUNs)
+ llvm::erase_if(UNs, [&](auto &UN) {
+ unsigned Freq = UNFrequency[UN.Id];
+ return Freq <= 1 || 2 * Freq > IdToUNs.size();
+ });
}
for (auto &[Id, UNs] : IdToUNs)
- Nodes.emplace_back(Id, UNs, IdToUNWeights[Id]);
+ Nodes.push_back(BPFunctionNode::createWithWeightedUtilities(Id, UNs));
// Since BalancedPartitioning is sensitive to the initial order, we explicitly
// order nodes by their earliest timestamp.
@@ -1228,13 +1216,13 @@ uint32_t getNumValueKindsInstrProf(const void *Record) {
}
uint32_t getNumValueSitesInstrProf(const void *Record, uint32_t VKind) {
- return reinterpret_cast<const InstrProfRecord *>(Record)
- ->getNumValueSites(VKind);
+ return reinterpret_cast<const InstrProfRecord *>(Record)->getNumValueSites(
+ VKind);
}
uint32_t getNumValueDataInstrProf(const void *Record, uint32_t VKind) {
- return reinterpret_cast<const InstrProfRecord *>(Record)
- ->getNumValueData(VKind);
+ return reinterpret_cast<const InstrProfRecord *>(Record)->getNumValueData(
+ VKind);
}
uint32_t getNumValueDataForSiteInstrProf(const void *R, uint32_t VK,
@@ -1435,9 +1423,8 @@ void annotateValueSite(Module &M, Instruction &Inst,
}
void annotateValueSite(Module &M, Instruction &Inst,
- ArrayRef<InstrProfValueData> VDs,
- uint64_t Sum, InstrProfValueKind ValueKind,
- uint32_t MaxMDCount) {
+ ArrayRef<InstrProfValueData> VDs, uint64_t Sum,
+ InstrProfValueKind ValueKind, uint32_t MaxMDCount) {
if (VDs.empty())
return;
LLVMContext &Ctx = M.getContext();
diff --git a/llvm/lib/Support/BalancedPartitioning.cpp b/llvm/lib/Support/BalancedPartitioning.cpp
index c18ec77e2a3a3..781930261c65f 100644
--- a/llvm/lib/Support/BalancedPartitioning.cpp
+++ b/llvm/lib/Support/BalancedPartitioning.cpp
@@ -22,8 +22,15 @@ using namespace llvm;
#define DEBUG_TYPE "balanced-partitioning"
void BPFunctionNode::dump(raw_ostream &OS) const {
- OS << formatv("{{ID={0} Utilities={{{1:$[,]}} Bucket={2}}", Id,
- make_range(UtilityNodes.begin(), UtilityNodes.end()), Bucket);
+ OS << "{ID=" << Id << " Utilities={";
+ for (const auto &UN : UtilityNodes)
+ OS << "(" << UN.Id << ", " << UN.Weight << ")";
+ OS << "} Bucket=";
+ if (Bucket)
+ OS << *Bucket;
+ else
+ OS << "none";
+ OS << "}";
}
template <typename Func>
@@ -172,40 +179,33 @@ void BalancedPartitioning::runIterations(const FunctionNodeRange Nodes,
DenseMap<BPFunctionNode::UtilityNodeT, unsigned> UtilityNodeIndex;
for (auto &N : Nodes)
for (auto &UN : N.UtilityNodes)
- ++UtilityNodeIndex[UN];
+ ++UtilityNodeIndex[UN.Id];
// Remove utility nodes if they have just one edge or are connected to all
// functions
- for (auto &N : Nodes) {
- size_t WriteIndex = 0;
- for (size_t ReadIndex = 0; ReadIndex < N.UtilityNodes.size(); ++ReadIndex) {
- auto UN = N.UtilityNodes[ReadIndex];
- unsigned UNI = UtilityNodeIndex[UN];
- if (UNI == 1 || UNI == NumNodes)
- continue;
- N.UtilityNodes[WriteIndex] = UN;
- N.UtilityNodeWeights[WriteIndex] = N.UtilityNodeWeights[ReadIndex];
- ++WriteIndex;
- }
- N.UtilityNodes.resize(WriteIndex);
- N.UtilityNodeWeights.resize(WriteIndex);
- }
+ for (auto &N : Nodes)
+ llvm::erase_if(N.UtilityNodes, [&](auto &UN) {
+ unsigned UNI = UtilityNodeIndex[UN.Id];
+ return UNI == 1 || UNI == NumNodes;
+ });
// Renumber utility nodes so they can be used to index into Signatures
UtilityNodeIndex.clear();
for (auto &N : Nodes)
for (auto &UN : N.UtilityNodes)
- UN = UtilityNodeIndex.insert({UN, UtilityNodeIndex.size()}).first->second;
+ UN.Id = UtilityNodeIndex.insert({UN.Id, UtilityNodeIndex.size()})
+ .first->second;
// Initialize signatures
SignaturesT Signatures(/*Size=*/UtilityNodeIndex.size());
for (auto &N : Nodes) {
- for (auto [UN, Weight] : llvm::zip(N.UtilityNodes, N.UtilityNodeWeights)) {
- assert(UN < Signatures.size());
- auto &Signature = Signatures[UN];
- if (Signature.LeftCount != 0 || Signature.RightCount != 0)
- assert(Signature.Weight == Weight &&
+ for (auto &UN : N.UtilityNodes) {
+ assert(UN.Id < Signatures.size());
+ auto &Signature = Signatures[UN.Id];
+ if (Signature.LeftCount == 0 && Signature.RightCount == 0)
+ Signature.Weight = UN.Weight;
+ else
+ assert(Signature.Weight == UN.Weight &&
"a utility node must have a consistent weight");
- Signature.Weight = Weight;
if (N.Bucket == LeftBucket) {
Signature.LeftCount++;
} else {
@@ -307,14 +307,14 @@ bool BalancedPartitioning::moveFunctionNode(BPFunctionNode &N,
// Update signatures and invalidate gain cache
if (FromLeftToRight) {
for (auto &UN : N.UtilityNodes) {
- auto &Signature = Signatures[UN];
+ auto &Signature = Signatures[UN.Id];
Signature.LeftCount--;
Signature.RightCount++;
Signature.CachedGainIsValid = false;
}
} else {
for (auto &UN : N.UtilityNodes) {
- auto &Signature = Signatures[UN];
+ auto &Signature = Signatures[UN.Id];
Signature.LeftCount++;
Signature.RightCount--;
Signature.CachedGainIsValid = false;
@@ -343,8 +343,8 @@ float BalancedPartitioning::moveGain(const BPFunctionNode &N,
const SignaturesT &Signatures) {
float Gain = 0.f;
for (auto &UN : N.UtilityNodes)
- Gain += (FromLeftToRight ? Signatures[UN].CachedGainLR
- : Signatures[UN].CachedGainRL);
+ Gain += (FromLeftToRight ? Signatures[UN.Id].CachedGainLR
+ : Signatures[UN.Id].CachedGainRL);
return Gain;
}
diff --git a/llvm/test/tools/llvm-profdata/order-weighted-traces.test b/llvm/test/tools/llvm-profdata/order-weighted-traces.test
new file mode 100644
index 0000000000000..496567751413c
--- /dev/null
+++ b/llvm/test/tools/llvm-profdata/order-weighted-traces.test
@@ -0,0 +1,87 @@
+# RUN: rm -rf %t && split-file %s %t
+# RUN: llvm-profdata merge %t/weighted.proftext %t/functions.proftext -o %t/weighted.profdata
+# RUN: llvm-profdata merge %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext \
+# RUN: %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext %t/primary.proftext \
+# RUN: %t/competing.proftext %t/functions.proftext -o %t/replicated.profdata
+# RUN: llvm-profdata merge %t/primary.proftext %t/competing.proftext %t/functions.proftext -o %t/unweighted.profdata
+# RUN: llvm-profdata order %t/weighted.profdata -o %t/weighted.order
+# RUN: llvm-profdata order %t/replicated.profdata -o %t/replicated.order
+# RUN: llvm-profdata order %t/unweighted.profdata -o %t/unweighted.order
+
+# A weight of 10 has the same ordering objective as ten copies of the trace.
+# RUN: cmp %t/weighted.order %t/replicated.order
+# RUN: not cmp %t/weighted.order %t/unweighted.order
+# RUN: FileCheck %s --input-file=%t/weighted.order --check-prefix=FUNCTIONS
+
+# FUNCTIONS-DAG: A
+# FUNCTIONS-DAG: B
+# FUNCTIONS-DAG: C
+# FUNCTIONS-DAG: D
+# FUNCTIONS-DAG: E
+# FUNCTIONS-DAG: F
+
+#--- weighted.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
+2
+# Trace Stream Size
+11
+# Weight
+10
+A, B, C, D, E, F
+# Weight
+1
+A, D, B, C, E, F
+
+#--- primary.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
+1
+# Trace Stream Size
+1
+# Weight
+1
+A, B, C, D, E, F
+
+#--- competing.proftext
+:ir
+:temporal_prof_traces
+# Num Traces
+1
+# Trace Stream Size
+1
+# Weight
+1
+A, D, B, C, E, F
+
+#--- functions.proftext
+:ir
+A
+# Func Hash
+1
+# Num Counters
+1
+# Counter Values
+1
+B
+2
+1
+1
+C
+3
+1
+1
+D
+4
+1
+1
+E
+5
+1
+1
+F
+6
+1
+1
diff --git a/llvm/test/tools/llvm-profdata/show-order-weight-saturation.proftext b/llvm/test/tools/llvm-profdata/show-order-weight-saturation.proftext
new file mode 100644
index 0000000000000..b9def54a9d5f6
--- /dev/null
+++ b/llvm/test/tools/llvm-profdata/show-order-weight-saturation.proftext
@@ -0,0 +1,31 @@
+# RUN: llvm-profdata order %s --num-test-traces=1 | FileCheck %s
+
+# CHECK: # Total area under the page fault curve: 1.844674e+19
+
+:ir
+:temporal_prof_traces
+# Num Traces
+3
+# Trace Stream Size
+3
+# Weight
+1
+a, b
+# Weight
+1
+b, a
+# Weight
+18446744073709551615
+a, b
+
+a
+# Func Hash
+1
+# Num Counters
+1
+# Counter Values
+1
+b
+2
+1
+1
diff --git a/llvm/test/tools/llvm-profdata/show-order.proftext b/llvm/test/tools/llvm-profdata/show-order.proftext
index 28eb1b9b42af7..c3bcd544d8fc3 100644
--- a/llvm/test/tools/llvm-profdata/show-order.proftext
+++ b/llvm/test/tools/llvm-profdata/show-order.proftext
@@ -1,6 +1,6 @@
# RUN: llvm-profdata order %s --num-test-traces=1 | FileCheck %s
-# CHECK: # Total area under the page fault curve: 4.000000e+00
+# CHECK: # Total area under the page fault curve: 4.000000e+01
# CHECK: a
# CHECK: b
@@ -24,7 +24,7 @@ a, x, main.c:b, c
1
a, main.c:b, c
# Weight
-1
+10
a, main.c:b, c, x
a
diff --git a/llvm/tools/llvm-profdata/llvm-profdata.cpp b/llvm/tools/llvm-profdata/llvm-profdata.cpp
index d3c83457a5d81..c73bdbbf3db18 100644
--- a/llvm/tools/llvm-profdata/llvm-profdata.cpp
+++ b/llvm/tools/llvm-profdata/llvm-profdata.cpp
@@ -38,6 +38,7 @@
#include "llvm/Support/FormattedStream.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/MD5.h"
+#include "llvm/Support/MathExtras.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/Regex.h"
@@ -638,7 +639,7 @@ class SymbolRemapper {
return New.empty() ? Name : FunctionId(New);
}
};
-}
+} // namespace
struct WeightedFile {
std::string Filename;
@@ -894,13 +895,11 @@ getFuncName(const StringMap<InstrProfWriter::ProfilingData>::value_type &Val) {
return Val.first();
}
-static std::string
-getFuncName(const SampleProfileMap::value_type &Val) {
+static std::string getFuncName(const SampleProfileMap::value_type &Val) {
return Val.second.getContext().toString();
}
-template <typename T>
-static void filterFunctions(T &ProfileMap) {
+template <typename T> static void filterFunctions(T &ProfileMap) {
bool hasFilter = !FuncNameFilter.empty();
bool hasNegativeFilter = !FuncNameNegativeFilter.empty();
if (!hasFilter && !hasNegativeFilter)
@@ -1499,8 +1498,8 @@ remapSamples(const sampleprof::FunctionSamples &Samples,
BodySample.second.getSamples());
for (const auto &Target : BodySample.second.getCallTargets()) {
Result.addCalledTargetSamples(BodySample.first.LineOffset,
- MaskedDiscriminator,
- Remapper(Target.first), Target.second);
+ MaskedDiscriminator, Remapper(Target.first),
+ Target.second);
}
}
for (const auto &CallsiteSamples : Samples.getCallsiteSamples()) {
@@ -1517,12 +1516,8 @@ remapSamples(const sampleprof::FunctionSamples &Samples,
}
static sampleprof::SampleProfileFormat FormatMap[] = {
- sampleprof::SPF_None,
- sampleprof::SPF_Text,
- sampleprof::SPF_None,
- sampleprof::SPF_Ext_Binary,
- sampleprof::SPF_GCC,
- sampleprof::SPF_Binary};
+ sampleprof::SPF_None, sampleprof::SPF_Text, sampleprof::SPF_None,
+ sampleprof::SPF_Ext_Binary, sampleprof::SPF_GCC, sampleprof::SPF_Binary};
static std::unique_ptr<MemoryBuffer>
getInputFileBuf(const StringRef &InputFile) {
@@ -3396,7 +3391,8 @@ static int show_main(StringRef ProgName) {
exitWithErrorCode(EC, OutputFilename);
if (ShowAllFunctions && !FuncNameFilter.empty())
- WithColor::warning() << "-function argument ignored: showing all functions\n";
+ WithColor::warning()
+ << "-function argument ignored: showing all functions\n";
if (!DebugInfoFilename.empty())
return showDebugInfoCorrelation(DebugInfoFilename, SFormat, OS);
@@ -3446,15 +3442,20 @@ static int order_main() {
IdToPageNumber[Node.Id] = IdToPageNumber.size() / 32;
SmallSet<unsigned, 0> TouchedPages;
- unsigned Area = 0;
+ uint64_t Area = 0;
for (auto &Trace : TestTraces) {
+ if (Trace.Weight == 0)
+ continue;
+ uint64_t TraceArea = 0;
for (auto Id : Trace.FunctionNameRefs) {
auto It = IdToPageNumber.find(Id);
if (It == IdToPageNumber.end())
continue;
TouchedPages.insert(It->getSecond());
- Area += TouchedPages.size();
+ TraceArea = SaturatingAdd(TraceArea,
+ static_cast<uint64_t>(TouchedPages.size()));
}
+ Area = SaturatingMultiplyAdd(TraceArea, Trace.Weight, Area);
TouchedPages.clear();
}
OS << "# Total area under the page fault curve: " << (float)Area << "\n";
diff --git a/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp b/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
index 058dd8d17e355..585343c661e8a 100644
--- a/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
+++ b/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
@@ -24,14 +24,10 @@ void PrintTo(const BPFunctionNode &Node, std::ostream *OS) {
TEST(BPFunctionNodeTest, Basic) {
auto NodeIs = [](BPFunctionNode::IDT Id,
- ArrayRef<BPFunctionNode::UtilityNodeT> UNs,
- ArrayRef<BPFunctionNode::UtilityNodeWeightT> Weights) {
+ ArrayRef<BPFunctionNode::WeightedUtilityNode> UNs) {
return AllOf(Field("Id", &BPFunctionNode::Id, Id),
Field("UtilityNodes", &BPFunctionNode::UtilityNodes,
- UnorderedElementsAreArray(UNs)),
- Field("UtilityNodeWeights",
- &BPFunctionNode::UtilityNodeWeights,
- UnorderedElementsAreArray(Weights)));
+ UnorderedElementsAreArray(UNs)));
};
std::vector<BPFunctionNode> Nodes;
@@ -39,20 +35,21 @@ TEST(BPFunctionNodeTest, Basic) {
{TemporalProfTraceTy({0, 1, 2, 3})}, Nodes, /*RemoveOutlierUNs=*/false);
// Utility nodes that are too infrequent or too prevalent are filtered out.
EXPECT_THAT(Nodes,
- UnorderedElementsAre(NodeIs(0, {0, 1, 2}, {1, 1, 1}),
- NodeIs(1, {1, 2}, {1, 1}),
- NodeIs(2, {2}, {1}), NodeIs(3, {2}, {1})));
+ UnorderedElementsAre(NodeIs(0, {{0, 1}, {1, 1}, {2, 1}}),
+ NodeIs(1, {{1, 1}, {2, 1}}),
+ NodeIs(2, {{2, 1}}), NodeIs(3, {{2, 1}})));
Nodes.clear();
TemporalProfTraceTy::createBPFunctionNodes(
{TemporalProfTraceTy({0, 1, 2, 3, 4}), TemporalProfTraceTy({4, 2})},
Nodes, /*RemoveOutlierUNs=*/false);
- EXPECT_THAT(Nodes, UnorderedElementsAre(NodeIs(0, {0, 1, 2, 3}, {1, 1, 1, 1}),
- NodeIs(1, {1, 2, 3}, {1, 1, 1}),
- NodeIs(2, {2, 3, 5}, {1, 1, 1}),
- NodeIs(3, {2, 3}, {1, 1}),
- NodeIs(4, {3, 4, 5}, {1, 1, 1})));
+ EXPECT_THAT(Nodes,
+ UnorderedElementsAre(NodeIs(0, {{0, 1}, {1, 1}, {2, 1}, {3, 1}}),
+ NodeIs(1, {{1, 1}, {2, 1}, {3, 1}}),
+ NodeIs(2, {{2, 1}, {3, 1}, {5, 1}}),
+ NodeIs(3, {{2, 1}, {3, 1}}),
+ NodeIs(4, {{3, 1}, {4, 1}, {5, 1}})));
Nodes.clear();
TemporalProfTraceTy::createBPFunctionNodes(
@@ -60,18 +57,20 @@ TEST(BPFunctionNodeTest, Basic) {
Nodes, /*RemoveOutlierUNs=*/true);
EXPECT_THAT(Nodes,
- UnorderedElementsAre(NodeIs(0, {1}, {1}), NodeIs(1, {1}, {1}),
- NodeIs(2, {5}, {1}), NodeIs(3, {}, {}),
- NodeIs(4, {5}, {1})));
+ UnorderedElementsAre(NodeIs(0, {{1, 1}}), NodeIs(1, {{1, 1}}),
+ NodeIs(2, {{5, 1}}), NodeIs(3, {}),
+ NodeIs(4, {{5, 1}})));
Nodes.clear();
TemporalProfTraceTy::createBPFunctionNodes(
- {TemporalProfTraceTy({0, 1, 2, 3}, 10), TemporalProfTraceTy({4, 5}, 0)},
+ {TemporalProfTraceTy({0, 1, 2, 3}, 10), TemporalProfTraceTy({0, 4, 1}, 1),
+ TemporalProfTraceTy({5, 6}, 0)},
Nodes, /*RemoveOutlierUNs=*/false);
- EXPECT_THAT(Nodes,
- UnorderedElementsAre(NodeIs(0, {0, 1, 2}, {10, 10, 10}),
- NodeIs(1, {1, 2}, {10, 10}),
- NodeIs(2, {2}, {10}), NodeIs(3, {2}, {10})));
+ EXPECT_THAT(
+ Nodes, UnorderedElementsAre(
+ NodeIs(0, {{0, 10}, {1, 10}, {2, 10}, {3, 1}, {4, 1}, {5, 1}}),
+ NodeIs(1, {{1, 10}, {2, 10}, {5, 1}}), NodeIs(2, {{2, 10}}),
+ NodeIs(3, {{2, 10}}), NodeIs(4, {{4, 1}, {5, 1}})));
}
} // end namespace llvm
diff --git a/llvm/unittests/Support/BalancedPartitioningTest.cpp b/llvm/unittests/Support/BalancedPartitioningTest.cpp
index a781d7b5d405e..ee9fc5a5cb99f 100644
--- a/llvm/unittests/Support/BalancedPartitioningTest.cpp
+++ b/llvm/unittests/Support/BalancedPartitioningTest.cpp
@@ -100,12 +100,18 @@ TEST_F(BalancedPartitioningTest, MoveGain) {
TEST_F(BalancedPartitioningTest, WeightedUtilitiesMatchReplication) {
Config.SkipProbability = 0;
Config.TaskSplitDepth = 0;
+ using WeightedUtilityNode = BPFunctionNode::WeightedUtilityNode;
std::vector<BPFunctionNode> WeightedNodes = {
- BPFunctionNode(0, {1}, {1}),
- BPFunctionNode(1, {1, 2}, {1, 1}),
- BPFunctionNode(2, {0, 1, 2}, {10, 1, 1}),
- BPFunctionNode(3, {0, 2}, {10, 1}),
+ BPFunctionNode::createWithWeightedUtilities(0,
+ {WeightedUtilityNode(1, 1)}),
+ BPFunctionNode::createWithWeightedUtilities(
+ 1, {WeightedUtilityNode(1, 1), WeightedUtilityNode(2, 1)}),
+ BPFunctionNode::createWithWeightedUtilities(
+ 2, {WeightedUtilityNode(0, 10), WeightedUtilityNode(1, 1),
+ WeightedUtilityNode(2, 1)}),
+ BPFunctionNode::createWithWeightedUtilities(
+ 3, {WeightedUtilityNode(0, 10), WeightedUtilityNode(2, 1)}),
};
std::vector<BPFunctionNode> ReplicatedNodes = {
BPFunctionNode(0, {10}),
@@ -120,10 +126,15 @@ TEST_F(BalancedPartitioningTest, WeightedUtilitiesMatchReplication) {
BPFunctionNode(3, {0, 2}),
};
std::vector<BPFunctionNode> WeightOneNodes = {
- BPFunctionNode(0, {1}, {1}),
- BPFunctionNode(1, {1, 2}, {1, 1}),
- BPFunctionNode(2, {0, 1, 2}, {1, 1, 1}),
- BPFunctionNode(3, {0, 2}, {1, 1}),
+ BPFunctionNode::createWithWeightedUtilities(0,
+ {WeightedUtilityNode(1, 1)}),
+ BPFunctionNode::createWithWeightedUtilities(
+ 1, {WeightedUtilityNode(1, 1), WeightedUtilityNode(2, 1)}),
+ BPFunctionNode::createWithWeightedUtilities(
+ 2, {WeightedUtilityNode(0, 1), WeightedUtilityNode(1, 1),
+ WeightedUtilityNode(2, 1)}),
+ BPFunctionNode::createWithWeightedUtilities(
+ 3, {WeightedUtilityNode(0, 1), WeightedUtilityNode(2, 1)}),
};
Bp.run(WeightedNodes);
>From 6c7e035bcb68b3199d4e0c9b3df4553e3117da01 Mon Sep 17 00:00:00 2001
From: Karim Alweheshy <karim.alweheshy at reddit.com>
Date: Wed, 5 Aug 2026 07:50:01 +0200
Subject: [PATCH 3/3] [Support] Keep unweighted BP utility storage compact
---
.../llvm/Support/BalancedPartitioning.h | 58 ++++++++++--
llvm/lib/Support/BalancedPartitioning.cpp | 92 ++++++++++++-------
.../ProfileData/BPFunctionNodeTest.cpp | 22 ++++-
.../Support/BalancedPartitioningTest.cpp | 37 ++++++++
4 files changed, 165 insertions(+), 44 deletions(-)
diff --git a/llvm/include/llvm/Support/BalancedPartitioning.h b/llvm/include/llvm/Support/BalancedPartitioning.h
index d088b78591b28..947a03db1c1d0 100644
--- a/llvm/include/llvm/Support/BalancedPartitioning.h
+++ b/llvm/include/llvm/Support/BalancedPartitioning.h
@@ -41,11 +41,13 @@
#include "raw_ostream.h"
#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Compiler.h"
#include <atomic>
#include <condition_variable>
+#include <memory>
#include <mutex>
#include <random>
#include <vector>
@@ -62,6 +64,7 @@ class BPFunctionNode {
using IDT = uint64_t;
using UtilityNodeT = uint32_t;
using UtilityNodeWeightT = uint64_t;
+ using UtilityNodeWeightsT = SmallVector<UtilityNodeWeightT, 4>;
struct WeightedUtilityNode {
UtilityNodeT Id;
@@ -78,9 +81,34 @@ class BPFunctionNode {
/// \param UtilityNodes the set of utility nodes (must be unique'd)
BPFunctionNode(IDT Id, ArrayRef<UtilityNodeT> UtilityNodes) : Id(Id) {
for (UtilityNodeT UtilityNode : UtilityNodes)
- this->UtilityNodes.emplace_back(UtilityNode, 1);
+ this->UtilityNodes.push_back(UtilityNode);
}
+ BPFunctionNode(const BPFunctionNode &Other)
+ : Id(Other.Id), UtilityNodes(Other.UtilityNodes), Bucket(Other.Bucket),
+ InputOrderIndex(Other.InputOrderIndex) {
+ if (Other.UtilityNodeWeights)
+ UtilityNodeWeights =
+ std::make_unique<UtilityNodeWeightsT>(*Other.UtilityNodeWeights);
+ }
+
+ BPFunctionNode &operator=(const BPFunctionNode &Other) {
+ if (this == &Other)
+ return *this;
+ Id = Other.Id;
+ UtilityNodes = Other.UtilityNodes;
+ UtilityNodeWeights =
+ Other.UtilityNodeWeights
+ ? std::make_unique<UtilityNodeWeightsT>(*Other.UtilityNodeWeights)
+ : nullptr;
+ Bucket = Other.Bucket;
+ InputOrderIndex = Other.InputOrderIndex;
+ return *this;
+ }
+
+ BPFunctionNode(BPFunctionNode &&) = default;
+ BPFunctionNode &operator=(BPFunctionNode &&) = default;
+
/// Create a node with weighted utility nodes (which must be unique'd). All
/// occurrences of a utility node must have the same weight. A utility's cost
/// is multiplied by its weight without materializing duplicate utility
@@ -89,10 +117,17 @@ class BPFunctionNode {
createWithWeightedUtilities(IDT Id,
ArrayRef<WeightedUtilityNode> UtilityNodes) {
BPFunctionNode Node(Id);
+ bool HasNonUnitWeight = llvm::any_of(UtilityNodes, [](auto UtilityNode) {
+ return UtilityNode.Weight != 0 && UtilityNode.Weight != 1;
+ });
+ if (HasNonUnitWeight)
+ Node.UtilityNodeWeights = std::make_unique<UtilityNodeWeightsT>();
for (WeightedUtilityNode UtilityNode : UtilityNodes) {
if (UtilityNode.Weight == 0)
continue;
- Node.UtilityNodes.push_back(UtilityNode);
+ Node.UtilityNodes.push_back(UtilityNode.Id);
+ if (HasNonUnitWeight)
+ Node.UtilityNodeWeights->push_back(UtilityNode.Weight);
}
return Node;
}
@@ -104,7 +139,11 @@ class BPFunctionNode {
protected:
/// The list of utility nodes associated with this node
- SmallVector<WeightedUtilityNode, 4> UtilityNodes;
+ SmallVector<UtilityNodeT, 4> UtilityNodes;
+ /// Optional weights parallel to UtilityNodes. This is absent for the common
+ /// all-weight-one case so unweighted inputs retain their compact edge
+ /// representation.
+ std::unique_ptr<UtilityNodeWeightsT> UtilityNodeWeights;
/// The bucket assigned by balanced partitioning
std::optional<unsigned> Bucket;
/// The index of the input order of the FunctionNodes
@@ -113,6 +152,9 @@ class BPFunctionNode {
friend class BPFunctionNodeTest_Basic_Test;
friend class BalancedPartitioningTest_Basic_Test;
friend class BalancedPartitioningTest_Large_Test;
+ friend class BalancedPartitioningTest_WeightedUtilitiesMatchReplication_Test;
+ friend class
+ BalancedPartitioningTest_WeightOneUtilitiesPreserveLargeOrder_Test;
private:
explicit BPFunctionNode(IDT Id) : Id(Id) {}
@@ -183,9 +225,11 @@ class BalancedPartitioning {
/// Run a bisection iteration to improve the optimization goal
/// \returns the total number of moved FunctionNodes
- unsigned runIteration(const FunctionNodeRange Nodes, unsigned LeftBucket,
- unsigned RightBucket, SignaturesT &Signatures,
- std::mt19937 &RNG) const;
+ unsigned
+ runIteration(const FunctionNodeRange Nodes, unsigned LeftBucket,
+ unsigned RightBucket, SignaturesT &Signatures,
+ ArrayRef<BPFunctionNode::UtilityNodeWeightT> UtilityNodeWeights,
+ std::mt19937 &RNG) const;
/// Try to move \p N from one bucket to another
/// \returns true iff \p N is moved
@@ -224,8 +268,6 @@ class BalancedPartitioning {
float CachedGainRL;
/// Whether \p CachedGainLR and \p CachedGainRL are valid
bool CachedGainIsValid = false;
- /// The number of times this utility contributes to the objective
- BPFunctionNode::UtilityNodeWeightT Weight = 1;
};
protected:
diff --git a/llvm/lib/Support/BalancedPartitioning.cpp b/llvm/lib/Support/BalancedPartitioning.cpp
index 781930261c65f..9f43f469b8d98 100644
--- a/llvm/lib/Support/BalancedPartitioning.cpp
+++ b/llvm/lib/Support/BalancedPartitioning.cpp
@@ -23,8 +23,9 @@ using namespace llvm;
void BPFunctionNode::dump(raw_ostream &OS) const {
OS << "{ID=" << Id << " Utilities={";
- for (const auto &UN : UtilityNodes)
- OS << "(" << UN.Id << ", " << UN.Weight << ")";
+ for (auto [I, UN] : llvm::enumerate(UtilityNodes))
+ OS << "(" << UN << ", "
+ << (UtilityNodeWeights ? (*UtilityNodeWeights)[I] : 1) << ")";
OS << "} Bucket=";
if (Bucket)
OS << *Bucket;
@@ -179,33 +180,58 @@ void BalancedPartitioning::runIterations(const FunctionNodeRange Nodes,
DenseMap<BPFunctionNode::UtilityNodeT, unsigned> UtilityNodeIndex;
for (auto &N : Nodes)
for (auto &UN : N.UtilityNodes)
- ++UtilityNodeIndex[UN.Id];
+ ++UtilityNodeIndex[UN];
// Remove utility nodes if they have just one edge or are connected to all
// functions
- for (auto &N : Nodes)
- llvm::erase_if(N.UtilityNodes, [&](auto &UN) {
- unsigned UNI = UtilityNodeIndex[UN.Id];
- return UNI == 1 || UNI == NumNodes;
- });
+ for (auto &N : Nodes) {
+ auto KeepUtility = [&](auto UN) {
+ unsigned UNI = UtilityNodeIndex[UN];
+ return UNI != 1 && UNI != NumNodes;
+ };
+ if (!N.UtilityNodeWeights) {
+ llvm::erase_if(N.UtilityNodes, [&](auto UN) { return !KeepUtility(UN); });
+ continue;
+ }
+ size_t Out = 0;
+ for (size_t I = 0; I < N.UtilityNodes.size(); ++I) {
+ if (!KeepUtility(N.UtilityNodes[I]))
+ continue;
+ N.UtilityNodes[Out] = N.UtilityNodes[I];
+ (*N.UtilityNodeWeights)[Out] = (*N.UtilityNodeWeights)[I];
+ ++Out;
+ }
+ N.UtilityNodes.resize(Out);
+ N.UtilityNodeWeights->resize(Out);
+ if (llvm::all_of(*N.UtilityNodeWeights,
+ [](auto Weight) { return Weight == 1; }))
+ N.UtilityNodeWeights.reset();
+ }
// Renumber utility nodes so they can be used to index into Signatures
UtilityNodeIndex.clear();
for (auto &N : Nodes)
for (auto &UN : N.UtilityNodes)
- UN.Id = UtilityNodeIndex.insert({UN.Id, UtilityNodeIndex.size()})
- .first->second;
+ UN = UtilityNodeIndex.insert({UN, UtilityNodeIndex.size()}).first->second;
// Initialize signatures
SignaturesT Signatures(/*Size=*/UtilityNodeIndex.size());
+ bool HasWeightedUtilities = llvm::any_of(
+ Nodes, [](const auto &N) { return N.UtilityNodeWeights != nullptr; });
+ SmallVector<BPFunctionNode::UtilityNodeWeightT, 4> UtilityNodeWeights;
+ if (HasWeightedUtilities)
+ UtilityNodeWeights.resize(Signatures.size());
for (auto &N : Nodes) {
- for (auto &UN : N.UtilityNodes) {
- assert(UN.Id < Signatures.size());
- auto &Signature = Signatures[UN.Id];
- if (Signature.LeftCount == 0 && Signature.RightCount == 0)
- Signature.Weight = UN.Weight;
- else
- assert(Signature.Weight == UN.Weight &&
- "a utility node must have a consistent weight");
+ for (auto [I, UN] : llvm::enumerate(N.UtilityNodes)) {
+ assert(UN < Signatures.size());
+ auto &Signature = Signatures[UN];
+ if (HasWeightedUtilities) {
+ uint64_t Weight = N.UtilityNodeWeights ? (*N.UtilityNodeWeights)[I] : 1;
+ if (Signature.LeftCount == 0 && Signature.RightCount == 0)
+ UtilityNodeWeights[UN] = Weight;
+ else
+ assert(UtilityNodeWeights[UN] == Weight &&
+ "a utility node must have a consistent weight");
+ }
if (N.Bucket == LeftBucket) {
Signature.LeftCount++;
} else {
@@ -215,20 +241,20 @@ void BalancedPartitioning::runIterations(const FunctionNodeRange Nodes,
}
for (unsigned I = 0; I < Config.IterationsPerSplit; I++) {
- unsigned NumMovedNodes =
- runIteration(Nodes, LeftBucket, RightBucket, Signatures, RNG);
+ unsigned NumMovedNodes = runIteration(Nodes, LeftBucket, RightBucket,
+ Signatures, UtilityNodeWeights, RNG);
if (NumMovedNodes == 0)
break;
}
}
-unsigned BalancedPartitioning::runIteration(const FunctionNodeRange Nodes,
- unsigned LeftBucket,
- unsigned RightBucket,
- SignaturesT &Signatures,
- std::mt19937 &RNG) const {
+unsigned BalancedPartitioning::runIteration(
+ const FunctionNodeRange Nodes, unsigned LeftBucket, unsigned RightBucket,
+ SignaturesT &Signatures,
+ ArrayRef<BPFunctionNode::UtilityNodeWeightT> UtilityNodeWeights,
+ std::mt19937 &RNG) const {
// Init signature cost caches
- for (auto &Signature : Signatures) {
+ for (auto [I, Signature] : llvm::enumerate(Signatures)) {
if (Signature.CachedGainIsValid)
continue;
unsigned L = Signature.LeftCount;
@@ -245,9 +271,9 @@ unsigned BalancedPartitioning::runIteration(const FunctionNodeRange Nodes,
// Weight distinct utility nodes with the same signature, but avoids
// expanding the graph. Weight == 1 deliberately takes no FP operation so
// unweighted inputs preserve their previous behavior.
- if (Signature.Weight != 1) {
- Signature.CachedGainLR *= static_cast<float>(Signature.Weight);
- Signature.CachedGainRL *= static_cast<float>(Signature.Weight);
+ if (!UtilityNodeWeights.empty() && UtilityNodeWeights[I] != 1) {
+ Signature.CachedGainLR *= static_cast<float>(UtilityNodeWeights[I]);
+ Signature.CachedGainRL *= static_cast<float>(UtilityNodeWeights[I]);
}
Signature.CachedGainIsValid = true;
}
@@ -307,14 +333,14 @@ bool BalancedPartitioning::moveFunctionNode(BPFunctionNode &N,
// Update signatures and invalidate gain cache
if (FromLeftToRight) {
for (auto &UN : N.UtilityNodes) {
- auto &Signature = Signatures[UN.Id];
+ auto &Signature = Signatures[UN];
Signature.LeftCount--;
Signature.RightCount++;
Signature.CachedGainIsValid = false;
}
} else {
for (auto &UN : N.UtilityNodes) {
- auto &Signature = Signatures[UN.Id];
+ auto &Signature = Signatures[UN];
Signature.LeftCount++;
Signature.RightCount--;
Signature.CachedGainIsValid = false;
@@ -343,8 +369,8 @@ float BalancedPartitioning::moveGain(const BPFunctionNode &N,
const SignaturesT &Signatures) {
float Gain = 0.f;
for (auto &UN : N.UtilityNodes)
- Gain += (FromLeftToRight ? Signatures[UN.Id].CachedGainLR
- : Signatures[UN.Id].CachedGainRL);
+ Gain += (FromLeftToRight ? Signatures[UN].CachedGainLR
+ : Signatures[UN].CachedGainRL);
return Gain;
}
diff --git a/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp b/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
index 585343c661e8a..4a7665d00519a 100644
--- a/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
+++ b/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
@@ -12,6 +12,7 @@
#include "gtest/gtest.h"
using testing::Field;
+using testing::Truly;
using testing::UnorderedElementsAre;
using testing::UnorderedElementsAreArray;
@@ -25,9 +26,24 @@ void PrintTo(const BPFunctionNode &Node, std::ostream *OS) {
TEST(BPFunctionNodeTest, Basic) {
auto NodeIs = [](BPFunctionNode::IDT Id,
ArrayRef<BPFunctionNode::WeightedUtilityNode> UNs) {
- return AllOf(Field("Id", &BPFunctionNode::Id, Id),
- Field("UtilityNodes", &BPFunctionNode::UtilityNodes,
- UnorderedElementsAreArray(UNs)));
+ SmallVector<BPFunctionNode::WeightedUtilityNode> Expected(UNs);
+ return AllOf(
+ Field("Id", &BPFunctionNode::Id, Id),
+ Truly([Expected = std::move(Expected)](const BPFunctionNode &Node) {
+ if (Node.UtilityNodes.size() != Expected.size())
+ return false;
+ for (auto ExpectedUN : Expected) {
+ bool Found = false;
+ for (auto [I, UN] : llvm::enumerate(Node.UtilityNodes)) {
+ uint64_t Weight =
+ Node.UtilityNodeWeights ? (*Node.UtilityNodeWeights)[I] : 1;
+ Found |= UN == ExpectedUN.Id && Weight == ExpectedUN.Weight;
+ }
+ if (!Found)
+ return false;
+ }
+ return true;
+ }));
};
std::vector<BPFunctionNode> Nodes;
diff --git a/llvm/unittests/Support/BalancedPartitioningTest.cpp b/llvm/unittests/Support/BalancedPartitioningTest.cpp
index ee9fc5a5cb99f..091b865b0a398 100644
--- a/llvm/unittests/Support/BalancedPartitioningTest.cpp
+++ b/llvm/unittests/Support/BalancedPartitioningTest.cpp
@@ -13,6 +13,7 @@
using testing::Each;
using testing::Field;
+using testing::IsNull;
using testing::Not;
using testing::UnorderedElementsAre;
using testing::UnorderedElementsAreArray;
@@ -137,6 +138,10 @@ TEST_F(BalancedPartitioningTest, WeightedUtilitiesMatchReplication) {
3, {WeightedUtilityNode(0, 1), WeightedUtilityNode(2, 1)}),
};
+ EXPECT_THAT(WeightOneNodes,
+ Each(Field("UtilityNodeWeights",
+ &BPFunctionNode::UtilityNodeWeights, IsNull())));
+
Bp.run(WeightedNodes);
Bp.run(ReplicatedNodes);
Bp.run(UnweightedNodes);
@@ -147,4 +152,36 @@ TEST_F(BalancedPartitioningTest, WeightedUtilitiesMatchReplication) {
EXPECT_EQ(getIds(WeightOneNodes), getIds(UnweightedNodes));
}
+TEST_F(BalancedPartitioningTest, WeightOneUtilitiesPreserveLargeOrder) {
+ const int ProblemSize = 1000;
+ std::vector<BPFunctionNode::UtilityNodeT> AllUNs;
+ for (int I = 0; I < ProblemSize; ++I)
+ AllUNs.push_back(I);
+
+ std::mt19937 RNG;
+ std::vector<BPFunctionNode> UnweightedNodes;
+ std::vector<BPFunctionNode> WeightOneNodes;
+ for (int I = 0; I < ProblemSize; ++I) {
+ std::vector<BPFunctionNode::UtilityNodeT> UNs;
+ int SampleSize =
+ std::uniform_int_distribution<int>(0, AllUNs.size() - 1)(RNG);
+ std::sample(AllUNs.begin(), AllUNs.end(), std::back_inserter(UNs),
+ SampleSize, RNG);
+ UnweightedNodes.emplace_back(I, UNs);
+
+ SmallVector<BPFunctionNode::WeightedUtilityNode> WeightedUNs;
+ for (auto UN : UNs)
+ WeightedUNs.emplace_back(UN, 1);
+ WeightOneNodes.push_back(
+ BPFunctionNode::createWithWeightedUtilities(I, WeightedUNs));
+ }
+
+ EXPECT_THAT(WeightOneNodes,
+ Each(Field("UtilityNodeWeights",
+ &BPFunctionNode::UtilityNodeWeights, IsNull())));
+ Bp.run(UnweightedNodes);
+ Bp.run(WeightOneNodes);
+ EXPECT_EQ(getIds(WeightOneNodes), getIds(UnweightedNodes));
+}
+
} // end namespace llvm
More information about the llvm-commits
mailing list