310 lines · cpp
1//===- SplitModule.cpp - Split a module into partitions -------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file defines the function llvm::SplitModule, which splits a module10// into multiple linkable partitions. It can be used to implement parallel code11// generation for link-time optimization.12//13//===----------------------------------------------------------------------===//14 15#include "llvm/Transforms/Utils/SplitModule.h"16#include "llvm/ADT/DenseMap.h"17#include "llvm/ADT/EquivalenceClasses.h"18#include "llvm/ADT/SmallPtrSet.h"19#include "llvm/ADT/SmallVector.h"20#include "llvm/ADT/StringRef.h"21#include "llvm/IR/Comdat.h"22#include "llvm/IR/Constant.h"23#include "llvm/IR/Constants.h"24#include "llvm/IR/Function.h"25#include "llvm/IR/GlobalAlias.h"26#include "llvm/IR/GlobalObject.h"27#include "llvm/IR/GlobalValue.h"28#include "llvm/IR/GlobalVariable.h"29#include "llvm/IR/Instruction.h"30#include "llvm/IR/Module.h"31#include "llvm/IR/User.h"32#include "llvm/IR/Value.h"33#include "llvm/Support/Casting.h"34#include "llvm/Support/Debug.h"35#include "llvm/Support/ErrorHandling.h"36#include "llvm/Support/MD5.h"37#include "llvm/Support/raw_ostream.h"38#include "llvm/Transforms/Utils/Cloning.h"39#include "llvm/Transforms/Utils/ValueMapper.h"40#include <cassert>41#include <iterator>42#include <memory>43#include <queue>44#include <utility>45#include <vector>46 47using namespace llvm;48 49#define DEBUG_TYPE "split-module"50 51namespace {52 53using ClusterMapType = EquivalenceClasses<const GlobalValue *>;54using ComdatMembersType = DenseMap<const Comdat *, const GlobalValue *>;55using ClusterIDMapType = DenseMap<const GlobalValue *, unsigned>;56 57bool compareClusters(const std::pair<unsigned, unsigned> &A,58 const std::pair<unsigned, unsigned> &B) {59 if (A.second || B.second)60 return A.second > B.second;61 return A.first > B.first;62}63 64using BalancingQueueType =65 std::priority_queue<std::pair<unsigned, unsigned>,66 std::vector<std::pair<unsigned, unsigned>>,67 decltype(compareClusters) *>;68 69} // end anonymous namespace70 71static void addNonConstUser(ClusterMapType &GVtoClusterMap,72 const GlobalValue *GV, const User *U) {73 assert((!isa<Constant>(U) || isa<GlobalValue>(U)) && "Bad user");74 75 if (const Instruction *I = dyn_cast<Instruction>(U)) {76 const GlobalValue *F = I->getParent()->getParent();77 GVtoClusterMap.unionSets(GV, F);78 } else if (const GlobalValue *GVU = dyn_cast<GlobalValue>(U)) {79 GVtoClusterMap.unionSets(GV, GVU);80 } else {81 llvm_unreachable("Underimplemented use case");82 }83}84 85// Adds all GlobalValue users of V to the same cluster as GV.86static void addAllGlobalValueUsers(ClusterMapType &GVtoClusterMap,87 const GlobalValue *GV, const Value *V) {88 for (const auto *U : V->users()) {89 SmallVector<const User *, 4> Worklist;90 Worklist.push_back(U);91 while (!Worklist.empty()) {92 const User *UU = Worklist.pop_back_val();93 // For each constant that is not a GV (a pure const) recurse.94 if (isa<Constant>(UU) && !isa<GlobalValue>(UU)) {95 Worklist.append(UU->user_begin(), UU->user_end());96 continue;97 }98 addNonConstUser(GVtoClusterMap, GV, UU);99 }100 }101}102 103static const GlobalObject *getGVPartitioningRoot(const GlobalValue *GV) {104 const GlobalObject *GO = GV->getAliaseeObject();105 if (const auto *GI = dyn_cast_or_null<GlobalIFunc>(GO))106 GO = GI->getResolverFunction();107 return GO;108}109 110// Find partitions for module in the way that no locals need to be111// globalized.112// Try to balance pack those partitions into N files since this roughly equals113// thread balancing for the backend codegen step.114static void findPartitions(Module &M, ClusterIDMapType &ClusterIDMap,115 unsigned N) {116 // At this point module should have the proper mix of globals and locals.117 // As we attempt to partition this module, we must not change any118 // locals to globals.119 LLVM_DEBUG(dbgs() << "Partition module with (" << M.size()120 << ") functions\n");121 ClusterMapType GVtoClusterMap;122 ComdatMembersType ComdatMembers;123 124 auto recordGVSet = [&GVtoClusterMap, &ComdatMembers](GlobalValue &GV) {125 if (GV.isDeclaration())126 return;127 128 if (!GV.hasName())129 GV.setName("__llvmsplit_unnamed");130 131 // Comdat groups must not be partitioned. For comdat groups that contain132 // locals, record all their members here so we can keep them together.133 // Comdat groups that only contain external globals are already handled by134 // the MD5-based partitioning.135 if (const Comdat *C = GV.getComdat()) {136 auto &Member = ComdatMembers[C];137 if (Member)138 GVtoClusterMap.unionSets(Member, &GV);139 else140 Member = &GV;141 }142 143 // Aliases should not be separated from their aliasees and ifuncs should144 // not be separated from their resolvers regardless of linkage.145 if (const GlobalObject *Root = getGVPartitioningRoot(&GV))146 if (&GV != Root)147 GVtoClusterMap.unionSets(&GV, Root);148 149 if (const Function *F = dyn_cast<Function>(&GV)) {150 for (const BasicBlock &BB : *F) {151 BlockAddress *BA = BlockAddress::lookup(&BB);152 if (!BA || !BA->isConstantUsed())153 continue;154 addAllGlobalValueUsers(GVtoClusterMap, F, BA);155 }156 }157 158 if (GV.hasLocalLinkage())159 addAllGlobalValueUsers(GVtoClusterMap, &GV, &GV);160 };161 162 llvm::for_each(M.functions(), recordGVSet);163 llvm::for_each(M.globals(), recordGVSet);164 llvm::for_each(M.aliases(), recordGVSet);165 166 // Assigned all GVs to merged clusters while balancing number of objects in167 // each.168 BalancingQueueType BalancingQueue(compareClusters);169 // Pre-populate priority queue with N slot blanks.170 for (unsigned i = 0; i < N; ++i)171 BalancingQueue.push(std::make_pair(i, 0));172 173 SmallPtrSet<const GlobalValue *, 32> Visited;174 175 // To guarantee determinism, we have to sort SCC according to size.176 // When size is the same, use leader's name.177 for (const auto &C : GVtoClusterMap) {178 if (!C->isLeader())179 continue;180 181 unsigned CurrentClusterID = BalancingQueue.top().first;182 unsigned CurrentClusterSize = BalancingQueue.top().second;183 BalancingQueue.pop();184 185 LLVM_DEBUG(dbgs() << "Root[" << CurrentClusterID << "] cluster_size("186 << std::distance(GVtoClusterMap.member_begin(*C),187 GVtoClusterMap.member_end())188 << ") ----> " << C->getData()->getName() << "\n");189 190 for (ClusterMapType::member_iterator MI = GVtoClusterMap.findLeader(*C);191 MI != GVtoClusterMap.member_end(); ++MI) {192 if (!Visited.insert(*MI).second)193 continue;194 LLVM_DEBUG(dbgs() << "----> " << (*MI)->getName()195 << ((*MI)->hasLocalLinkage() ? " l " : " e ") << "\n");196 Visited.insert(*MI);197 ClusterIDMap[*MI] = CurrentClusterID;198 CurrentClusterSize++;199 }200 // Add this set size to the number of entries in this cluster.201 BalancingQueue.push(std::make_pair(CurrentClusterID, CurrentClusterSize));202 }203}204 205static void externalize(GlobalValue *GV) {206 if (GV->hasLocalLinkage()) {207 GV->setLinkage(GlobalValue::ExternalLinkage);208 GV->setVisibility(GlobalValue::HiddenVisibility);209 }210 211 // Unnamed entities must be named consistently between modules. setName will212 // give a distinct name to each such entity.213 if (!GV->hasName())214 GV->setName("__llvmsplit_unnamed");215}216 217// Returns whether GV should be in partition (0-based) I of N.218static bool isInPartition(const GlobalValue *GV, unsigned I, unsigned N) {219 if (const GlobalObject *Root = getGVPartitioningRoot(GV))220 GV = Root;221 222 StringRef Name;223 if (const Comdat *C = GV->getComdat())224 Name = C->getName();225 else226 Name = GV->getName();227 228 // Partition by MD5 hash. We only need a few bits for evenness as the number229 // of partitions will generally be in the 1-2 figure range; the low 16 bits230 // are enough.231 MD5 H;232 MD5::MD5Result R;233 H.update(Name);234 H.final(R);235 return (R[0] | (R[1] << 8)) % N == I;236}237 238void llvm::SplitModule(239 Module &M, unsigned N,240 function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback,241 bool PreserveLocals, bool RoundRobin) {242 if (!PreserveLocals) {243 for (Function &F : M)244 externalize(&F);245 for (GlobalVariable &GV : M.globals())246 externalize(&GV);247 for (GlobalAlias &GA : M.aliases())248 externalize(&GA);249 for (GlobalIFunc &GIF : M.ifuncs())250 externalize(&GIF);251 }252 253 // This performs splitting without a need for externalization, which might not254 // always be possible.255 ClusterIDMapType ClusterIDMap;256 findPartitions(M, ClusterIDMap, N);257 258 // Find functions not mapped to modules in ClusterIDMap and count functions259 // per module. Map unmapped functions using round-robin so that they skip260 // being distributed by isInPartition() based on function name hashes below.261 // This provides better uniformity of distribution of functions to modules262 // in some cases - for example when the number of functions equals to N.263 if (RoundRobin) {264 DenseMap<unsigned, unsigned> ModuleFunctionCount;265 SmallVector<const GlobalValue *> UnmappedFunctions;266 for (const auto &F : M.functions()) {267 if (F.isDeclaration() ||268 F.getLinkage() != GlobalValue::LinkageTypes::ExternalLinkage)269 continue;270 auto It = ClusterIDMap.find(&F);271 if (It == ClusterIDMap.end())272 UnmappedFunctions.push_back(&F);273 else274 ++ModuleFunctionCount[It->second];275 }276 BalancingQueueType BalancingQueue(compareClusters);277 for (unsigned I = 0; I < N; ++I) {278 if (auto It = ModuleFunctionCount.find(I);279 It != ModuleFunctionCount.end())280 BalancingQueue.push(*It);281 else282 BalancingQueue.push({I, 0});283 }284 for (const auto *const F : UnmappedFunctions) {285 const unsigned I = BalancingQueue.top().first;286 const unsigned Count = BalancingQueue.top().second;287 BalancingQueue.pop();288 ClusterIDMap.insert({F, I});289 BalancingQueue.push({I, Count + 1});290 }291 }292 293 // FIXME: We should be able to reuse M as the last partition instead of294 // cloning it. Note that the callers at the moment expect the module to295 // be preserved, so will need some adjustments as well.296 for (unsigned I = 0; I < N; ++I) {297 ValueToValueMapTy VMap;298 std::unique_ptr<Module> MPart(299 CloneModule(M, VMap, [&](const GlobalValue *GV) {300 if (auto It = ClusterIDMap.find(GV); It != ClusterIDMap.end())301 return It->second == I;302 else303 return isInPartition(GV, I, N);304 }));305 if (I != 0)306 MPart->setModuleInlineAsm("");307 ModuleCallback(std::move(MPart));308 }309}310