================
@@ -0,0 +1,360 @@
+#include "llvm/Transforms/Utils/SplitModuleCG.h"
+#include "llvm/Bitcode/BitcodeReader.h"
+#include "llvm/Bitcode/BitcodeWriter.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Value.h"
+#include "llvm/Support/MD5.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+#include <thread>
+using namespace llvm;
+
+#define DEBUG_TYPE "split-module-cg"
+
+namespace {
+
+static cl::opt<bool>
+    enablePrintSimplifiedCallGraph("enable-print-simplified-callgraph",
+                                   cl::Hidden, cl::init(false),
+                                   cl::desc("print SimplifiedCallGraph"));
+
+using PartitionID = unsigned;
+
+/// Returns whether duplicate definitions of \p F across partitions may be
+/// downgraded to available_externally. This is safe for external functions
+/// (either originally external or promoted by externalize), and for
+/// weak_odr/linkonce_odr functions whose equivalent definitions can be
+/// deduplicated to reduce codegen. Interposable linkages (weak/linkonce
+/// non-ODR) are excluded since downgrading them would change their
+/// optimization semantics.
+static bool canDowngradeToAvailableExternally(const Function &F) {
+  return !F.isDeclaration() &&
+         (F.hasExternalLinkage() || F.hasWeakODRLinkage() ||
+          F.hasLinkOnceODRLinkage());
+}
+
+} // namespace
+
+std::vector<DenseSet<const Function *>> SplitModuleCG::doPartitioning() {
+  LLVM_DEBUG(dbgs() << "\n--Partitioning Starts--\n");
+  // Performs all of the partitioning work on M.
+  assert(N != 0 && "Partition count must be at least 1");
+  std::vector<DenseSet<const Function *>> Partitions;
+  Partitions.resize(N);
+
+  auto ComparePartitions = [](const std::pair<PartitionID, CostType> &LHS,
+                              const std::pair<PartitionID, CostType> &RHS) {
+    // When two partitions have the same cost, assign to the one with the
+    // biggest ID first. This allows us to put things in P0 last, because P0 
may
+    // have other stuff added later.
+    if (LHS.second == RHS.second)
+      return LHS.first < RHS.first;
+    return LHS.second > RHS.second;
+  };
+
+  std::vector<std::pair<PartitionID, CostType>> BalancingQueue;
+  for (unsigned I = 0; I < N; ++I)
+    BalancingQueue.emplace_back(I, 0);
+
+  for (auto &CurFn : FWDWorkList) {
+    // Normal "load-balancing", assign to partition with least pressure.
+    auto [PID, _] = BalancingQueue.back();
+
+    // Insert the root function and its dependencies into the partition,
+    // tracking the cost of newly inserted functions so the balancing queue
+    // can be updated. CurFn.Dependencies includes the root F itself.
+    auto &FnsInPart = Partitions[PID];
+    CostType AddedCost = 0;
+    for (const Function *Dep : CurFn.Dependencies)
+      if (FnsInPart.insert(Dep).second)
+        AddedCost += FuncsCosts.lookup(Dep);
+
+    // Update the balancing queue. We scan backwards because in the common
+    // case the target partition is at the end of the sorted queue.
+    for (auto &[QueuePID, Cost] : reverse(BalancingQueue)) {
+      if (QueuePID != PID)
+        continue;
+      Cost += AddedCost;
+    }
+
+    sort(BalancingQueue, ComparePartitions);
+  }
+
+  return Partitions;
+}
+
+void SplitModuleCG::calculateFunctionCosts() {
+  ModuleCost = 0;
+  for (auto &Fn : M) {
+    if (Fn.isDeclaration())
+      continue;
+
+    CostType FnCost = 0;
+    for (const auto &BB : Fn)
+      FnCost += std::distance(BB.begin(), BB.end());
+    assert(FnCost != 0);
+    FuncsCosts[&Fn] = FnCost;
+    assert((ModuleCost + FnCost) >= ModuleCost && "Overflow!");
+    ModuleCost += FnCost;
+  }
+}
+
+void SplitModuleCG::dealWithMpart(Module &MPart, unsigned I) {
+  // Downgrade duplicate definitions of external functions to
+  // available_externally. The first partition to define such a function keeps
+  // the real definition; all other partitions get available_externally copies.
+  for (auto &PartFunc : MPart.functions()) {
+    if (PartFunc.isDeclaration())
+      continue;
+    // Look up the corresponding function in the original module M to check
+    // its ExternalFunction status.
+    auto *OrigFn = M.getFunction(PartFunc.getName());
+    if (!ExternalFunction.contains(OrigFn))
+      continue;
+    if (!ExternalFunction[OrigFn]) {
+      PartFunc.setLinkage(GlobalValue::AvailableExternallyLinkage);
+      PartFunc.setComdat(nullptr);
+    } else {
+      ExternalFunction[OrigFn] = false;
+    }
+  }
+
+  // Rename GlobalValues whose linkage was promoted from local to external,
+  // to avoid duplicate symbols across partitions in ThinLTO. Use the naming
+  // convention "name.llvm.<suffix>" so the promoted local cannot clash with
+  // an external that happens to share the same name. The suffix is derived
+  // from the module via getUniqueModuleId, so it is consistent across all
+  // partitions.
+  std::string Suffix = getUniqueModuleId(&M);
+  for (auto &GV : MPart.global_values()) {
+    // Only rename symbols that were promoted from local to external: skip
+    // those that are still local, and those that were already external in
+    // the source module (recorded in OriginalExternals).
+    if (GV.hasLocalLinkage() || OriginalExternals.contains(GV.getName()))
+      continue;
+    // Skip declarations of functions that were not explicitly externalized
+    // (e.g. skipped by the hasOneUse check). Their definitions in other
+    // partitions remain internal and are not renamed, so declarations must
+    // keep the original name to stay consistent.
+    auto *Fn = dyn_cast<Function>(&GV);
+    if (Fn && Fn->isDeclaration() &&
+        !ExternalFunction.contains(M.getFunction(Fn->getName())))
+      continue;
+    GV.setName((GV.getName() + ".llvm" + Suffix).str());
+  }
+
+#ifndef NDEBUG
+  LLVM_DEBUG(dbgs() << MPart.getModuleIdentifier() << "  : \n");
+  for (auto &F : MPart)
+    if (!F.isDeclaration())
+      LLVM_DEBUG(dbgs() << "   [Function: ] " << I << "  " << F.getName() << " 
"
+                        << F.getLinkage() << "\n");
+#endif
+}
+
+void SplitModuleCG::createWorkList() {
+  // First, find all the entry functions with an in-degree of 0
+  // (i.e., those that are not called by any function).
+  for (auto &SCGNode : SCG->values()) {
+    Function *F = SCGNode->getFunction();
+    if (F && SCGNode->getNumReferences() == 0)
+      EntryFuncs.insert(F);
+  }
+
+  // Second, find all the dependencies of each entry function.
+  for (auto *F : EntryFuncs) {
----------------
mmjjpp wrote:

Done

https://github.com/llvm/llvm-project/pull/198702
_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to