llvm/lib/Transforms/Utils/SplitModule.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-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file defines the function llvm::SplitModule, which splits a module
  10 // into multiple linkable partitions. It can be used to implement parallel code
  11 // 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
  47 using namespace llvm;
  48
  49 #define DEBUG_TYPE "split-module"
  50
  51 namespace {
  52
  53 using ClusterMapType = EquivalenceClasses<const GlobalValue *>;
  54 using ComdatMembersType = DenseMap<const Comdat *, const GlobalValue *>;
  55 using ClusterIDMapType = DenseMap<const GlobalValue *, unsigned>;
  56
  57 bool 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
  64 using BalancingQueueType =
  65     std::priority_queue<std::pair<unsigned, unsigned>,
  66                         std::vector<std::pair<unsigned, unsigned>>,
  67                         decltype(compareClusters) *>;
  68
  69 } // end anonymous namespace
  70
  71 static 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.
  86 static 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
 103 static 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 be
 111 // globalized.
 112 // Try to balance pack those partitions into N files since this roughly equals
 113 // thread balancing for the backend codegen step.
 114 static 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 any
 118   // 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 contain
 132     // locals, record all their members here so we can keep them together.
 133     // Comdat groups that only contain external globals are already handled by
 134     // 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       else
 140         Member = &GV;
 141     }
 142
 143     // Aliases should not be separated from their aliasees and ifuncs should
 144     // 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 in
 167   // 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   using SortType = std::pair<unsigned, ClusterMapType::iterator>;
 174
 175   SmallVector<SortType, 64> Sets;
 176   SmallPtrSet<const GlobalValue *, 32> Visited;
 177
 178   // To guarantee determinism, we have to sort SCC according to size.
 179   // When size is the same, use leader's name.
 180   for (ClusterMapType::iterator I = GVtoClusterMap.begin(),
 181                                 E = GVtoClusterMap.end();
 182        I != E; ++I)
 183     if (I->isLeader())
 184       Sets.push_back(
 185           std::make_pair(std::distance(GVtoClusterMap.member_begin(I),
 186                                        GVtoClusterMap.member_end()),
 187                          I));
 188
 189   llvm::sort(Sets, [](const SortType &a, const SortType &b) {
 190     if (a.first == b.first)
 191       return a.second->getData()->getName() > b.second->getData()->getName();
 192     else
 193       return a.first > b.first;
 194   });
 195
 196   for (auto &I : Sets) {
 197     unsigned CurrentClusterID = BalancingQueue.top().first;
 198     unsigned CurrentClusterSize = BalancingQueue.top().second;
 199     BalancingQueue.pop();
 200
 201     LLVM_DEBUG(dbgs() << "Root[" << CurrentClusterID << "] cluster_size("
 202                       << I.first << ") ----> " << I.second->getData()->getName()
 203                       << "\n");
 204
 205     for (ClusterMapType::member_iterator MI =
 206              GVtoClusterMap.findLeader(I.second);
 207          MI != GVtoClusterMap.member_end(); ++MI) {
 208       if (!Visited.insert(*MI).second)
 209         continue;
 210       LLVM_DEBUG(dbgs() << "----> " << (*MI)->getName()
 211                         << ((*MI)->hasLocalLinkage() ? " l " : " e ") << "\n");
 212       Visited.insert(*MI);
 213       ClusterIDMap[*MI] = CurrentClusterID;
 214       CurrentClusterSize++;
 215     }
 216     // Add this set size to the number of entries in this cluster.
 217     BalancingQueue.push(std::make_pair(CurrentClusterID, CurrentClusterSize));
 218   }
 219 }
 220
 221 static void externalize(GlobalValue *GV) {
 222   if (GV->hasLocalLinkage()) {
 223     GV->setLinkage(GlobalValue::ExternalLinkage);
 224     GV->setVisibility(GlobalValue::HiddenVisibility);
 225   }
 226
 227   // Unnamed entities must be named consistently between modules. setName will
 228   // give a distinct name to each such entity.
 229   if (!GV->hasName())
 230     GV->setName("__llvmsplit_unnamed");
 231 }
 232
 233 // Returns whether GV should be in partition (0-based) I of N.
 234 static bool isInPartition(const GlobalValue *GV, unsigned I, unsigned N) {
 235   if (const GlobalObject *Root = getGVPartitioningRoot(GV))
 236     GV = Root;
 237
 238   StringRef Name;
 239   if (const Comdat *C = GV->getComdat())
 240     Name = C->getName();
 241   else
 242     Name = GV->getName();
 243
 244   // Partition by MD5 hash. We only need a few bits for evenness as the number
 245   // of partitions will generally be in the 1-2 figure range; the low 16 bits
 246   // are enough.
 247   MD5 H;
 248   MD5::MD5Result R;
 249   H.update(Name);
 250   H.final(R);
 251   return (R[0] | (R[1] << 8)) % N == I;
 252 }
 253
 254 void llvm::SplitModule(
 255     Module &M, unsigned N,
 256     function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback,
 257     bool PreserveLocals, bool RoundRobin) {
 258   if (!PreserveLocals) {
 259     for (Function &F : M)
 260       externalize(&F);
 261     for (GlobalVariable &GV : M.globals())
 262       externalize(&GV);
 263     for (GlobalAlias &GA : M.aliases())
 264       externalize(&GA);
 265     for (GlobalIFunc &GIF : M.ifuncs())
 266       externalize(&GIF);
 267   }
 268
 269   // This performs splitting without a need for externalization, which might not
 270   // always be possible.
 271   ClusterIDMapType ClusterIDMap;
 272   findPartitions(M, ClusterIDMap, N);
 273
 274   // Find functions not mapped to modules in ClusterIDMap and count functions
 275   // per module. Map unmapped functions using round-robin so that they skip
 276   // being distributed by isInPartition() based on function name hashes below.
 277   // This provides better uniformity of distribution of functions to modules
 278   // in some cases - for example when the number of functions equals to N.
 279   if (RoundRobin) {
 280     DenseMap<unsigned, unsigned> ModuleFunctionCount;
 281     SmallVector<const GlobalValue *> UnmappedFunctions;
 282     for (const auto &F : M.functions()) {
 283       if (F.isDeclaration() ||
 284           F.getLinkage() != GlobalValue::LinkageTypes::ExternalLinkage)
 285         continue;
 286       auto It = ClusterIDMap.find(&F);
 287       if (It == ClusterIDMap.end())
 288         UnmappedFunctions.push_back(&F);
 289       else
 290         ++ModuleFunctionCount[It->second];
 291     }
 292     BalancingQueueType BalancingQueue(compareClusters);
 293     for (unsigned I = 0; I < N; ++I) {
 294       if (auto It = ModuleFunctionCount.find(I);
 295           It != ModuleFunctionCount.end())
 296         BalancingQueue.push(*It);
 297       else
 298         BalancingQueue.push({I, 0});
 299     }
 300     for (const auto *const F : UnmappedFunctions) {
 301       const unsigned I = BalancingQueue.top().first;
 302       const unsigned Count = BalancingQueue.top().second;
 303       BalancingQueue.pop();
 304       ClusterIDMap.insert({F, I});
 305       BalancingQueue.push({I, Count + 1});
 306     }
 307   }
 308
 309   // FIXME: We should be able to reuse M as the last partition instead of
 310   // cloning it. Note that the callers at the moment expect the module to
 311   // be preserved, so will need some adjustments as well.
 312   for (unsigned I = 0; I < N; ++I) {
 313     ValueToValueMapTy VMap;
 314     std::unique_ptr<Module> MPart(
 315         CloneModule(M, VMap, [&](const GlobalValue *GV) {
 316           if (auto It = ClusterIDMap.find(GV); It != ClusterIDMap.end())
 317             return It->second == I;
 318           else
 319             return isInPartition(GV, I, N);
 320         }));
 321     if (I != 0)
 322       MPart->setModuleInlineAsm("");
 323     ModuleCallback(std::move(MPart));
 324   }
 325 }