[llvm] [RFC][LTO] Allow target-specific module splittting (PR #83128)
Pierre van Houtryve via llvm-commits
llvm-commits at lists.llvm.org
Tue Feb 27 04:26:51 PST 2024
https://github.com/Pierre-vh created https://github.com/llvm/llvm-project/pull/83128
Allow targets to implement custom module splitting logic for --lto-partitions.
>From c950e8a84d5a450dfb051eda0a2bfa48921deeda Mon Sep 17 00:00:00 2001
From: pvanhout <pierre.vanhoutryve at amd.com>
Date: Tue, 27 Feb 2024 13:26:18 +0100
Subject: [PATCH] [RFC][LTO] Allow target-specific module splittting
Allow targets to implement custom module splitting logic for
--lto-partitions.
---
llvm/include/llvm/Target/TargetMachine.h | 10 +++
llvm/lib/LTO/LTOBackend.cpp | 12 ++-
llvm/tools/llvm-split/CMakeLists.txt | 7 ++
llvm/tools/llvm-split/llvm-split.cpp | 102 +++++++++++++++++------
4 files changed, 102 insertions(+), 29 deletions(-)
diff --git a/llvm/include/llvm/Target/TargetMachine.h b/llvm/include/llvm/Target/TargetMachine.h
index d7ce088cad49f6..43342c3e474615 100644
--- a/llvm/include/llvm/Target/TargetMachine.h
+++ b/llvm/include/llvm/Target/TargetMachine.h
@@ -418,6 +418,16 @@ class TargetMachine {
virtual unsigned getAddressSpaceForPseudoSourceKind(unsigned Kind) const {
return 0;
}
+
+ /// Entry point for module splitting. Targets can implement custom module
+ /// splitting logic, mainly used by LTO for --lto-partitions.
+ ///
+ /// \returns `true` if the module was split, `false` otherwise. When `false` is returned, it is assumed that \p ModuleCallback has never been called and \p M has not been modified.
+ virtual bool splitModule(
+ Module &M, unsigned NumParts,
+ function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback) const {
+ return false;
+ }
};
/// This class describes a target machine that is implemented with the LLVM
diff --git a/llvm/lib/LTO/LTOBackend.cpp b/llvm/lib/LTO/LTOBackend.cpp
index 6cfe67779b1a7d..494a8b0798900b 100644
--- a/llvm/lib/LTO/LTOBackend.cpp
+++ b/llvm/lib/LTO/LTOBackend.cpp
@@ -436,8 +436,7 @@ static void splitCodeGen(const Config &C, TargetMachine *TM,
unsigned ThreadCount = 0;
const Target *T = &TM->getTarget();
- SplitModule(
- Mod, ParallelCodeGenParallelismLevel,
+ const auto HandleModulePartition =
[&](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
@@ -469,8 +468,13 @@ static void splitCodeGen(const Config &C, TargetMachine *TM,
// Pass BC using std::move to ensure that it get moved rather than
// copied into the thread's context.
std::move(BC), ThreadCount++);
- },
- false);
+ };
+
+ // Try target-specific module splitting first, then fallback to the default.
+ if (!TM->splitModule(Mod, ParallelCodeGenParallelismLevel,
+ HandleModulePartition)) {
+ SplitModule(Mod, ParallelCodeGenParallelismLevel, HandleModulePartition, false);
+ }
// Because the inner lambda (which runs in a worker thread) captures our local
// variables, we need to wait for the worker threads to terminate before we
diff --git a/llvm/tools/llvm-split/CMakeLists.txt b/llvm/tools/llvm-split/CMakeLists.txt
index 52eedeb9f53f32..0579298462d113 100644
--- a/llvm/tools/llvm-split/CMakeLists.txt
+++ b/llvm/tools/llvm-split/CMakeLists.txt
@@ -1,9 +1,16 @@
set(LLVM_LINK_COMPONENTS
+ AllTargetsAsmParsers
+ AllTargetsCodeGens
+ AllTargetsDescs
+ AllTargetsInfos
TransformUtils
BitWriter
+ CodeGen
Core
IRReader
+ MC
Support
+ Target
)
add_llvm_tool(llvm-split
diff --git a/llvm/tools/llvm-split/llvm-split.cpp b/llvm/tools/llvm-split/llvm-split.cpp
index c6e20e0373c717..f463135c67504f 100644
--- a/llvm/tools/llvm-split/llvm-split.cpp
+++ b/llvm/tools/llvm-split/llvm-split.cpp
@@ -1,4 +1,4 @@
-//===-- llvm-split: command line tool for testing module splitter ---------===//
+//===-- llvm-split: command line tool for testing module splitting --------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,7 +6,8 @@
//
//===----------------------------------------------------------------------===//
//
-// This program can be used to test the llvm::SplitModule function.
+// This program can be used to test the llvm::SplitModule and
+// TargetMachine::splitModule functions.
//
//===----------------------------------------------------------------------===//
@@ -15,12 +16,17 @@
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Verifier.h"
#include "llvm/IRReader/IRReader.h"
+#include "llvm/MC/TargetRegistry.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/ToolOutputFile.h"
-#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/WithColor.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/TargetParser/Triple.h"
#include "llvm/Transforms/Utils/SplitModule.h"
using namespace llvm;
@@ -47,12 +53,45 @@ static cl::opt<bool>
cl::desc("Split without externalizing locals"),
cl::cat(SplitCategory));
+static cl::opt<std::string>
+ MTarget("mtarget",
+ cl::desc("Target triple. When present, a TargetMachine is created "
+ "and TargetMachine::splitModule is used instead of the "
+ "common SplitModule logic."),
+ cl::value_desc("triple"), cl::cat(SplitCategory));
+
+static cl::opt<std::string>
+ MCPU("mcpu", cl::desc("Target CPU, ignored if -mtarget is not used"),
+ cl::value_desc("cpu"), cl::cat(SplitCategory));
+
int main(int argc, char **argv) {
+ InitLLVM X(argc, argv);
+
+ // NOTE: If mtarget is not present we could avoid initializing targets to save
+ // time, but this is a testing tool and it's likely not worth the added
+ // complexity.
+ InitializeAllTargets();
+ InitializeAllTargetMCs();
+
LLVMContext Context;
SMDiagnostic Err;
cl::HideUnrelatedOptions({&SplitCategory, &getColorCategory()});
cl::ParseCommandLineOptions(argc, argv, "LLVM module splitter\n");
+ TargetMachine *TM = nullptr;
+ if (!MTarget.empty()) {
+ std::string Error;
+ const Target *T = TargetRegistry::lookupTarget(MTarget, Error);
+ if (!T) {
+ errs() << "unknown target '" << MTarget << "': " << Error << "\n";
+ return 1;
+ }
+
+ TargetOptions Options;
+ TM = T->createTargetMachine(MTarget, MCPU, /*FS*/ "", Options, std::nullopt,
+ std::nullopt);
+ }
+
std::unique_ptr<Module> M = parseIRFile(InputFilename, Err, Context);
if (!M) {
@@ -61,28 +100,41 @@ int main(int argc, char **argv) {
}
unsigned I = 0;
- SplitModule(
- *M, NumOutputs,
- [&](std::unique_ptr<Module> MPart) {
- std::error_code EC;
- std::unique_ptr<ToolOutputFile> Out(new ToolOutputFile(
- OutputFilename + utostr(I++), EC, sys::fs::OF_None));
- if (EC) {
- errs() << EC.message() << '\n';
- exit(1);
- }
-
- if (verifyModule(*MPart, &errs())) {
- errs() << "Broken module!\n";
- exit(1);
- }
-
- WriteBitcodeToFile(*MPart, Out->os());
-
- // Declare success.
- Out->keep();
- },
- PreserveLocals);
+ const auto HandleModulePart = [&](std::unique_ptr<Module> MPart) {
+ std::error_code EC;
+ std::unique_ptr<ToolOutputFile> Out(
+ new ToolOutputFile(OutputFilename + utostr(I++), EC, sys::fs::OF_None));
+ if (EC) {
+ errs() << EC.message() << '\n';
+ exit(1);
+ }
+
+ if (verifyModule(*MPart, &errs())) {
+ errs() << "Broken module!\n";
+ exit(1);
+ }
+
+ WriteBitcodeToFile(*MPart, Out->os());
+
+ // Declare success.
+ Out->keep();
+ };
+
+ if (!TM) {
+ SplitModule(*M, NumOutputs, HandleModulePart, PreserveLocals);
+ return 0;
+ }
+
+ if (PreserveLocals) {
+ errs() << "warning: -preserve-locals has no effect when using "
+ "TargetMachine::splitModule\n";
+ }
+
+ if (!TM->splitModule(*M, NumOutputs, HandleModulePart)) {
+ errs() << "error: '" << MTarget
+ << "' does not implement TargetMachine::splitModule\n";
+ return 1;
+ }
return 0;
}
More information about the llvm-commits
mailing list