diff --git a/llvm/include/llvm/Target/TargetMachine.h b/llvm/include/llvm/Target/TargetMachine.h index ceb371bdc734803..48ea3cfe02775bc 100644 --- a/llvm/include/llvm/Target/TargetMachine.h +++ b/llvm/include/llvm/Target/TargetMachine.h @@ -418,6 +418,18 @@ 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 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 71e8849dc3cc915..d4b89ede2d7134f 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 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,14 @@ 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/test/tools/llvm-split/target-specific-split.ll b/llvm/test/tools/llvm-split/target-specific-split.ll new file mode 100644 index 000000000000000..ea4d38084d759a7 --- /dev/null +++ b/llvm/test/tools/llvm-split/target-specific-split.ll @@ -0,0 +1,10 @@ +; RUN: llvm-split -o %t %s -mtriple x86_64 -preserve-locals 2>&1 | FileCheck %s + +; Basic test for a target that doesn't support target-specific module splitting. + +; CHECK: warning: -preserve-locals has no effect when using TargetMachine::splitModule +; CHECK: warning: TargetMachine::splitModule failed, falling back to default splitModule implementation + +define void @bar() { + ret void +} diff --git a/llvm/tools/llvm-split/CMakeLists.txt b/llvm/tools/llvm-split/CMakeLists.txt index 52eedeb9f53f32f..0579298462d1134 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 c6e20e0373c7178..41f392bcc53117b 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,42 @@ static cl::opt cl::desc("Split without externalizing locals"), cl::cat(SplitCategory)); +static cl::opt + MTriple("mtriple", + 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 + MCPU("mcpu", cl::desc("Target CPU, ignored if -mtriple is not used"), + cl::value_desc("cpu"), cl::cat(SplitCategory)); + int main(int argc, char **argv) { + InitLLVM X(argc, argv); + LLVMContext Context; SMDiagnostic Err; cl::HideUnrelatedOptions({&SplitCategory, &getColorCategory()}); cl::ParseCommandLineOptions(argc, argv, "LLVM module splitter\n"); + TargetMachine *TM = nullptr; + if (!MTriple.empty()) { + InitializeAllTargets(); + InitializeAllTargetMCs(); + + std::string Error; + const Target *T = TargetRegistry::lookupTarget(MTriple, Error); + if (!T) { + errs() << "unknown target '" << MTriple << "': " << Error << "\n"; + return 1; + } + + TargetOptions Options; + TM = T->createTargetMachine(MTriple, MCPU, /*FS*/ "", Options, std::nullopt, + std::nullopt); + } + std::unique_ptr M = parseIRFile(InputFilename, Err, Context); if (!M) { @@ -61,28 +97,40 @@ int main(int argc, char **argv) { } unsigned I = 0; - SplitModule( - *M, NumOutputs, - [&](std::unique_ptr MPart) { - std::error_code EC; - std::unique_ptr 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 MPart) { + std::error_code EC; + std::unique_ptr 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) { + if (PreserveLocals) { + errs() << "warning: -preserve-locals has no effect when using " + "TargetMachine::splitModule\n"; + } + + if (TM->splitModule(*M, NumOutputs, HandleModulePart)) + return 0; + + errs() << "warning: " + "TargetMachine::splitModule failed, falling back to default " + "splitModule implementation\n"; + } + SplitModule(*M, NumOutputs, HandleModulePart, PreserveLocals); return 0; }