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/SmallString.h"
15#include "llvm/ADT/SmallVector.h"
16#include "llvm/ADT/StringExtras.h"
17#include "llvm/Bitcode/BitcodeWriter.h"
18#include "llvm/IR/LLVMContext.h"
19#include "llvm/IR/PassInstrumentation.h"
20#include "llvm/IR/PassManager.h"
21#include "llvm/IR/Verifier.h"
22#include "llvm/IRReader/IRReader.h"
23#include "llvm/MC/TargetRegistry.h"
24#include "llvm/Support/CommandLine.h"
25#include "llvm/Support/FileSystem.h"
26#include "llvm/Support/FormatVariadic.h"
27#include "llvm/Support/InitLLVM.h"
28#include "llvm/Support/SourceMgr.h"
29#include "llvm/Support/TargetSelect.h"
30#include "llvm/Support/ToolOutputFile.h"
31#include "llvm/Support/WithColor.h"
32#include "llvm/Support/raw_ostream.h"
33#include "llvm/Target/TargetMachine.h"
34#include "llvm/TargetParser/Triple.h"
35#include "llvm/Transforms/IPO/GlobalDCE.h"
36#include "llvm/Transforms/Utils/SplitModule.h"
37#include "llvm/Transforms/Utils/SplitModuleByCategory.h"
38
39using namespace llvm;
40
41static cl::OptionCategory SplitCategory("Split Options");
42
43static cl::opt<std::string> InputFilename(cl::Positional,
44 cl::desc("<input bitcode file>"),
45 cl::init(Val: "-"),
46 cl::value_desc("filename"),
47 cl::cat(SplitCategory));
48
49static cl::opt<std::string> OutputFilename("o",
50 cl::desc("Override output filename"),
51 cl::value_desc("filename"),
52 cl::cat(SplitCategory));
53
54static cl::opt<unsigned> NumOutputs("j", cl::Prefix, cl::init(Val: 2),
55 cl::desc("Number of output files"),
56 cl::cat(SplitCategory));
57
58static cl::opt<bool>
59 PreserveLocals("preserve-locals", cl::Prefix, cl::init(Val: false),
60 cl::desc("Split without externalizing locals"),
61 cl::cat(SplitCategory));
62
63static cl::opt<bool>
64 RoundRobin("round-robin", cl::Prefix, cl::init(Val: false),
65 cl::desc("Use round-robin distribution of functions to "
66 "modules instead of the default name-hash-based one"),
67 cl::cat(SplitCategory));
68
69static cl::opt<std::string>
70 MTriple("mtriple",
71 cl::desc("Target triple. When present, a TargetMachine is created "
72 "and TargetMachine::splitModule is used instead of the "
73 "common SplitModule logic."),
74 cl::value_desc("triple"), cl::cat(SplitCategory));
75
76static cl::opt<std::string>
77 MCPU("mcpu", cl::desc("Target CPU, ignored if --mtriple is not used"),
78 cl::value_desc("cpu"), cl::cat(SplitCategory));
79
80enum class SplitByCategoryType {
81 SBCT_ByAttribute,
82 SBCT_ByKernel,
83 SBCT_None,
84};
85
86static cl::opt<SplitByCategoryType> SplitByCategory(
87 "split-by-category",
88 cl::desc("Split by category. If present, splitting by category is used "
89 "with the specified categorization type."),
90 cl::Optional, cl::init(Val: SplitByCategoryType::SBCT_None),
91 cl::values(clEnumValN(SplitByCategoryType::SBCT_ByAttribute, "attribute",
92 "one output module per unique value of the function "
93 "attribute named by --category-attribute"),
94 clEnumValN(SplitByCategoryType::SBCT_ByKernel, "kernel",
95 "one output module per kernel")),
96 cl::cat(SplitCategory));
97
98static cl::opt<std::string>
99 CategoryAttribute("category-attribute",
100 cl::desc("Function attribute name to use when splitting "
101 "with -split-by-category=attribute"),
102 cl::value_desc("name"), cl::cat(SplitCategory));
103
104static cl::opt<bool> OutputAssembly{
105 "S", cl::desc("Write output as LLVM assembly"), cl::cat(SplitCategory)};
106
107void writeStringToFile(StringRef Content, StringRef Path) {
108 std::error_code EC;
109 raw_fd_ostream OS(Path, EC);
110 if (EC) {
111 errs() << formatv(Fmt: "error opening file: {0}, error: {1}\n", Vals&: Path,
112 Vals: EC.message());
113 exit(status: 1);
114 }
115
116 OS << Content << "\n";
117}
118
119void writeModuleToFile(Module &M, StringRef Path, bool OutputAssembly) {
120 int FD = -1;
121 if (std::error_code EC = sys::fs::openFileForWrite(Name: Path, ResultFD&: FD)) {
122 errs() << formatv(Fmt: "error opening file: {0}, error: {1}", Vals&: Path, Vals: EC.message())
123 << '\n';
124 exit(status: 1);
125 }
126
127 raw_fd_ostream OS(FD, /*ShouldClose*/ true);
128 if (OutputAssembly) {
129 M.renumberMetadataForAssembly();
130 M.print(OS, /*AssemblyAnnotationWriter*/ AAW: nullptr);
131 } else
132 WriteBitcodeToFile(M, Out&: OS);
133}
134
135/// EntryPointCategorizer is used for splitting by category either by a named
136/// function attribute or by kernels. It doesn't provide categories for
137/// functions other than kernels. Categorizer computes a string key for the
138/// given Function and records the association between the string key and an
139/// integer category. If a string key already belongs to some category then the
140/// corresponding integer category is returned.
141class EntryPointCategorizer {
142public:
143 EntryPointCategorizer(SplitByCategoryType Type, StringRef AttributeName)
144 : Type(Type), AttributeName(AttributeName) {}
145
146 EntryPointCategorizer() = delete;
147 EntryPointCategorizer(EntryPointCategorizer &) = delete;
148 EntryPointCategorizer &operator=(const EntryPointCategorizer &) = delete;
149 EntryPointCategorizer(EntryPointCategorizer &&) = default;
150 EntryPointCategorizer &operator=(EntryPointCategorizer &&) = default;
151
152 /// Returns integer specifying the category for the given \p F.
153 /// If the given function isn't a kernel then returns std::nullopt.
154 std::optional<int> operator()(const Function &F) {
155 if (!isEntryPoint(F))
156 return std::nullopt; // skip the function.
157
158 auto StringKey = computeFunctionCategory(Type, F);
159 if (auto it = StrKeyToID.find(Val: StringRef(StringKey)); it != StrKeyToID.end())
160 return it->second;
161
162 int ID = static_cast<int>(StrKeyToID.size());
163 return StrKeyToID.try_emplace(Key: std::move(StringKey), Args&: ID).first->second;
164 }
165
166private:
167 static bool isEntryPoint(const Function &F) {
168 if (F.isDeclaration())
169 return false;
170
171 return F.hasKernelCallingConv();
172 }
173
174 SmallString<0> computeFunctionCategory(SplitByCategoryType Type,
175 const Function &F) {
176 SmallString<0> Key;
177 switch (Type) {
178 case SplitByCategoryType::SBCT_ByKernel:
179 Key = F.getName().str();
180 break;
181 case SplitByCategoryType::SBCT_ByAttribute:
182 Key = F.getFnAttribute(Kind: AttributeName).getValueAsString().str();
183 break;
184 default:
185 llvm_unreachable("unexpected mode.");
186 }
187
188 return Key;
189 }
190
191private:
192 struct KeyInfo {
193 static bool isEqual(const SmallString<0> &LHS, const SmallString<0> &RHS) {
194 return LHS == RHS;
195 }
196
197 static unsigned getHashValue(const SmallString<0> &S) {
198 return llvm::hash_value(S: StringRef(S));
199 }
200 };
201
202 SplitByCategoryType Type;
203 std::string AttributeName;
204 DenseMap<SmallString<0>, int, KeyInfo> StrKeyToID;
205};
206
207void cleanupModule(Module &M) {
208 ModuleAnalysisManager MAM;
209 MAM.registerPass(PassBuilder: [&] { return PassInstrumentationAnalysis(); });
210 ModulePassManager MPM;
211 MPM.addPass(Pass: GlobalDCEPass()); // Delete unreachable globals.
212 MPM.run(IR&: M, AM&: MAM);
213}
214
215Error runSplitModuleByCategory(std::unique_ptr<Module> M) {
216 if (SplitByCategory == SplitByCategoryType::SBCT_ByAttribute &&
217 CategoryAttribute.empty())
218 return createStringError(
219 Fmt: "-split-by-category=attribute requires --category-attribute=<name>");
220
221 size_t OutputID = 0;
222 auto PostSplitCallback = [&](std::unique_ptr<Module> MPart) -> Error {
223 if (verifyModule(M: *MPart)) {
224 errs() << "Broken Module!\n";
225 exit(status: 1);
226 }
227
228 // TODO: DCE is a crucial pass since it removes unused declarations.
229 // At the moment, LIT checking can't be perfomed without DCE.
230 cleanupModule(M&: *MPart);
231 size_t ID = OutputID;
232 ++OutputID;
233 StringRef ModuleSuffix = OutputAssembly ? ".ll" : ".bc";
234 std::string ModulePath =
235 (Twine(OutputFilename) + "_" + Twine(ID) + ModuleSuffix).str();
236 writeModuleToFile(M&: *MPart, Path: ModulePath, OutputAssembly);
237 return Error::success();
238 };
239
240 auto Categorizer = EntryPointCategorizer(SplitByCategory, CategoryAttribute);
241 return splitModuleTransitiveFromEntryPoints(M: std::move(M), EntryPointCategorizer: Categorizer,
242 Callback: PostSplitCallback);
243}
244
245int main(int argc, char **argv) {
246 InitLLVM X(argc, argv);
247
248 LLVMContext Context;
249 SMDiagnostic Err;
250 cl::HideUnrelatedOptions(Categories: {&SplitCategory, &getColorCategory()});
251 cl::ParseCommandLineOptions(argc, argv, Overview: "LLVM module splitter\n");
252
253 Triple TT(MTriple);
254
255 std::unique_ptr<TargetMachine> TM;
256 if (!MTriple.empty()) {
257 InitializeAllTargets();
258 InitializeAllTargetMCs();
259
260 std::string Error;
261 const Target *T = TargetRegistry::lookupTarget(TheTriple: TT, Error);
262 if (!T) {
263 errs() << "unknown target '" << MTriple << "': " << Error << "\n";
264 return 1;
265 }
266
267 TargetOptions Options;
268 TM = std::unique_ptr<TargetMachine>(T->createTargetMachine(
269 TT, CPU: MCPU, /*FS*/ Features: "", Options, RM: std::nullopt, CM: std::nullopt));
270 }
271
272 std::unique_ptr<Module> M = parseIRFile(Filename: InputFilename, Err, Context);
273
274 if (!M) {
275 Err.print(ProgName: argv[0], S&: errs());
276 return 1;
277 }
278
279 unsigned I = 0;
280 const auto HandleModulePart = [&](std::unique_ptr<Module> MPart) {
281 std::error_code EC;
282 std::unique_ptr<ToolOutputFile> Out(
283 new ToolOutputFile(OutputFilename + utostr(X: I++), EC, sys::fs::OF_None));
284 if (EC) {
285 errs() << EC.message() << '\n';
286 exit(status: 1);
287 }
288
289 if (verifyModule(M: *MPart, OS: &errs())) {
290 errs() << "Broken module!\n";
291 exit(status: 1);
292 }
293
294 WriteBitcodeToFile(M: *MPart, Out&: Out->os());
295
296 // Declare success.
297 Out->keep();
298 };
299
300 if (SplitByCategory != SplitByCategoryType::SBCT_None) {
301 auto E = runSplitModuleByCategory(M: std::move(M));
302 if (E) {
303 errs() << "error: " << toString(E: std::move(E)) << "\n";
304 return 1;
305 }
306
307 return 0;
308 }
309
310 if (TM) {
311 if (PreserveLocals) {
312 errs() << "warning: --preserve-locals has no effect when using "
313 "TargetMachine::splitModule\n";
314 }
315 if (RoundRobin)
316 errs() << "warning: --round-robin has no effect when using "
317 "TargetMachine::splitModule\n";
318
319 if (TM->splitModule(M&: *M, NumParts: NumOutputs, ModuleCallback: HandleModulePart))
320 return 0;
321
322 errs() << "warning: "
323 "TargetMachine::splitModule failed, falling back to default "
324 "splitModule implementation\n";
325 }
326
327 SplitModule(M&: *M, N: NumOutputs, ModuleCallback: HandleModulePart, PreserveLocals, RoundRobin);
328 return 0;
329}
330