[InstrProf] Adding utility weights to BalancedPartitioning #72717
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
spupyrev
force-pushed
the
pika-bp-weights
branch
3 times, most recently
from
565ef7e
to
0d355d8
Compare
|
@llvm/pr-subscribers-pgo @llvm/pr-subscribers-llvm-support Author: None (spupyrev) ChangesAdding weights to utility nodes in BP so that we can give more importance to Full diff: https://github.com/llvm/llvm-project/pull/72717.diff 5 Files Affected:
diff --git a/llvm/include/llvm/Support/BalancedPartitioning.h b/llvm/include/llvm/Support/BalancedPartitioning.h
index a8464ac0fe60e58..11d085badc10a0c 100644
--- a/llvm/include/llvm/Support/BalancedPartitioning.h
+++ b/llvm/include/llvm/Support/BalancedPartitioning.h
@@ -57,8 +57,14 @@ class BPFunctionNode {
friend class BalancedPartitioning;
public:
+ /// The type of ID
using IDT = uint64_t;
- using UtilityNodeT = uint32_t;
+ /// The type of UtilityNode
+ struct UtilityNodeT {
+ UtilityNodeT(uint32_t id, uint32_t weight = 1) : id(id), weight(weight) {}
+ uint32_t id;
+ uint32_t weight;
+ };
/// \param UtilityNodes the set of utility nodes (must be unique'd)
BPFunctionNode(IDT Id, ArrayRef<UtilityNodeT> UtilityNodes)
@@ -69,6 +75,8 @@ class BPFunctionNode {
void dump(raw_ostream &OS) const;
+ std::optional<unsigned> getBucket() const { return Bucket; };
+
protected:
/// The list of utility nodes associated with this node
SmallVector<UtilityNodeT, 4> UtilityNodes;
@@ -188,6 +196,8 @@ class BalancedPartitioning {
float CachedGainRL;
/// Whether \p CachedGainLR and \p CachedGainRL are valid
bool CachedGainIsValid = false;
+ /// The weight of this utility node
+ uint32_t Weight = 1;
};
protected:
diff --git a/llvm/lib/ProfileData/InstrProf.cpp b/llvm/lib/ProfileData/InstrProf.cpp
index 236b083a1e2155b..02d679cc920d34d 100644
--- a/llvm/lib/ProfileData/InstrProf.cpp
+++ b/llvm/lib/ProfileData/InstrProf.cpp
@@ -916,15 +916,15 @@ std::vector<BPFunctionNode> TemporalProfTraceTy::createBPFunctionNodes(
int N = std::ceil(std::log2(LargestTraceSize));
- // TODO: We need to use the Trace.Weight field to give more weight to more
+ // TODO: We may use the Trace.Weight field to give more weight to more
// important utilities
DenseMap<IDT, SmallVector<UtilityNodeT, 4>> FuncGroups;
- for (size_t TraceIdx = 0; TraceIdx < Traces.size(); TraceIdx++) {
+ for (uint32_t TraceIdx = 0; TraceIdx < Traces.size(); TraceIdx++) {
auto &Trace = Traces[TraceIdx].FunctionNameRefs;
for (size_t Timestamp = 0; Timestamp < Trace.size(); Timestamp++) {
for (int I = std::floor(std::log2(Timestamp + 1)); I < N; I++) {
auto &FunctionId = Trace[Timestamp];
- UtilityNodeT GroupId = TraceIdx * N + I;
+ UtilityNodeT GroupId = {TraceIdx * N + I};
FuncGroups[FunctionId].push_back(GroupId);
}
}
@@ -933,8 +933,15 @@ std::vector<BPFunctionNode> TemporalProfTraceTy::createBPFunctionNodes(
std::vector<BPFunctionNode> Nodes;
for (auto &Id : FunctionIds) {
auto &UNs = FuncGroups[Id];
- llvm::sort(UNs);
- UNs.erase(std::unique(UNs.begin(), UNs.end()), UNs.end());
+ llvm::sort(UNs.begin(), UNs.end(),
+ [](const UtilityNodeT &L, const UtilityNodeT &R) {
+ return L.id < R.id;
+ });
+ UNs.erase(std::unique(UNs.begin(), UNs.end(),
+ [](const UtilityNodeT &L, const UtilityNodeT &R) {
+ return L.id == R.id;
+ }),
+ UNs.end());
Nodes.emplace_back(Id, UNs);
}
return Nodes;
diff --git a/llvm/lib/Support/BalancedPartitioning.cpp b/llvm/lib/Support/BalancedPartitioning.cpp
index 5843be949911514..6d6f34c8b4dc86c 100644
--- a/llvm/lib/Support/BalancedPartitioning.cpp
+++ b/llvm/lib/Support/BalancedPartitioning.cpp
@@ -20,6 +20,15 @@
using namespace llvm;
#define DEBUG_TYPE "balanced-partitioning"
+namespace llvm {
+template <> struct format_provider<BPFunctionNode::UtilityNodeT> {
+ static void format(const BPFunctionNode::UtilityNodeT &V, raw_ostream &OS,
+ StringRef Style) {
+ OS << "(" << V.id << "-" << V.weight << ")";
+ }
+};
+} // namespace llvm
+
void BPFunctionNode::dump(raw_ostream &OS) const {
OS << formatv("{{ID={0} Utilities={{{1:$[,]}} Bucket={2}}", Id,
make_range(UtilityNodes.begin(), UtilityNodes.end()), Bucket);
@@ -167,37 +176,38 @@ void BalancedPartitioning::runIterations(const FunctionNodeRange Nodes,
unsigned RightBucket,
std::mt19937 &RNG) const {
unsigned NumNodes = std::distance(Nodes.begin(), Nodes.end());
- DenseMap<BPFunctionNode::UtilityNodeT, unsigned> UtilityNodeDegree;
+ DenseMap<uint32_t, unsigned> UtilityNodeDegree;
for (auto &N : Nodes)
for (auto &UN : N.UtilityNodes)
- ++UtilityNodeDegree[UN];
+ ++UtilityNodeDegree[UN.id];
// 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) {
- return UtilityNodeDegree[UN] <= 1 || UtilityNodeDegree[UN] >= NumNodes;
+ return UtilityNodeDegree[UN.id] <= 1 ||
+ UtilityNodeDegree[UN.id] >= NumNodes;
});
// Renumber utility nodes so they can be used to index into Signatures
- DenseMap<BPFunctionNode::UtilityNodeT, unsigned> UtilityNodeIndex;
+ DenseMap<uint32_t, unsigned> UtilityNodeIndex;
for (auto &N : Nodes)
for (auto &UN : N.UtilityNodes)
- if (!UtilityNodeIndex.count(UN))
- UtilityNodeIndex[UN] = UtilityNodeIndex.size();
+ if (!UtilityNodeIndex.count(UN.id))
+ UtilityNodeIndex[UN.id] = UtilityNodeIndex.size();
for (auto &N : Nodes)
for (auto &UN : N.UtilityNodes)
- UN = UtilityNodeIndex[UN];
+ UN.id = UtilityNodeIndex[UN.id];
// Initialize signatures
SignaturesT Signatures(/*Size=*/UtilityNodeIndex.size());
for (auto &N : Nodes) {
for (auto &UN : N.UtilityNodes) {
- assert(UN < Signatures.size());
- if (N.Bucket == LeftBucket) {
- Signatures[UN].LeftCount++;
- } else {
- Signatures[UN].RightCount++;
- }
+ assert(UN.id < Signatures.size());
+ if (N.Bucket == LeftBucket)
+ Signatures[UN.id].LeftCount++;
+ else
+ Signatures[UN.id].RightCount++;
+ Signatures[UN.id].Weight = UN.weight;
}
}
@@ -225,9 +235,11 @@ unsigned BalancedPartitioning::runIteration(const FunctionNodeRange Nodes,
Signature.CachedGainLR = 0.f;
Signature.CachedGainRL = 0.f;
if (L > 0)
- Signature.CachedGainLR = Cost - logCost(L - 1, R + 1);
+ Signature.CachedGainLR =
+ (Cost - logCost(L - 1, R + 1)) * Signature.Weight;
if (R > 0)
- Signature.CachedGainRL = Cost - logCost(L + 1, R - 1);
+ Signature.CachedGainRL =
+ (Cost - logCost(L + 1, R - 1)) * Signature.Weight;
Signature.CachedGainIsValid = true;
}
@@ -286,14 +298,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;
@@ -322,8 +334,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/unittests/ProfileData/BPFunctionNodeTest.cpp b/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
index 97ff5c1edf5efc9..ecfcc722c3cbcdd 100644
--- a/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
+++ b/llvm/unittests/ProfileData/BPFunctionNodeTest.cpp
@@ -29,17 +29,18 @@ TEST(BPFunctionNodeTest, Basic) {
TemporalProfTraceTy({4, 2}),
});
- auto NodeIs = [](BPFunctionNode::IDT Id,
- ArrayRef<BPFunctionNode::UtilityNodeT> UNs) {
- return AllOf(Field("Id", &BPFunctionNode::Id, Id),
- Field("UtilityNodes", &BPFunctionNode::UtilityNodes,
- UnorderedElementsAreArray(UNs)));
- };
-
- EXPECT_THAT(Nodes,
- UnorderedElementsAre(NodeIs(0, {0, 1, 2}), NodeIs(1, {1, 2}),
- NodeIs(2, {1, 2, 4, 5}), NodeIs(3, {2}),
- NodeIs(4, {2, 3, 4, 5})));
+ // TODO
+ // auto NodeIs = [](BPFunctionNode::IDT Id,
+ // ArrayRef<BPFunctionNode::UtilityNodeT> UNs) {
+ // return AllOf(Field("Id", &BPFunctionNode::Id, Id),
+ // Field("UtilityNodes", &BPFunctionNode::UtilityNodes,
+ // UnorderedElementsAreArray(UNs)));
+ // };
+
+ // EXPECT_THAT(Nodes,
+ // UnorderedElementsAre(NodeIs(0, {0, 1, 2}), NodeIs(1, {1, 2}),
+ // NodeIs(2, {1, 2, 4, 5}), NodeIs(3, {2}),
+ // NodeIs(4, {2, 3, 4, 5})));
}
} // end namespace llvm
diff --git a/llvm/unittests/Support/BalancedPartitioningTest.cpp b/llvm/unittests/Support/BalancedPartitioningTest.cpp
index ebe518a8e89cacf..aac7e6edc921812 100644
--- a/llvm/unittests/Support/BalancedPartitioningTest.cpp
+++ b/llvm/unittests/Support/BalancedPartitioningTest.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "llvm/Support/BalancedPartitioning.h"
+#include "llvm/Support/MathExtras.h"
#include "llvm/Testing/Support/SupportHelpers.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
@@ -37,6 +38,14 @@ class BalancedPartitioningTest : public ::testing::Test {
Ids.push_back(N.Id);
return Ids;
}
+
+ static std::vector<std::pair<BPFunctionNode::IDT, uint32_t>>
+ getBuckets(std::vector<BPFunctionNode> &Nodes) {
+ std::vector<std::pair<BPFunctionNode::IDT, uint32_t>> Res;
+ for (auto &N : Nodes)
+ Res.emplace_back(std::make_pair(N.Id, *N.getBucket()));
+ return Res;
+ }
};
TEST_F(BalancedPartitioningTest, Basic) {
@@ -97,4 +106,56 @@ TEST_F(BalancedPartitioningTest, MoveGain) {
30.f);
}
+TEST_F(BalancedPartitioningTest, Weight1) {
+ std::vector<BPFunctionNode::UtilityNodeT> UNs = {
+ BPFunctionNode::UtilityNodeT(0, 100),
+ BPFunctionNode::UtilityNodeT(1, 100),
+ BPFunctionNode::UtilityNodeT(2, 100),
+ BPFunctionNode::UtilityNodeT(3, 1),
+ BPFunctionNode::UtilityNodeT(4, 1),
+ };
+ std::vector<BPFunctionNode> Nodes = {
+ BPFunctionNode(0, {UNs[0], UNs[3]}), BPFunctionNode(1, {UNs[1], UNs[3]}),
+ BPFunctionNode(2, {UNs[2], UNs[3]}), BPFunctionNode(3, {UNs[0], UNs[4]}),
+ BPFunctionNode(4, {UNs[1], UNs[4]}), BPFunctionNode(5, {UNs[2], UNs[4]}),
+ };
+
+ Bp.run(Nodes);
+
+ // Verify that the created groups are [(0,3) -- (1,4) -- (2,5)].
+ auto Res = getBuckets(Nodes);
+ llvm::sort(Res);
+ EXPECT_THAT(AbsoluteDifference(Res[0].second, Res[3].second), 1);
+ EXPECT_THAT(AbsoluteDifference(Res[1].second, Res[4].second), 1);
+ EXPECT_THAT(AbsoluteDifference(Res[2].second, Res[5].second), 1);
+}
+
+TEST_F(BalancedPartitioningTest, Weight2) {
+ std::vector<BPFunctionNode::UtilityNodeT> UNs = {
+ BPFunctionNode::UtilityNodeT(0, 1),
+ BPFunctionNode::UtilityNodeT(1, 1),
+ BPFunctionNode::UtilityNodeT(2, 1),
+ BPFunctionNode::UtilityNodeT(3, 100),
+ BPFunctionNode::UtilityNodeT(4, 100),
+ };
+ std::vector<BPFunctionNode> Nodes = {
+ BPFunctionNode(0, {UNs[0], UNs[3]}), BPFunctionNode(1, {UNs[1], UNs[4]}),
+ BPFunctionNode(2, {UNs[2], UNs[3]}), BPFunctionNode(3, {UNs[0], UNs[4]}),
+ BPFunctionNode(4, {UNs[1], UNs[3]}), BPFunctionNode(5, {UNs[2], UNs[4]}),
+ };
+
+ Bp.run(Nodes);
+
+ // Verify that the created groups are [(0,2,4) -- (1,3,5)].
+ auto Res = getBuckets(Nodes);
+ llvm::sort(Res);
+ EXPECT_LE(AbsoluteDifference(Res[0].second, Res[2].second), 2u);
+ EXPECT_LE(AbsoluteDifference(Res[0].second, Res[4].second), 2u);
+ EXPECT_LE(AbsoluteDifference(Res[2].second, Res[4].second), 2u);
+
+ EXPECT_LE(AbsoluteDifference(Res[1].second, Res[3].second), 2u);
+ EXPECT_LE(AbsoluteDifference(Res[1].second, Res[5].second), 2u);
+ EXPECT_LE(AbsoluteDifference(Res[3].second, Res[5].second), 2u);
+}
+
} // end namespace llvm
|
|
@ellishg can you help me adjust |
ellishg
reviewed
Nov 20, 2023
llvm/lib/ProfileData/InstrProf.cpp
Outdated
| auto &Trace = Traces[TraceIdx].FunctionNameRefs; | ||
| for (size_t Timestamp = 0; Timestamp < Trace.size(); Timestamp++) { | ||
| for (int I = std::floor(std::log2(Timestamp + 1)); I < N; I++) { | ||
| auto &FunctionId = Trace[Timestamp]; | ||
| UtilityNodeT GroupId = TraceIdx * N + I; | ||
| UtilityNodeT GroupId = {TraceIdx * N + I}; |
There was a problem hiding this comment.
Suggested change
| UtilityNodeT GroupId = {TraceIdx * N + I}; | |
| UtilityNodeT GroupId(TraceIdx * N + I); |
I think a more explicit initialization would be better
llvm/lib/ProfileData/InstrProf.cpp
Outdated
| @@ -933,8 +933,15 @@ std::vector<BPFunctionNode> TemporalProfTraceTy::createBPFunctionNodes( | |||
| std::vector<BPFunctionNode> Nodes; | |||
| for (auto &Id : FunctionIds) { | |||
| auto &UNs = FuncGroups[Id]; | |||
| llvm::sort(UNs); | |||
| UNs.erase(std::unique(UNs.begin(), UNs.end()), UNs.end()); | |||
| llvm::sort(UNs.begin(), UNs.end(), | |||
There was a problem hiding this comment.
Suggested change
| llvm::sort(UNs.begin(), UNs.end(), | |
| llvm::sort(UNs, |
llvm/lib/ProfileData/InstrProf.cpp
Outdated
Comment on lines
940
to
950
| UNs.erase(std::unique(UNs.begin(), UNs.end(), | ||
| [](const UtilityNodeT &L, const UtilityNodeT &R) { | ||
| return L.id == R.id; | ||
| }), | ||
| UNs.end()); |
There was a problem hiding this comment.
Now that we have weights, do you think it makes sense to accumulate the weights of duplicated nodes instead? For example, if we had nodes {(A, 1), (B, 2), (B, 3), (C, 2)} then we would end up with {(A, 1), (B, 5), (C, 2)}.
I imagine this might be useful for compression. If a function has many duplicate instructions, it makes sense to weight that more than other instructions that appear once.
There was a problem hiding this comment.
That's a good suggestion but perhaps needs some experiments? In this diff I kept the behavior unchanged.
There was a problem hiding this comment.
Lets at least add a comment here so we don't forget.
There was a problem hiding this comment.
Also, lets factor this out into it's own helper function like void sortAndDeduplicate(SmallVector<UtilityNodeT> UNs).
| Signatures[UN.id].LeftCount++; | ||
| else | ||
| Signatures[UN.id].RightCount++; | ||
| Signatures[UN.id].Weight = UN.weight; |
There was a problem hiding this comment.
I think a signature could be affected by multiple utility nodes, each of which could have different weights. So I'm not sure this is exactly correct. In fact, I think with this change CachedGainLR is now dependent on which utility node moves between buckets, so we can no longer cache this value.
There was a problem hiding this comment.
All utility nodes with the same id must have the same weight -- we just store the weight together with the utility.
I re-wrote the code, please see if the new version makes more sense.
spupyrev
force-pushed
the
pika-bp-weights
branch
from
0d355d8
to
0c3051f
Compare
|
@ellishg any thoughts on this diff? |
Thanks for the reminder! I'll take a look next week. |
ellishg
reviewed
Jan 17, 2024
Comment on lines
65
to
66
| uint32_t id; | ||
| uint32_t weight; |
There was a problem hiding this comment.
Please use CamelCase for variable names.
https://llvm.org/docs/CodingStandards.html#name-types-functions-variables-and-enumerators-properly
Suggested change
| uint32_t id; | |
| uint32_t weight; | |
| uint32_t Id; | |
| uint32_t Weight; |
Comment on lines
210
to
217
| // Identical utility nodes (having the same UN.id) are guaranteed to have | ||
| // the same weight; thus, we can get a new weight only when the signature | ||
| // is uninitialized. | ||
| if (Signatures[UN.id].Weight != UN.weight) { | ||
| assert(Signatures[UN.id].Weight == 1); | ||
| Signatures[UN.id].Weight = UN.weight; |
There was a problem hiding this comment.
Since utility nodes usually come from hashes, it's possible utility nodes with the same id have different weights due to hash collisions. I don't think we need to worry too much about this case, but I'd like to remove assert and the word "guaranteed".
llvm/lib/ProfileData/InstrProf.cpp
Outdated
Comment on lines
940
to
950
| UNs.erase(std::unique(UNs.begin(), UNs.end(), | ||
| [](const UtilityNodeT &L, const UtilityNodeT &R) { | ||
| return L.id == R.id; | ||
| }), | ||
| UNs.end()); |
There was a problem hiding this comment.
Lets at least add a comment here so we don't forget.
spupyrev
force-pushed
the
pika-bp-weights
branch
from
0c3051f
to
1647d53
Compare
ellishg
reviewed
Jan 18, 2024
llvm/lib/ProfileData/InstrProf.cpp
Outdated
Comment on lines
908
to
921
| // Deduplicate utility nodes for a given function. | ||
| // TODO: One may experiment with accumulating the weights of duplicates. | ||
| void sortAndDeduplicate(SmallVector<BPFunctionNode::UtilityNodeT, 4> &UNs) { | ||
| using UtilityNodeT = BPFunctionNode::UtilityNodeT; | ||
| llvm::sort(UNs, [](const UtilityNodeT &L, const UtilityNodeT &R) { | ||
| return L.Id < R.Id; | ||
| }); | ||
| UNs.erase(std::unique(UNs.begin(), UNs.end(), | ||
| [](const UtilityNodeT &L, const UtilityNodeT &R) { | ||
| return L.Id == R.Id; | ||
| }), | ||
| UNs.end()); | ||
| } | ||
|
|
There was a problem hiding this comment.
Since this can be useful outside of InstrProf.cpp, I think we should move this to be a static function of UtilityNodeT.
Comment on lines
207
to
208
| if (Signatures[UN.Id].Weight != UN.Weight) | ||
| Signatures[UN.Id].Weight = UN.Weight; |
There was a problem hiding this comment.
Can we simply remove the if statement and always assign the weight?
ellishg
approved these changes
Jan 19, 2024
spupyrev
force-pushed
the
pika-bp-weights
branch
from
423c42b
to
1e48795
Compare
|
Hi @spupyrev your change seems to cause a build bot failure on Windows, can you take a look? https://lab.llvm.org/buildbot/#/builders/216/builds/33165 |
spupyrev
added a commit
to spupyrev/llvm-project
that referenced
this pull request
Jan 19, 2024
…lvm#72717)" This reverts commit 5954b9d due to broken Windows build
|
Reverted by 30aa9fb. Will investigate and re-commit later |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Adding weights to utility nodes in BP so that we can give more importance to
certain utilities. This is useful when we optimize several objectives jointly.