Add to the split module utility an SCC based method which allows not to globalize any local variables.

[slarin]

Summary

Currently llvm::SplitModule as the first step globalizes all local objects, which might not be desirable in some scenarios.
This change adds a new utility llvm::SplitModuleSelective that uses SCC approach to search for a balanced partition without the need to externalize symbols.
Such partition might not be possible or fully balanced for a given number of partitions, and is a function of the module properties (global/local dependencies within the module).

Joint development Tobias Edler von Koch ([e-mail]) and Sergei Larin ([e-mail])

[pcc]

As far as I can tell, it seems that in a typical LTO scenario with internalization, this partitioning would be strongly biased towards a single partition containing most of the code. Do you have a plan to address this?

I also have some lower level comments which will become more relevant if you can address the above.

include/llvm/CodeGen/ParallelCG.h

37–40

I don't think we need this extra flexibility. If this new way is better I think we should replace the existing SplitModule implementation with it.

lib/Transforms/Utils/SplitModule.cpp

47

I believe that what this function does is already implemented by llvm::EquivalenceClasses.

113–123

I think you can do all this stuff in one go with a loop that iterates over the whole map after visiting all the globals.

161

You can simplify this code by removing this line.

172–175

Same here

304

Should we do something better with module inline asm? Maybe add referenced globals to an equivalence class? (Probably doesn't need to happen in the initial version though.)

[mehdi_amini]

What is the issue with the current globalization? You mentioned that it "might not be desirable in some scenarios", can you elaborate?

include/llvm/CodeGen/ParallelCG.h

41

The added argument would probably benefit some documentation.

[tobiasvk]

In D16124#325170, @pcc wrote:

As far as I can tell, it seems that in a typical LTO scenario with internalization, this partitioning would be strongly biased towards a single partition containing most of the code. Do you have a plan to address this?

That's correct. Locally, we remember which symbols were originally non-local and restore their linkage type just prior to splitting / CodeGen. Constant merging and similar optimizations also reduce the ability to partition which is a bigger issue (that I haven't got a great solution for yet). But as it is, this is just as much about correctness as it is about getting the 'best' split.

In D16124#325188, @joker.eph wrote:

What is the issue with the current globalization? You mentioned that it "might not be desirable in some scenarios", can you elaborate?

Well, an obvious example would applying LTO to a library or to only part of your program. If locals become global, you may end up with name collisions or symbols you don't want to export.

[slarin]

Peter,

You are right in general - we have seen this bias. To address the issue, we re-globalize things that were internalized for LTO leaving only _real_ dependencies. If that step is interesting for anyone - we can definitely upstream it... but it will not go all the way. Constant merging sometimes introduces additional dependencies that might be counterproductive for splitting, and we are currently looking at our options to address that.

Nevertheless, without this patch the whole (split) feature is kind of useless for us, and I will try to explain why when I answer to Mehdi's question...

include/llvm/CodeGen/ParallelCG.h

37–40

I can see it both ways. In my tree I have one more change that is not (yet) reflected in this patch- partially I did not want to change several things at once... for that case having custom sorting function as argument makes a lot of sense.

On the other hand, I wanted to preserve as much from the original implementation as possible since i do not have a good visibility into all possible use cases.

I guess, if we can simply add "split mode" flag to splitCodeGen it might do the trick in fewer lines of code. Your are the owner of this code - in some way it is your call...

41

Do you mean in comments? Sure - based on the final form-factor we would want to use.

lib/Transforms/Utils/SplitModule.cpp

47

Maybe... I need to look if it can be used for this - one question would be the use of std::shared_ptr in this case - memory and time is an issue for us. This code is highly stressed.

If you are OK with the general concept - I can always improve it with additional patches - I am planning to stay here for a while :)

113–123

Yes, I looked into that, and it seemed to me that doing it like this was marginally faster. Is your concern with code clarity or implementation efficiency?

161

I am not sure about that - I will let Tobias to comment on this...

304

Generally agree... we can do it in following patches...

[mehdi_amini]

It looks like we should set visibility hidden when turning a local into a global.

[slarin]

Mehdi,

I think Tobias beat me to the answer - at the least we can have linker errors if we try to blindly globalize locals with identical names. Renaming them will not always work either, since our customers might rely on presence of specific symbols in the final image... There are also other reasons for that, but correctness of at least some corner cases seems like a hard justification on its own :-)

Welcome to the embedded world :) :) :)

Sergei

---

Qualcomm Innovation Center, Inc. is a member of Code Aurora Forum, hosted by The Linux Foundation

[slarin]

In the end it depends on what your linker does. If your linker tracks global symbols, turning locals into globals will not play well even with hidden attribute.

Besides, I want to "add" new functionality, not to change existing one (patch reflects that). It seems like a generic enough feature that people can find useful outside my use case...

---

Qualcomm Innovation Center, Inc. is a member of Code Aurora Forum, hosted by The Linux Foundation

[pcc]

Constant merging sometimes introduces additional dependencies that might be counterproductive for splitting, and we are currently looking at our options to address that.

How bad is the problem? I think that as long as your re-globalization is conservatively correct (i.e. does not pollute the global namespace) it would still be a useful contribution.

Nevertheless, without this patch the whole (split) feature is kind of useless for us

I can understand that, but to be frank, your patch isn't very useful for upstream without the re-globalization step.

include/llvm/CodeGen/ParallelCG.h

37–40

I can see it both ways. In my tree I have one more change that is not (yet) reflected in this patch- partially I did not want to change several things at once... for that case having custom sorting function as argument makes a lot of sense.

I'd like to understand why.

On the other hand, I wanted to preserve as much from the original implementation as possible since i do not have a good visibility into all possible use cases.

The reason for externalizing everything was mostly "implementation simplicity at the cost of some correctness". If we can fix the correctness problem then I'd rather not have two implementations.

lib/Transforms/Utils/SplitModule.cpp

47

It certainly seems that EquivalenceClasses would be more efficient than your custom implementation here, I'd also strongly prefer not to have a second implementation of equivalence classes in the tree if at all possible.

113–123

Mostly code clarity. If it's faster, that's all the better.

[slarin]

Peter,

Constant merging sometimes introduces additional dependencies that

might be counterproductive for splitting, and we are currently looking at our
options to address that.

How bad is the problem? I think that as long as your re-globalization is
conservatively correct (i.e. does not pollute the global namespace) it would
still be a useful contribution.

It is entirely property of your code. If you do have a lot of reusable constants in the LTO module, const merging acts as a spaghetti explosion connecting otherwise unrelated SCC. We are considering two ways for dealing with it - do not merge heuristic (likely to be slow), or re-split some constants smartly - this could be done here, as we build SCCs. For our case it works even as is, with acceptable misbalance, but can be improved.

I can understand that, but to be frank, your patch isn't very useful for
upstream without the re-globalization step.

I am glad you think so. We will try to upstream it ASAP. The only question would be - same patch or an independent feature? Under common flag?

If we can fix the correctness problem then I'd rather not have two implementations.

Ok. I can probably make it work with an enum flag for llvm::SplitModule or something.

Let me refactor the patch with EquivalenceClasses and single SplitModule function. Meanwhile Tobias will work on re-globalization patch. Stay tuned :)

Thanks.

Sergei

---

Qualcomm Innovation Center, Inc. is a member of Code Aurora Forum, hosted by The Linux Foundation

[slarin]

Peter,

Here is one issue with EquivalenceClasses. It uses

std::set<ECValue> TheMapping;

for members. We purposefully used SetVector since we are dealing with pointers... As a result I have non deterministic mapping - which is not cool.

I will see if I can do some sorting prior to partition assignment, but if it fails I do not think this will work.

If you have better suggestions - please speak up.

Sergei

[pcc]

The only question would be - same patch or an independent feature? Under common flag?

Up to you.

I will see if I can do some sorting prior to partition assignment, but if it fails I do not think this will work.

Yes, you will need to sort the equivalence classes. (It may also be beneficial to sort by partition size as well -- this is the LPT multiprocessor scheduling algorithm -- see https://en.wikipedia.org/wiki/Multiprocessor_scheduling). One thing that could help you is that the identity of the leader should be deterministic, so you can sort on its name.

[slarin]

Add to the split module utility an SCC based method which allows not to globalize any local variables.

Currently llvm::SplitModule as the first step globalizes all local objects, which might not be desirable in some scenarios.
This change adds a new utility llvm::SplitModuleSelective that uses SCC approach to search for a balanced partition without the need to externalize symbols.
Such partition might not be possible or fully balanced for a given number of partitions, and is a function of the module properties (global/local dependencies within the module).

Joint development Tobias Edler von Koch ([e-mail]) and Sergei Larin ([e-mail])

[slarin]

This would be even better if EquivalenceClasses would guarantee deterministic order... but anyhow - this should reflect Peter's comments.

The re-globalization patch will likely appear as an independent change later today.

[pcc]

lib/Transforms/Utils/SplitModule.cpp

63–78

I realized that you can reimplement this function in a simpler way that also avoids non-determinism in your loop (unordered_multimap can re-order values non-deterministically).

llvm::DenseMap<Comdat *, GlobalVariable *> ComdatMembers;
[...]
auto &Member = ComdatMembers[C];
if (Member)
  GVToClusterMap.unionSets(Member, GV);
else
  Member = GV;

161

Did you look into this?

186

If you store a ClusterMapType::member_iterator in Sets instead of a GlobalValue * you won't have to look it up again.

[slarin]

Excellent suggestions. New patch will be up shortly.

lib/Transforms/Utils/SplitModule.cpp

63–78

Yes, I like this!

161

Yes.. It is needed. Tobias is describing why now...

[slarin]

Implemented Peter's suggestions.

[tobiasvk]

lib/Transforms/Utils/SplitModule.cpp

161

This is to handle the case where a function uses e.g. a bitcast of a GlobalValue... that bitcast would be ConstantExpr, which is shared between functions, but is not a GlobalValue itself. Constants can be nested, hence the worklist.

Oh I see, perhaps you mean that we could just always add U to the worklist, and for GlobalValues we simply wouldn't 'recurse' any deeper. If you prefer that it's fine, it just adds extra work when we already know there's only ever going to be one item on the worklist.

[pcc]

LGTM with this one change. Thanks!

lib/Transforms/Utils/SplitModule.cpp

161

Sorry, I meant what you wrote in your second paragraph. It isn't all that much work and makes the code slightly simpler.

[tobiasvk]

In D16124#328371, @pcc wrote:

LGTM with this one change. Thanks!

Ok, no problem. Sergei will make the change.

[Eugene.Zelenko]

Committed in rL258083.

include/llvm/CodeGen/ParallelCG.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_PARALLELCG_H
 #define LLVM_CODEGEN_PARALLELCG_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Target/TargetMachine.h"
+#include <functional>
 
 namespace llvm {
 
 class Module;
 class TargetOptions;
 class raw_pwrite_stream;
 
 /// Split M into OSs.size() partitions, and generate code for each. Writes
 /// OSs.size() output files to the output streams in OSs. The resulting output
 /// files if linked together are intended to be equivalent to the single output
 /// file that would have been code generated from M.
 ///
 /// \returns M if OSs.size() == 1, otherwise returns std::unique_ptr<Module>().
-std::unique_ptr<Module>
-splitCodeGen(std::unique_ptr<Module> M, ArrayRef<raw_pwrite_stream *> OSs,
-             StringRef CPU, StringRef Features, const TargetOptions &Options,
-             Reloc::Model RM = Reloc::Default,
-             CodeModel::Model CM = CodeModel::Default,
-             CodeGenOpt::Level OL = CodeGenOpt::Default,
-             TargetMachine::CodeGenFileType FT = TargetMachine::CGFT_ObjectFile);
+std::unique_ptr<Module> splitCodeGen(
+    std::unique_ptr<Module> M, ArrayRef<raw_pwrite_stream *> OSs, StringRef CPU,
+    StringRef Features, const TargetOptions &Options,
+    std::function<
+        void(std::unique_ptr<Module> M, unsigned N,
+             std::function<void(std::unique_ptr<Module> MPart)> ModuleCallback)>
+        SplitFunction,
+    Reloc::Model RM = Reloc::Default, CodeModel::Model CM = CodeModel::Default,
 
 } // namespace llvm
 
 #endif

include/llvm/Transforms/Utils/SplitModule.h

 /// - Internal symbols should not collide with symbols defined outside the
 ///   module.
 /// - Internal symbols defined in module-level inline asm should be visible to
 ///   each partition.
 void SplitModule(
     std::unique_ptr<Module> M, unsigned N,
     std::function<void(std::unique_ptr<Module> MPart)> ModuleCallback);
 
+/// Similar to the above, but splitting without globalizing locals.
+void SplitModuleSelective(
+    std::unique_ptr<Module> M, unsigned N,
+    std::function<void(std::unique_ptr<Module> MPart)> ModuleCallback);
+
 } // End llvm namespace
 
 #endif

lib/CodeGen/ParallelCG.cpp

std::unique_ptr<TargetMachine> TM(TheTarget->createTargetMachine(
       M->getTargetTriple(), CPU, Features, Options, RM, CM, OL));
 
   legacy::PassManager CodeGenPasses;
   if (TM->addPassesToEmitFile(CodeGenPasses, OS, FileType))
     report_fatal_error("Failed to setup codegen");
   CodeGenPasses.run(*M);
 }
 
-std::unique_ptr<Module>
-llvm::splitCodeGen(std::unique_ptr<Module> M,
-                   ArrayRef<llvm::raw_pwrite_stream *> OSs, StringRef CPU,
-                   StringRef Features, const TargetOptions &Options,
-                   Reloc::Model RM, CodeModel::Model CM, CodeGenOpt::Level OL,
-                   TargetMachine::CodeGenFileType FileType) {
+std::unique_ptr<Module> llvm::splitCodeGen(
+    std::unique_ptr<Module> M, ArrayRef<llvm::raw_pwrite_stream *> OSs,
+    StringRef CPU, StringRef Features, const TargetOptions &Options,
+    std::function<
+        void(std::unique_ptr<Module> M, unsigned N,
+             std::function<void(std::unique_ptr<Module> MPart)> ModuleCallback)>
+        SplitFunction,
   StringRef TripleStr = M->getTargetTriple();
   std::string ErrMsg;
   const Target *TheTarget = TargetRegistry::lookupTarget(TripleStr, ErrMsg);
   if (!TheTarget)
     report_fatal_error(Twine("Target not found: ") + ErrMsg);
 
   if (OSs.size() == 1) {
     codegen(M.get(), *OSs[0], TheTarget, CPU, Features, Options, RM, CM,
             OL, FileType);
     return M;
   }
 
   std::vector<thread> Threads;
-  SplitModule(std::move(M), OSs.size(), [&](std::unique_ptr<Module> MPart) {
+  SplitFunction(std::move(M), OSs.size(), [&](std::unique_ptr<Module> MPart) {
     // We want to clone the module in a new context to multi-thread the codegen.
     // We do it by serializing partition modules to bitcode (while still on the
     // main thread, in order to avoid data races) and spinning up new threads
     // which deserialize the partitions into separate contexts.
     // FIXME: Provide a more direct way to do this in LLVM.
     SmallVector<char, 0> BC;
     raw_svector_ostream BCOS(BC);
     WriteBitcodeToFile(MPart.get(), BCOS);

lib/LTO/LTOCodeGenerator.cpp

 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetLowering.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Target/TargetRegisterInfo.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 #include "llvm/Transforms/ObjCARC.h"
+#include "llvm/Transforms/Utils/SplitModule.h"
 #include <system_error>
 using namespace llvm;
 
 const char* LTOCodeGenerator::getVersionString() {
 #ifdef LLVM_VERSION_INFO
   return PACKAGE_NAME " version " PACKAGE_VERSION ", " LLVM_VERSION_INFO;
 #else
   return PACKAGE_NAME " version " PACKAGE_VERSION;
bool LTOCodeGenerator::compileOptimized(ArrayRef<raw_pwrite_stream *> Out) {
 
   // Do code generation. We need to preserve the module in case the client calls
   // writeMergedModules() after compilation, but we only need to allow this at
   // parallelism level 1. This is achieved by having splitCodeGen return the
   // original module at parallelism level 1 which we then assign back to
   // MergedModule.
   MergedModule =
       splitCodeGen(std::move(MergedModule), Out, MCpu, FeatureStr, Options,
-                   RelocModel, CodeModel::Default, CGOptLevel, FileType);
+                   SplitModule, RelocModel, CodeModel::Default, CGOptLevel,
+                   FileType);
 
   return true;
 }
 
 /// setCodeGenDebugOptions - Set codegen debugging options to aid in debugging
 /// LTO problems.
 void LTOCodeGenerator::setCodeGenDebugOptions(const char *Options) {
   for (std::pair<StringRef, StringRef> o = getToken(Options); !o.first.empty();

lib/Transforms/Utils/SplitModule.cpp

 //===- SplitModule.cpp - Split a module into partitions -------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the function llvm::SplitModule, which splits a module
 // into multiple linkable partitions. It can be used to implement parallel code
 // generation for link-time optimization.
 //
 //===----------------------------------------------------------------------===//
 
+#define DEBUG_TYPE "split-module"
+
 #include "llvm/Transforms/Utils/SplitModule.h"
 #include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalObject.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include <queue>
+#include <set>
+#include <unordered_map>
 
 using namespace llvm;
 
+namespace {
+typedef SetVector<const GlobalValue *> GVSetType;
+typedef MapVector<const GlobalValue *, std::shared_ptr<GVSetType>>
+    ClusterMapType;
+typedef std::unordered_multimap<const Comdat *, const GlobalValue *>
+    ComdatMembersType;
+typedef DenseMap<const GlobalValue *, unsigned> ClusterIDMapType;
+}
+
+static void addGlobalValueUser(ClusterMapType &GVtoClusterMap,
+                               const GlobalValue *GV,
+                               const GlobalValue *UserGV) {
+  assert(GV->hasLocalLinkage() && "Only process local GVs here");
+
+  // This entry does not have a map yet...
+  if (!GVtoClusterMap.count(GV)) {
+    if (GVtoClusterMap.count(UserGV)) {
+      assert(!GVtoClusterMap[UserGV]->empty() &&
+             "There should not be an empty set here");
+      // But it's use already does, and it is not empty and is possibly shared
+      // by multiple GVs.
+      GVtoClusterMap[GV] = GVtoClusterMap[UserGV];
+      // Now add GV itself to the user set.
+      GVtoClusterMap[UserGV]->insert(GV);
+    } else {
+      // Create brand new mapping.
+      GVtoClusterMap[GV] = std::make_shared<GVSetType>();
+      GVtoClusterMap[GV]->insert(UserGV);
+      // Add the key to its own set.
+      GVtoClusterMap[GV]->insert(GV);
+      GVtoClusterMap[UserGV] = GVtoClusterMap[GV];
+    }
+  } else {
+    // This entry already exist, and it needs to be updated.
+    GVtoClusterMap[GV]->insert(UserGV);
+    if (GVtoClusterMap.count(UserGV)) {
+      // These two sets need to be merged.
+      if (GVtoClusterMap[UserGV] != GVtoClusterMap[GV]) {
+        GVtoClusterMap[GV]->insert(GVtoClusterMap[UserGV]->begin(),
+                                   GVtoClusterMap[UserGV]->end());
+        std::shared_ptr<GVSetType> OldSet = GVtoClusterMap[UserGV];
+
+        for (const GlobalValue *Entry : *OldSet)
+          if (GVtoClusterMap.count(Entry))
+            GVtoClusterMap[Entry] = GVtoClusterMap[GV];
+      }
+    } else {
+      GVtoClusterMap[UserGV] = GVtoClusterMap[GV];
+    }
+  }
+}
+
+static void addNonConstUser(ClusterMapType &GVtoClusterMap,
+                            const GlobalValue *GV, const User *U) {
+  assert((!isa<Constant>(U) || isa<GlobalValue>(U)) && "Bad user");
+
+  if (const Instruction *I = dyn_cast<Instruction>(U)) {
+    const GlobalValue *F = I->getParent()->getParent();
+    addGlobalValueUser(GVtoClusterMap, GV, F);
+  } else if (isa<GlobalAlias>(U) || isa<Function>(U) ||
+             isa<GlobalVariable>(U)) {
+    addGlobalValueUser(GVtoClusterMap, GV, cast<GlobalValue>(U));
+  } else {
+    llvm_unreachable("Underimplemented use case");
+  }
+}
+
+static void processComdat(ClusterMapType &GVtoClusterMap,
+                          ComdatMembersType &ComdatMembers,
+                          const GlobalValue *GV, const Comdat *C) {
+  // If this variable shares comdat marker with some other function, it is the
+  // same as having a "use" for our purposes.
+  // Note that it includes both external and local objects.
+  ComdatMembers.insert(std::make_pair(C, GV));
+
+  for (auto &&CM : make_range(ComdatMembers.equal_range(C))) {
+    if (CM.second == GV)
+      continue;
+    DEBUG(dbgs() << "In the same comdat: (" << GV->getName() << ") and ("
+                 << CM.second->getName() << ")\n");
+    if (GV->hasLocalLinkage()) {
+      addGlobalValueUser(GVtoClusterMap, GV, CM.second);
+    } else if (CM.second->hasLocalLinkage()) {
+      addGlobalValueUser(GVtoClusterMap, CM.second, GV);
+    }
+  }
+}
+
+// Find partitions for module in the way that no locals need to be
+// globalized.
+// Try to balance pack those partitions into N files since this roughly equals
+// thread balancing for the backend codegen step.
+static void findPartitions(Module *M, ClusterIDMapType &ClusterIDMap,
+                           unsigned N) {
+  // At this point module should have the proper mix of globals and locals.
+  // As we attempt to partition this module, we must not change any
+  // locals to globals.
+
+  DEBUG(dbgs() << "Partition module with (" << M->size() << ")functions\n");
+
+  ClusterMapType GVtoClusterMap;
+  ComdatMembersType ComdatMembers;
+
+  auto recordGVSet = [&GVtoClusterMap, &ComdatMembers](GlobalValue &GV) {
+    if (GV.isDeclaration())
+      return;
+
+    if (!GV.hasName())
+      GV.setName("__llvmsplit_unnamed");
+
+    // Comdat groups must not be partitioned. For comdat groups that contain
+    // locals, record all their members here so we can keep them together.
+    // Comdat groups that only contain external globals are already handled by
+    // the MD5-based partitioning.
+    if (const Comdat *C = GV.getComdat())
+      processComdat(GVtoClusterMap, ComdatMembers, &GV, C);
+
+    // Further only iterate over local GVs.
+    if (!GV.hasLocalLinkage())
+      return;
+
+    for (auto *U : GV.users()) {
+      // For each constant that is not a GV (a pure const):
+      if (isa<Constant>(U) && !isa<GlobalValue>(U)) {
+        SmallVector<const User *, 8> Worklist;
+        Worklist.push_back(U);
+        while (!Worklist.empty()) {
+          const User *UU = Worklist.pop_back_val();
+          if (isa<Constant>(UU) && !isa<GlobalValue>(UU)) {
+            Worklist.append(UU->user_begin(), UU->user_end());
+            continue;
+          }
+          addNonConstUser(GVtoClusterMap, &GV, UU);
+        }
+      } else {
+        // User is an instruction, alias or GlobalValue.
+        addNonConstUser(GVtoClusterMap, &GV, U);
+      }
+    }
+  };
+
+  std::for_each(M->begin(), M->end(), recordGVSet);
+  std::for_each(M->global_begin(), M->global_end(), recordGVSet);
+  std::for_each(M->alias_begin(), M->alias_end(), recordGVSet);
+  DEBUG(dbgs() << "Recorded " << GVtoClusterMap.size() << " local GVs.\n");
+
+  // 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)>
+      BalancinQueue(CompareClusters);
+  // Pre-populate priority queue with N slot blanks.
+  for (unsigned i = 0; i < N; ++i)
+    BalancinQueue.push(std::make_pair(i, 0));
+
+  SmallPtrSet<const GlobalValue *, 64> Visited;
+  for (auto &I : GVtoClusterMap) {
+    if (!Visited.insert(I.first).second)
+      continue;
+    // The top() always returns cluster ID with fewest number of entries.
+    unsigned CurrentClusterID = BalancinQueue.top().first;
+    unsigned CurrentClusterSize = BalancinQueue.top().second;
+    BalancinQueue.pop();
+
+    DEBUG(dbgs() << "Root[" << CurrentClusterID << "] size(" << I.second->size()
+                 << ") cluster_size(" << CurrentClusterSize << ") ----> "
+                 << I.first->getName() << "\n");
+
+    assert(I.second && I.second->count(I.first) && "Missing root in the set");
+    for (auto &II : *I.second) {
+      DEBUG(dbgs() << "----> " << II->getName()
+                   << (II->hasLocalLinkage() ? " l " : " e ") << "\n");
+      Visited.insert(II);
+      ClusterIDMap[II] = CurrentClusterID;
+    }
+    // Add this set size to the number of entries in this cluster.
+    BalancinQueue.push(std::make_pair(CurrentClusterID,
+                                      CurrentClusterSize + I.second->size()));
+  }
+}
+
 static void externalize(GlobalValue *GV) {
   if (GV->hasLocalLinkage()) {
     GV->setLinkage(GlobalValue::ExternalLinkage);
     GV->setVisibility(GlobalValue::HiddenVisibility);
   }
 
   // Unnamed entities must be named consistently between modules. setName will
   // give a distinct name to each such entity.
std::unique_ptr<Module> MPart(
         CloneModule(M.get(), VMap, [=](const GlobalValue *GV) {
           return isInPartition(GV, I, N);
         }));
     if (I != 0)
       MPart->setModuleInlineAsm("");
     ModuleCallback(std::move(MPart));
   }
 }
+
+// This performs splitting without a need for externalization, which might not
+// always be possible.
+void llvm::SplitModuleSelective(
+    std::unique_ptr<Module> M, unsigned N,
+    std::function<void(std::unique_ptr<Module> MPart)> ModuleCallback) {
+  ClusterIDMapType ClusterIDMap;
+  findPartitions(M.get(), ClusterIDMap, N);
+
+  for (unsigned I = 0; I < N; ++I) {
+    ValueToValueMapTy VMap;
+    std::unique_ptr<Module> MPart(
+        CloneModule(M.get(), VMap, [&](const GlobalValue *GV) {
+          if (ClusterIDMap.count(GV))
+            return (ClusterIDMap[GV] == I);
+          else
+            return isInPartition(GV, I, N);
+        }));
+    if (I != 0)
+      MPart->setModuleInlineAsm("");
+    ModuleCallback(std::move(MPart));
+  }
+}

test/tools/llvm-split/preserve-locals.ll

+; RUN: llvm-split -preserve-locals -o %t %s
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK0 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK1 %s
+
+; The local_var and local_func must not be separated.
+; CHECK0: @local_var
+; CHECK0: define internal fastcc void @local_func
+; The main and a must not be separated.
+; The main and local_func must not be together.
+; CHECK1: @a
+; CHECK1: define i32 @main
+; CHECK1: declare fastcc void @local_func
+
+@a = internal global i32 0, align 4
+@global_storage = common global i32 0, align 4
+@local_var = internal global i32 0, align 4
+
+; Function Attrs: nounwind
+define i32 @main(i32 %x) {
+entry:
+  %call = call fastcc i32 @foo(i32 %x, i32* nonnull @a)
+  %call1 = call fastcc i32 @baz(i32 %x)
+  %add = add nsw i32 %call, %call1
+  ret i32 %add
+}
+
+; Function Attrs: nounwind
+define fastcc i32 @bar(i32 %b) {
+entry:
+  %call = call fastcc i32 @baz(i32 %b)
+  ret i32 %call
+}
+
+; Function Attrs: nounwind
+define fastcc i32 @baz(i32 %x) {
+entry:
+  store i32 %x, i32* @global_storage, align 4
+  %shl = shl i32 %x, %x
+  ret i32 %shl
+}
+
+; Function Attrs: noinline nounwind
+define fastcc i32 @foo(i32 %a, i32* nocapture %b) {
+entry:
+  call fastcc void @local_func()
+  %call = call fastcc i32 @bar(i32 %a)
+  %0 = load i32, i32* @global_storage, align 4
+  %call1 = call fastcc i32 @baz(i32 %0)
+  %add = add nsw i32 %call, %call1
+  store i32 %add, i32* %b, align 4
+  %call.i = call fastcc i32 @baz(i32 %add) #2
+  %add.i = add nsw i32 %call.i, 2
+  %1 = load volatile i32, i32* @local_var, align 4
+  %add3 = add nsw i32 %add.i, %1
+  ret i32 %add3
+}
+
+; Function Attrs: noinline nounwind
+define internal fastcc void @local_func() section ".text" {
+entry:
+  %0 = load i32, i32* @global_storage, align 4
+  %call = call fastcc i32 @foo(i32 %0, i32* null)
+  store volatile i32 %call, i32* @local_var, align 4
+  ret void
+}

test/tools/llvm-split/scc-alias.ll

+; All of the functions in this module must end up
+; in the same partition without change of scope.
+; RUN: llvm-split -j=2 -preserve-locals -o %t %s
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK1 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK0 %s
+
+; CHECK0: declare i32 @funInternal
+; CHECK0: declare i32 @funExternal
+; CHECK0: declare i32 @funInternal2
+; CHECK0: declare i32 @funExternal2
+
+; All functions are in the same file.
+; Local functions are still local.
+; CHECK1: define internal i32 @funInternal
+; CHECK1: define i32 @funExternal
+; CHECK1: define internal i32 @funInternal2
+; CHECK1: define i32 @funExternal2
+
+
+@funInternalAlias = internal alias i32 (), i32 ()* @funInternal
+
+define internal i32 @funInternal() {
+entry:
+  ret i32 0
+}
+
+; Direct call to local alias
+
+define i32 @funExternal() {
+entry:
+  %x = call i32 @funInternalAlias()
+  ret i32 %x
+}
+
+; Call to local function that calls local alias
+
+define internal i32 @funInternal2() {
+entry:
+  %x = call i32 @funInternalAlias()
+  ret i32 %x
+}
+
+define i32 @funExternal2() {
+entry:
+  %x = call i32 @funInternal2()
+  ret i32 %x
+}
+

test/tools/llvm-split/scc-callchain.ll

+; All of the functions in this module must end up
+; in the same partition.
+
+; RUN: llvm-split -j=2 -preserve-locals -o %t %s
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK1 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK0 %s
+
+; CHECK0: declare i32 @funInternal0
+; CHECK0: declare i32 @funInternal1
+; CHECK0: declare i32 @funInternal2
+; CHECK0: declare i32 @funExternal
+
+; All functions are in the same file.
+; Local functions are still local.
+; CHECK1: define internal i32 @funInternal0
+; CHECK1: define internal i32 @funInternal1
+; CHECK1: define internal i32 @funInternal2
+; CHECK1: define i32 @funExternal
+; CHECK1: define i32 @funExternal2
+
+define internal i32 @funInternal0() {
+entry:
+  ret i32 0
+}
+
+define internal i32 @funInternal1() {
+entry:
+  %x = call i32 @funInternal0()
+  ret i32 %x
+}
+
+define internal i32 @funInternal2() {
+entry:
+  %x = call i32 @funInternal1()
+  ret i32 %x
+}
+
+define i32 @funExternal() {
+entry:
+  %x = call i32 @funInternal2()
+  ret i32 %x
+}
+
+define i32 @funExternal2() {
+entry:
+  %x = call i32 @funInternal0()
+  ret i32 %x
+}

test/tools/llvm-split/scc-comdat.ll

+; All functions in the same comdat group must
+; be in the same module
+
+; RUN: llvm-split -j=2 -preserve-locals -o %t %s
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK1 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK0 %s
+
+; CHECK0: declare i32 @fun1
+; CHECK0: declare i32 @fun2
+; CHECK0: declare i32 @fun3
+
+; CHECK1: define internal i32 @fun1
+; CHECK1: define internal i32 @fun2
+; CHECK1: define i32 @fun3
+
+$fun = comdat any
+
+define internal i32 @fun1() section ".text.funs" comdat($fun) {
+entry:
+  ret i32 0
+}
+
+define internal i32 @fun2() section ".text.funs" comdat($fun) {
+entry:
+  ret i32 0
+}
+
+define i32 @fun3() section ".text.funs" comdat($fun) {
+entry:
+  ret i32 0
+}
+

test/tools/llvm-split/scc-constants.ll

+; All of the functions in this module must end up
+; in the same partition.
+
+; RUN: llvm-split -j=2 -preserve-locals -o %t %s
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK1 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK0 %s
+
+; CHECK0: declare i32 @foo
+; CHECK0: declare i32 @baz
+; CHECK0: declare i32 @bar
+; CHECK0: declare i32 @bar2
+
+; CHECK1: @bla
+; CHECK1: @bla2
+; CHECK1: define internal i32 @foo
+; CHECK1: define internal i32 @baz
+; CHECK1: define i32 @bar
+; CHECK1: define i32 @bar2
+
+%struct.anon = type { i64, i64 }
+
+@bla = internal global %struct.anon { i64 1, i64 2 }, align 8
+@bla2 = internal global %struct.anon { i64 1, i64 2 }, align 8
+
+define internal i32 @foo() {
+entry:
+  store i64 5, i64* getelementptr inbounds (%struct.anon, %struct.anon* @bla, i32 0, i32 0), align 8
+  store i32 -1, i32* bitcast (i64* getelementptr inbounds (%struct.anon, %struct.anon* @bla2, i32 0, i32 1) to i32*), align 8
+  ret i32 0
+}
+
+define internal i32 @baz() {
+entry:
+  store i64 5, i64* getelementptr inbounds (%struct.anon, %struct.anon* @bla, i32 0, i32 0), align 8
+  store i32 -1, i32* bitcast (i64* getelementptr inbounds (%struct.anon, %struct.anon* @bla2, i32 0, i32 1) to i32*), align 8
+  ret i32 0
+}
+
+define i32 @bar() {
+  %call = call i32 @foo()
+  ret i32 0
+}
+
+define i32 @bar2() {
+  %call = call i32 @baz()
+  ret i32 0
+}
+

test/tools/llvm-split/scc-cycle.ll

+; All of the functions in this module must end up
+; in the same partition.
+
+; Mutually recursive calls
+; RUN: llvm-split -j=2 -preserve-locals -o %t %s
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK1 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK0 %s
+
+; CHECK0: declare i32 @funInternal0
+; CHECK0: declare i32 @funInternal1
+; CHECK0: declare i32 @funExternal0
+; CHECK0: declare i32 @funExternal1
+
+; CHECK1: define internal i32 @funInternal0
+; CHECK1: define internal i32 @funInternal1
+; CHECK1: define i32 @funExternal0
+; CHECK1: define i32 @funExternal1
+
+define internal i32 @funInternal0() {
+entry:
+  %x = call i32 @funInternal1()
+  ret i32 %x
+}
+
+define internal i32 @funInternal1() {
+entry:
+  %x = call i32 @funInternal0()
+  ret i32 %x
+}
+
+; Extrnal functions
+
+define i32 @funExternal0() {
+entry:
+  %x = call i32 @funInternal0()
+  ret i32 %x
+}
+
+define i32 @funExternal1() {
+entry:
+  %x = call i32 @funInternal1()
+  ret i32 %x
+}
+

test/tools/llvm-split/scc-global2global.ll

+; All of the functions and globals in this module must end up
+; in the same partition.
+
+; RUN: llvm-split -j=2 -preserve-locals -o %t %s
+; RUN: llvm-dis -o - %t0 | FileCheck --check-prefix=CHECK1 %s
+; RUN: llvm-dis -o - %t1 | FileCheck --check-prefix=CHECK0 %s
+
+; CHECK0: declare %struct.anon* @local0
+; CHECK0: declare i8** @local1
+
+; CHECK1: @bla
+; CHECK1: @ptr
+; CHECK1: define internal %struct.anon* @local0
+; CHECK1: define internal i8** @local1
+
+%struct.anon = type { i64, i64 }
+
+@bla = internal global %struct.anon { i64 1, i64 2 }, align 8
+@ptr = internal global i8* bitcast (%struct.anon* @bla to i8*), align 4
+
+define internal %struct.anon* @local0() {
+  ret %struct.anon* @bla
+}
+
+define internal i8** @local1() {
+  ret i8** @ptr
+}
+

tools/gold/gold-plugin.cpp

 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 #include "llvm/Transforms/Utils/GlobalStatus.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
+#include "llvm/Transforms/Utils/SplitModule.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 #include <list>
 #include <plugin-api.h>
 #include <system_error>
 #include <vector>
 
 #ifndef LDPO_PIE
 // FIXME: remove this declaration when we stop maintaining Ubuntu Quantal and
for (unsigned I = 0; I != options::Parallelism; ++I) {
           message(LDPL_FATAL, "Could not open file: %s", EC.message().c_str());
       }
       OSs.emplace_back(FD, true);
       OSPtrs[I] = &OSs.back();
     }
 
     // Run backend threads.
     splitCodeGen(std::move(M), OSPtrs, options::mcpu, Features.getString(),
-                 Options, RelocationModel, CodeModel::Default, CGOptLevel);
+                 Options, SplitModule, RelocationModel, CodeModel::Default,
+                 CGOptLevel);
   }
 
   for (auto &Filename : Filenames) {
     if (add_input_file(Filename.c_str()) != LDPS_OK)
       message(LDPL_FATAL,
               "Unable to add .o file to the link. File left behind in: %s",
               Filename.c_str());
     if (TempOutFile)

tools/llvm-split/llvm-split.cpp

 
 static cl::opt<std::string>
 OutputFilename("o", cl::desc("Override output filename"),
                cl::value_desc("filename"));
 
 static cl::opt<unsigned> NumOutputs("j", cl::Prefix, cl::init(2),
                                     cl::desc("Number of output files"));
 
+static cl::opt<bool>
+    PreserveLocals("preserve-locals", cl::Prefix, cl::init(false),
+                   cl::desc("Split without externalizing locals"));
+
 int main(int argc, char **argv) {
   LLVMContext &Context = getGlobalContext();
   SMDiagnostic Err;
   cl::ParseCommandLineOptions(argc, argv, "LLVM module splitter\n");
 
   std::unique_ptr<Module> M = parseIRFile(InputFilename, Err, Context);
 
   if (!M) {
     Err.print(argv[0], errs());
     return 1;
   }
 
   unsigned I = 0;
+  if (PreserveLocals) {
+    SplitModuleSelective(
+        std::move(M), NumOutputs, [&](std::unique_ptr<Module> MPart) {
+          std::error_code EC;
+          std::unique_ptr<tool_output_file> Out(new tool_output_file(
+              OutputFilename + utostr(I++), EC, sys::fs::F_None));
+          if (EC) {
+            errs() << EC.message() << '\n';
+            exit(1);
+          }
+
+          WriteBitcodeToFile(MPart.get(), Out->os());
+
+          Out->keep();
+        });
+  } else {
     SplitModule(std::move(M), NumOutputs, [&](std::unique_ptr<Module> MPart) {
       std::error_code EC;
       std::unique_ptr<tool_output_file> Out(new tool_output_file(
           OutputFilename + utostr(I++), EC, sys::fs::F_None));
       if (EC) {
         errs() << EC.message() << '\n';
         exit(1);
       }
 
       WriteBitcodeToFile(MPart.get(), Out->os());
 
       // Declare success.
       Out->keep();
     });
+  }
+
 
   return 0;
 }