[llvm] [llvm][transforms] Add a new algorithm to SplitModule (PR #95941)
Ilia Sergachev via llvm-commits
llvm-commits at lists.llvm.org
Wed Jul 3 04:35:34 PDT 2024
https://github.com/sergachev updated https://github.com/llvm/llvm-project/pull/95941
>From 2bb3a7448ea9cf6104d7fb96dbebbaa0a03d0a2c Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev at nvidia.com>
Date: Wed, 5 Jun 2024 16:24:28 +0200
Subject: [PATCH 1/4] [Transforms] Add a new algorithm to SplitModule
The new algorithm augments the hash-based distribution of functions to
modules by pairing unmapped functions to empty modules before using the
hash-based distribution. It allows getting less empty modules when the
number of functions is close to the number of requested modules.
It's not in use by default and is available under a new flag.
---
.../llvm/Transforms/Utils/SplitModule.h | 2 +-
llvm/lib/Transforms/Utils/SplitModule.cpp | 34 ++++++++++++++++++-
.../tools/llvm-split/avoid-empty-modules.ll | 23 +++++++++++++
.../name-hash-based-distribution.ll | 17 ++++++++++
llvm/tools/llvm-split/llvm-split.cpp | 12 ++++++-
5 files changed, 85 insertions(+), 3 deletions(-)
create mode 100644 llvm/test/tools/llvm-split/avoid-empty-modules.ll
create mode 100644 llvm/test/tools/llvm-split/name-hash-based-distribution.ll
diff --git a/llvm/include/llvm/Transforms/Utils/SplitModule.h b/llvm/include/llvm/Transforms/Utils/SplitModule.h
index a5450738060a8..7b303172d1bb0 100644
--- a/llvm/include/llvm/Transforms/Utils/SplitModule.h
+++ b/llvm/include/llvm/Transforms/Utils/SplitModule.h
@@ -35,7 +35,7 @@ class Module;
void SplitModule(
Module &M, unsigned N,
function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback,
- bool PreserveLocals = false);
+ bool PreserveLocals = false, bool TryToAvoidEmptyModules = false);
} // end namespace llvm
diff --git a/llvm/lib/Transforms/Utils/SplitModule.cpp b/llvm/lib/Transforms/Utils/SplitModule.cpp
index 55db3737a1c09..7a9ea443cc2e2 100644
--- a/llvm/lib/Transforms/Utils/SplitModule.cpp
+++ b/llvm/lib/Transforms/Utils/SplitModule.cpp
@@ -254,7 +254,7 @@ static bool isInPartition(const GlobalValue *GV, unsigned I, unsigned N) {
void llvm::SplitModule(
Module &M, unsigned N,
function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback,
- bool PreserveLocals) {
+ bool PreserveLocals, bool TryToAvoidEmptyModules) {
if (!PreserveLocals) {
for (Function &F : M)
externalize(&F);
@@ -271,6 +271,38 @@ void llvm::SplitModule(
ClusterIDMapType ClusterIDMap;
findPartitions(M, ClusterIDMap, N);
+ // Find empty modules and functions not mapped to modules in ClusterIDMap.
+ // Map these functions to the empty modules so that they skip being
+ // distributed by isInPartition() based on function name hashes below.
+ // This provides better uniformity of distribution of functions to modules
+ // in some cases - for example when the number of functions equals to N.
+ if (TryToAvoidEmptyModules) {
+ DenseSet<unsigned> NonEmptyModules;
+ SmallVector<const GlobalValue *> UnmappedFunctions;
+ for (const auto &F : M.functions()) {
+ if (F.isDeclaration() ||
+ F.getLinkage() != GlobalValue::LinkageTypes::ExternalLinkage)
+ continue;
+ auto It = ClusterIDMap.find(&F);
+ if (It == ClusterIDMap.end())
+ UnmappedFunctions.push_back(&F);
+ else
+ NonEmptyModules.insert(It->second);
+ }
+ SmallVector<unsigned> EmptyModules;
+ for (unsigned I = 0; I < N; ++I) {
+ if (!NonEmptyModules.contains(I))
+ EmptyModules.push_back(I);
+ }
+ auto NextEmptyModuleIt = EmptyModules.begin();
+ for (const auto F : UnmappedFunctions) {
+ if (NextEmptyModuleIt == EmptyModules.end())
+ break;
+ ClusterIDMap.insert({F, *NextEmptyModuleIt});
+ ++NextEmptyModuleIt;
+ }
+ }
+
// FIXME: We should be able to reuse M as the last partition instead of
// cloning it. Note that the callers at the moment expect the module to
// be preserved, so will need some adjustments as well.
diff --git a/llvm/test/tools/llvm-split/avoid-empty-modules.ll b/llvm/test/tools/llvm-split/avoid-empty-modules.ll
new file mode 100644
index 0000000000000..654b8410a704f
--- /dev/null
+++ b/llvm/test/tools/llvm-split/avoid-empty-modules.ll
@@ -0,0 +1,23 @@
+; RUN: llvm-split -o %t %s -j 2 -avoid-empty-modules
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 %s
+
+; CHECK0-NOT: define
+; CHECK0: declare extern_weak void @e
+; CHECK0: define void @A
+; CHECK0-NOT: define
+
+; CHECK1-NOT: define
+; CHECK1: declare extern_weak void @e
+; CHECK1: define void @C
+; CHECK1-NOT: define
+
+declare extern_weak void @e(...)
+
+define void @A() {
+ ret void
+}
+
+define void @C() {
+ ret void
+}
diff --git a/llvm/test/tools/llvm-split/name-hash-based-distribution.ll b/llvm/test/tools/llvm-split/name-hash-based-distribution.ll
new file mode 100644
index 0000000000000..d9c4fc1eae0ad
--- /dev/null
+++ b/llvm/test/tools/llvm-split/name-hash-based-distribution.ll
@@ -0,0 +1,17 @@
+; RUN: llvm-split -o %t %s -j 2
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 %s
+
+; CHECK0-NOT: define
+
+; CHECK1-NOT: declare
+; CHECK1: define void @A
+; CHECK1: define void @C
+
+define void @A() {
+ ret void
+}
+
+define void @C() {
+ ret void
+}
diff --git a/llvm/tools/llvm-split/llvm-split.cpp b/llvm/tools/llvm-split/llvm-split.cpp
index 39a89cb1d2e75..786e92dd75084 100644
--- a/llvm/tools/llvm-split/llvm-split.cpp
+++ b/llvm/tools/llvm-split/llvm-split.cpp
@@ -53,6 +53,12 @@ static cl::opt<bool>
cl::desc("Split without externalizing locals"),
cl::cat(SplitCategory));
+static cl::opt<bool>
+ TryToAvoidEmptyModules("avoid-empty-modules", cl::Prefix, cl::init(false),
+ cl::desc("Try to avoid generating empty modules by "
+ "modifying the distribution of functions"),
+ cl::cat(SplitCategory));
+
static cl::opt<std::string>
MTriple("mtriple",
cl::desc("Target triple. When present, a TargetMachine is created "
@@ -122,6 +128,9 @@ int main(int argc, char **argv) {
errs() << "warning: -preserve-locals has no effect when using "
"TargetMachine::splitModule\n";
}
+ if (TryToAvoidEmptyModules)
+ errs() << "warning: -avoid-empty-modules has no effect when using "
+ "TargetMachine::splitModule\n";
if (TM->splitModule(*M, NumOutputs, HandleModulePart))
return 0;
@@ -131,6 +140,7 @@ int main(int argc, char **argv) {
"splitModule implementation\n";
}
- SplitModule(*M, NumOutputs, HandleModulePart, PreserveLocals);
+ SplitModule(*M, NumOutputs, HandleModulePart, PreserveLocals,
+ TryToAvoidEmptyModules);
return 0;
}
>From f59e616112cd112df402c3d5f414a1bba838be12 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev at nvidia.com>
Date: Tue, 2 Jul 2024 21:30:21 +0200
Subject: [PATCH 2/4] Skip redundant work if EmptyModules is empty
address review comment
---
llvm/lib/Transforms/Utils/SplitModule.cpp | 14 ++++++++------
1 file changed, 8 insertions(+), 6 deletions(-)
diff --git a/llvm/lib/Transforms/Utils/SplitModule.cpp b/llvm/lib/Transforms/Utils/SplitModule.cpp
index 7a9ea443cc2e2..3f8c2a5bcc8c4 100644
--- a/llvm/lib/Transforms/Utils/SplitModule.cpp
+++ b/llvm/lib/Transforms/Utils/SplitModule.cpp
@@ -294,12 +294,14 @@ void llvm::SplitModule(
if (!NonEmptyModules.contains(I))
EmptyModules.push_back(I);
}
- auto NextEmptyModuleIt = EmptyModules.begin();
- for (const auto F : UnmappedFunctions) {
- if (NextEmptyModuleIt == EmptyModules.end())
- break;
- ClusterIDMap.insert({F, *NextEmptyModuleIt});
- ++NextEmptyModuleIt;
+ if (!EmptyModules.empty()) {
+ auto NextEmptyModuleIt = EmptyModules.begin();
+ for (const auto F : UnmappedFunctions) {
+ if (NextEmptyModuleIt == EmptyModules.end())
+ break;
+ ClusterIDMap.insert({F, *NextEmptyModuleIt});
+ ++NextEmptyModuleIt;
+ }
}
}
>From 9ac67a42266d165693a27ee125740087f389bce1 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev at nvidia.com>
Date: Tue, 2 Jul 2024 21:41:28 +0200
Subject: [PATCH 3/4] Document flags of SplitModule()
address review comment
---
llvm/include/llvm/Transforms/Utils/SplitModule.h | 3 +++
1 file changed, 3 insertions(+)
diff --git a/llvm/include/llvm/Transforms/Utils/SplitModule.h b/llvm/include/llvm/Transforms/Utils/SplitModule.h
index 7b303172d1bb0..3f4ffaf2d6b7e 100644
--- a/llvm/include/llvm/Transforms/Utils/SplitModule.h
+++ b/llvm/include/llvm/Transforms/Utils/SplitModule.h
@@ -24,6 +24,9 @@ class Module;
/// Splits the module M into N linkable partitions. The function ModuleCallback
/// is called N times passing each individual partition as the MPart argument.
+/// PreserveLocals: Split without externalizing locals.
+/// TryToAvoidEmptyModules: Try to avoid generating empty modules by
+/// modifying the distribution of functions.
///
/// FIXME: This function does not deal with the somewhat subtle symbol
/// visibility issues around module splitting, including (but not limited to):
>From 22f0aebf1df1d5b323146641b1bd4bd8d7df8a67 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev at nvidia.com>
Date: Wed, 3 Jul 2024 12:05:11 +0200
Subject: [PATCH 4/4] Extend the algorithm to do round-robin distribution
completely override the hash-based distribution by mapping all unmapped
functions to modules using round-robin
---
.../llvm/Transforms/Utils/SplitModule.h | 6 +-
llvm/lib/Transforms/Utils/SplitModule.cpp | 62 ++++++++++---------
.../name-hash-based-distribution.ll | 14 ++++-
...{avoid-empty-modules.ll => round-robin.ll} | 14 ++++-
llvm/tools/llvm-split/llvm-split.cpp | 15 +++--
5 files changed, 67 insertions(+), 44 deletions(-)
rename llvm/test/tools/llvm-split/{avoid-empty-modules.ll => round-robin.ll} (68%)
diff --git a/llvm/include/llvm/Transforms/Utils/SplitModule.h b/llvm/include/llvm/Transforms/Utils/SplitModule.h
index 3f4ffaf2d6b7e..e7e8ee6279ace 100644
--- a/llvm/include/llvm/Transforms/Utils/SplitModule.h
+++ b/llvm/include/llvm/Transforms/Utils/SplitModule.h
@@ -25,8 +25,8 @@ class Module;
/// Splits the module M into N linkable partitions. The function ModuleCallback
/// is called N times passing each individual partition as the MPart argument.
/// PreserveLocals: Split without externalizing locals.
-/// TryToAvoidEmptyModules: Try to avoid generating empty modules by
-/// modifying the distribution of functions.
+/// RoundRobin: Use round-robin distribution of functions to modules instead
+/// of the default name-hash-based one.
///
/// FIXME: This function does not deal with the somewhat subtle symbol
/// visibility issues around module splitting, including (but not limited to):
@@ -38,7 +38,7 @@ class Module;
void SplitModule(
Module &M, unsigned N,
function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback,
- bool PreserveLocals = false, bool TryToAvoidEmptyModules = false);
+ bool PreserveLocals = false, bool RoundRobin = false);
} // end namespace llvm
diff --git a/llvm/lib/Transforms/Utils/SplitModule.cpp b/llvm/lib/Transforms/Utils/SplitModule.cpp
index 3f8c2a5bcc8c4..a30afadf0365d 100644
--- a/llvm/lib/Transforms/Utils/SplitModule.cpp
+++ b/llvm/lib/Transforms/Utils/SplitModule.cpp
@@ -55,6 +55,18 @@ using ClusterMapType = EquivalenceClasses<const GlobalValue *>;
using ComdatMembersType = DenseMap<const Comdat *, const GlobalValue *>;
using ClusterIDMapType = DenseMap<const GlobalValue *, unsigned>;
+bool compareClusters(const std::pair<unsigned, unsigned> &A,
+ const std::pair<unsigned, unsigned> &B) {
+ if (A.second || B.second)
+ return A.second > B.second;
+ return A.first > B.first;
+}
+
+using BalancingQueueType =
+ std::priority_queue<std::pair<unsigned, unsigned>,
+ std::vector<std::pair<unsigned, unsigned>>,
+ decltype(compareClusters) *>;
+
} // end anonymous namespace
static void addNonConstUser(ClusterMapType &GVtoClusterMap,
@@ -154,18 +166,7 @@ static void findPartitions(Module &M, ClusterIDMapType &ClusterIDMap,
// Assigned all GVs to merged clusters while balancing number of objects in
// each.
- auto CompareClusters = [](const std::pair<unsigned, unsigned> &a,
- const std::pair<unsigned, unsigned> &b) {
- if (a.second || b.second)
- return a.second > b.second;
- else
- return a.first > b.first;
- };
-
- std::priority_queue<std::pair<unsigned, unsigned>,
- std::vector<std::pair<unsigned, unsigned>>,
- decltype(CompareClusters)>
- BalancingQueue(CompareClusters);
+ BalancingQueueType BalancingQueue(compareClusters);
// Pre-populate priority queue with N slot blanks.
for (unsigned i = 0; i < N; ++i)
BalancingQueue.push(std::make_pair(i, 0));
@@ -254,7 +255,7 @@ static bool isInPartition(const GlobalValue *GV, unsigned I, unsigned N) {
void llvm::SplitModule(
Module &M, unsigned N,
function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback,
- bool PreserveLocals, bool TryToAvoidEmptyModules) {
+ bool PreserveLocals, bool RoundRobin) {
if (!PreserveLocals) {
for (Function &F : M)
externalize(&F);
@@ -271,13 +272,13 @@ void llvm::SplitModule(
ClusterIDMapType ClusterIDMap;
findPartitions(M, ClusterIDMap, N);
- // Find empty modules and functions not mapped to modules in ClusterIDMap.
- // Map these functions to the empty modules so that they skip being
- // distributed by isInPartition() based on function name hashes below.
+ // Find functions not mapped to modules in ClusterIDMap and count functions
+ // per module. Map unmapped functions using round-robin so that they skip
+ // being distributed by isInPartition() based on function name hashes below.
// This provides better uniformity of distribution of functions to modules
// in some cases - for example when the number of functions equals to N.
- if (TryToAvoidEmptyModules) {
- DenseSet<unsigned> NonEmptyModules;
+ if (RoundRobin) {
+ DenseMap<unsigned, unsigned> ModuleFunctionCount;
SmallVector<const GlobalValue *> UnmappedFunctions;
for (const auto &F : M.functions()) {
if (F.isDeclaration() ||
@@ -287,21 +288,22 @@ void llvm::SplitModule(
if (It == ClusterIDMap.end())
UnmappedFunctions.push_back(&F);
else
- NonEmptyModules.insert(It->second);
+ ++ModuleFunctionCount[It->second];
}
- SmallVector<unsigned> EmptyModules;
+ BalancingQueueType BalancingQueue(compareClusters);
for (unsigned I = 0; I < N; ++I) {
- if (!NonEmptyModules.contains(I))
- EmptyModules.push_back(I);
+ if (auto It = ModuleFunctionCount.find(I);
+ It != ModuleFunctionCount.end())
+ BalancingQueue.push(*It);
+ else
+ BalancingQueue.push({I, 0});
}
- if (!EmptyModules.empty()) {
- auto NextEmptyModuleIt = EmptyModules.begin();
- for (const auto F : UnmappedFunctions) {
- if (NextEmptyModuleIt == EmptyModules.end())
- break;
- ClusterIDMap.insert({F, *NextEmptyModuleIt});
- ++NextEmptyModuleIt;
- }
+ for (const auto *const F : UnmappedFunctions) {
+ const unsigned I = BalancingQueue.top().first;
+ const unsigned Count = BalancingQueue.top().second;
+ BalancingQueue.pop();
+ ClusterIDMap.insert({F, I});
+ BalancingQueue.push({I, Count + 1});
}
}
diff --git a/llvm/test/tools/llvm-split/name-hash-based-distribution.ll b/llvm/test/tools/llvm-split/name-hash-based-distribution.ll
index d9c4fc1eae0ad..b0e2d5eba12f9 100644
--- a/llvm/test/tools/llvm-split/name-hash-based-distribution.ll
+++ b/llvm/test/tools/llvm-split/name-hash-based-distribution.ll
@@ -2,16 +2,28 @@
; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s
; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 %s
+; CHECK0-NOT: define
+; CHECK0: define void @D
; CHECK0-NOT: define
-; CHECK1-NOT: declare
+; CHECK1-NOT: define
; CHECK1: define void @A
+; CHECK1: define void @B
; CHECK1: define void @C
+; CHECK1-NOT: define
define void @A() {
ret void
}
+define void @B() {
+ ret void
+}
+
define void @C() {
ret void
}
+
+define void @D() {
+ ret void
+}
diff --git a/llvm/test/tools/llvm-split/avoid-empty-modules.ll b/llvm/test/tools/llvm-split/round-robin.ll
similarity index 68%
rename from llvm/test/tools/llvm-split/avoid-empty-modules.ll
rename to llvm/test/tools/llvm-split/round-robin.ll
index 654b8410a704f..385c896888710 100644
--- a/llvm/test/tools/llvm-split/avoid-empty-modules.ll
+++ b/llvm/test/tools/llvm-split/round-robin.ll
@@ -1,15 +1,17 @@
-; RUN: llvm-split -o %t %s -j 2 -avoid-empty-modules
+; RUN: llvm-split -o %t %s -j 2 -round-robin
; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s
; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 %s
; CHECK0-NOT: define
; CHECK0: declare extern_weak void @e
; CHECK0: define void @A
+; CHECK0: define void @C
; CHECK0-NOT: define
; CHECK1-NOT: define
; CHECK1: declare extern_weak void @e
-; CHECK1: define void @C
+; CHECK1: define void @B
+; CHECK1: define void @D
; CHECK1-NOT: define
declare extern_weak void @e(...)
@@ -18,6 +20,14 @@ define void @A() {
ret void
}
+define void @B() {
+ ret void
+}
+
define void @C() {
ret void
}
+
+define void @D() {
+ ret void
+}
diff --git a/llvm/tools/llvm-split/llvm-split.cpp b/llvm/tools/llvm-split/llvm-split.cpp
index 786e92dd75084..c456403e6bc68 100644
--- a/llvm/tools/llvm-split/llvm-split.cpp
+++ b/llvm/tools/llvm-split/llvm-split.cpp
@@ -54,10 +54,10 @@ static cl::opt<bool>
cl::cat(SplitCategory));
static cl::opt<bool>
- TryToAvoidEmptyModules("avoid-empty-modules", cl::Prefix, cl::init(false),
- cl::desc("Try to avoid generating empty modules by "
- "modifying the distribution of functions"),
- cl::cat(SplitCategory));
+ RoundRobin("round-robin", cl::Prefix, cl::init(false),
+ cl::desc("Use round-robin distribution of functions to "
+ "modules instead of the default name-hash-based one"),
+ cl::cat(SplitCategory));
static cl::opt<std::string>
MTriple("mtriple",
@@ -128,8 +128,8 @@ int main(int argc, char **argv) {
errs() << "warning: -preserve-locals has no effect when using "
"TargetMachine::splitModule\n";
}
- if (TryToAvoidEmptyModules)
- errs() << "warning: -avoid-empty-modules has no effect when using "
+ if (RoundRobin)
+ errs() << "warning: -round-robin has no effect when using "
"TargetMachine::splitModule\n";
if (TM->splitModule(*M, NumOutputs, HandleModulePart))
@@ -140,7 +140,6 @@ int main(int argc, char **argv) {
"splitModule implementation\n";
}
- SplitModule(*M, NumOutputs, HandleModulePart, PreserveLocals,
- TryToAvoidEmptyModules);
+ SplitModule(*M, NumOutputs, HandleModulePart, PreserveLocals, RoundRobin);
return 0;
}
More information about the llvm-commits
mailing list