1//===-- llvm-split: command line tool for testing module splitting --------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This program can be used to test the llvm::SplitModule and
10// TargetMachine::splitModule functions.
11//
12//===----------------------------------------------------------------------===//
13
14#include "llvm/ADT/StringExtras.h"
15#include "llvm/Bitcode/BitcodeWriter.h"
16#include "llvm/IR/LLVMContext.h"
17#include "llvm/IR/Verifier.h"
18#include "llvm/IRReader/IRReader.h"
19#include "llvm/MC/TargetRegistry.h"
20#include "llvm/Support/CommandLine.h"
21#include "llvm/Support/FileSystem.h"
22#include "llvm/Support/InitLLVM.h"
23#include "llvm/Support/SourceMgr.h"
24#include "llvm/Support/TargetSelect.h"
25#include "llvm/Support/ToolOutputFile.h"
26#include "llvm/Support/WithColor.h"
27#include "llvm/Support/raw_ostream.h"
28#include "llvm/Target/TargetMachine.h"
29#include "llvm/TargetParser/Triple.h"
30#include "llvm/Transforms/Utils/SplitModule.h"
31
32using namespace llvm;
33
34static cl::OptionCategory SplitCategory("Split Options");
35
36static cl::opt<std::string> InputFilename(cl::Positional,
37 cl::desc("<input bitcode file>"),
38 cl::init(Val: "-"),
39 cl::value_desc("filename"),
40 cl::cat(SplitCategory));
41
42static cl::opt<std::string> OutputFilename("o",
43 cl::desc("Override output filename"),
44 cl::value_desc("filename"),
45 cl::cat(SplitCategory));
46
47static cl::opt<unsigned> NumOutputs("j", cl::Prefix, cl::init(Val: 2),
48 cl::desc("Number of output files"),
49 cl::cat(SplitCategory));
50
51static cl::opt<bool>
52 PreserveLocals("preserve-locals", cl::Prefix, cl::init(Val: false),
53 cl::desc("Split without externalizing locals"),
54 cl::cat(SplitCategory));
55
56static cl::opt<bool>
57 RoundRobin("round-robin", cl::Prefix, cl::init(Val: false),
58 cl::desc("Use round-robin distribution of functions to "
59 "modules instead of the default name-hash-based one"),
60 cl::cat(SplitCategory));
61
62static cl::opt<std::string>
63 MTriple("mtriple",
64 cl::desc("Target triple. When present, a TargetMachine is created "
65 "and TargetMachine::splitModule is used instead of the "
66 "common SplitModule logic."),
67 cl::value_desc("triple"), cl::cat(SplitCategory));
68
69static cl::opt<std::string>
70 MCPU("mcpu", cl::desc("Target CPU, ignored if --mtriple is not used"),
71 cl::value_desc("cpu"), cl::cat(SplitCategory));
72
73int main(int argc, char **argv) {
74 InitLLVM X(argc, argv);
75
76 LLVMContext Context;
77 SMDiagnostic Err;
78 cl::HideUnrelatedOptions(Categories: {&SplitCategory, &getColorCategory()});
79 cl::ParseCommandLineOptions(argc, argv, Overview: "LLVM module splitter\n");
80
81 std::unique_ptr<TargetMachine> TM;
82 if (!MTriple.empty()) {
83 InitializeAllTargets();
84 InitializeAllTargetMCs();
85
86 std::string Error;
87 const Target *T = TargetRegistry::lookupTarget(TripleStr: MTriple, Error);
88 if (!T) {
89 errs() << "unknown target '" << MTriple << "': " << Error << "\n";
90 return 1;
91 }
92
93 TargetOptions Options;
94 TM = std::unique_ptr<TargetMachine>(T->createTargetMachine(
95 TT: Triple(MTriple), CPU: MCPU, /*FS*/ Features: "", Options, RM: std::nullopt, CM: std::nullopt));
96 }
97
98 std::unique_ptr<Module> M = parseIRFile(Filename: InputFilename, Err, Context);
99
100 if (!M) {
101 Err.print(ProgName: argv[0], S&: errs());
102 return 1;
103 }
104
105 unsigned I = 0;
106 const auto HandleModulePart = [&](std::unique_ptr<Module> MPart) {
107 std::error_code EC;
108 std::unique_ptr<ToolOutputFile> Out(
109 new ToolOutputFile(OutputFilename + utostr(X: I++), EC, sys::fs::OF_None));
110 if (EC) {
111 errs() << EC.message() << '\n';
112 exit(status: 1);
113 }
114
115 if (verifyModule(M: *MPart, OS: &errs())) {
116 errs() << "Broken module!\n";
117 exit(status: 1);
118 }
119
120 WriteBitcodeToFile(M: *MPart, Out&: Out->os());
121
122 // Declare success.
123 Out->keep();
124 };
125
126 if (TM) {
127 if (PreserveLocals) {
128 errs() << "warning: --preserve-locals has no effect when using "
129 "TargetMachine::splitModule\n";
130 }
131 if (RoundRobin)
132 errs() << "warning: --round-robin has no effect when using "
133 "TargetMachine::splitModule\n";
134
135 if (TM->splitModule(M&: *M, NumParts: NumOutputs, ModuleCallback: HandleModulePart))
136 return 0;
137
138 errs() << "warning: "
139 "TargetMachine::splitModule failed, falling back to default "
140 "splitModule implementation\n";
141 }
142
143 SplitModule(M&: *M, N: NumOutputs, ModuleCallback: HandleModulePart, PreserveLocals, RoundRobin);
144 return 0;
145}
146