[LTO] Allow target-specific module splittting (#83128)

Allow targets to implement custom module splitting logic for
--lto-partitions, see #89245

https://discourse.llvm.org/t/rfc-lto-target-specific-module-splittting/77252

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<void(std::unique_ptr<Module> MPart)> ModuleCallback) const {
+    return false;
+  }
 };
 
 /// This class describes a target machine that is implemented with the LLVM

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,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

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
+}

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

llvm/tools/llvm-split/llvm-split.cpp

@@ -1,12 +1,13 @@
-//===-- 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.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// 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<bool>
                    cl::desc("Split without externalizing locals"),
                    cl::cat(SplitCategory));
 
+static cl::opt<std::string>
+    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<std::string>
+    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<Module> 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<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) {
+    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;
 }