Skip to content

Commit a12a6f1

Browse files
author
Zak Kent
committed
[Immediate] Refactor Immediate Mode to use Materialization Unit
Refactor Immediate Mode to use the Materialization Unit API, materializing the entire `SILModule` corresponding to the script when the `main` symbol is looked up.
1 parent b03d8df commit a12a6f1

File tree

1 file changed

+121
-90
lines changed

1 file changed

+121
-90
lines changed

‎lib/Immediate/Immediate.cpp

Lines changed: 121 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -254,40 +254,116 @@ bool swift::immediate::autolinkImportedModules(ModuleDecl *M,
254254
return false;
255255
}
256256

257-
int swift::RunImmediately(CompilerInstance &CI,
258-
const ProcessCmdLine &CmdLine,
259-
const IRGenOptions &IRGenOpts,
260-
const SILOptions &SILOpts,
261-
std::unique_ptr<SILModule> &&SM) {
262-
// TODO: Use OptimizedIRRequest for this.
263-
ASTContext &Context = CI.getASTContext();
264-
265-
// IRGen the main module.
266-
auto *swiftModule = CI.getMainModule();
267-
const auto PSPs = CI.getPrimarySpecificPathsForAtMostOnePrimary();
268-
const auto &TBDOpts = CI.getInvocation().getTBDGenOptions();
269-
auto GenModule = performIRGeneration(
270-
swiftModule, IRGenOpts, TBDOpts, std::move(SM),
271-
swiftModule->getName().str(), PSPs, ArrayRef<std::string>());
257+
static llvm::Expected<std::unique_ptr<llvm::orc::LLJIT>>
258+
createJIT(const IRGenOptions &IRGenOpts, ASTContext &Ctx) {
259+
llvm::TargetOptions TargetOpt;
260+
std::string CPU;
261+
std::string Triple;
262+
std::vector<std::string> Features;
263+
std::tie(TargetOpt, CPU, Features, Triple) =
264+
getIRTargetOptions(IRGenOpts, Ctx);
265+
auto JTMB = llvm::orc::JITTargetMachineBuilder(llvm::Triple(Triple))
266+
.setRelocationModel(llvm::Reloc::PIC_)
267+
.setOptions(std::move(TargetOpt))
268+
.setCPU(std::move(CPU))
269+
.addFeatures(Features)
270+
.setCodeGenOptLevel(llvm::CodeGenOpt::Default);
271+
return llvm::orc::LLJITBuilder().setJITTargetMachineBuilder(JTMB).create();
272+
}
272273

273-
if (Context.hadError())
274-
return -1;
274+
class SILMaterializationUnit : public llvm::orc::MaterializationUnit {
275+
public:
276+
SILMaterializationUnit(llvm::orc::LLJIT &JIT, CompilerInstance &CI,
277+
const IRGenOptions &IRGenOpts,
278+
std::unique_ptr<SILModule> SM)
279+
: MaterializationUnit(getInterface(JIT, SM->getSwiftModule())), JIT(JIT),
280+
CI(CI), IRGenOpts(IRGenOpts), SM(std::move(SM)) {}
281+
282+
void materialize(
283+
std::unique_ptr<llvm::orc::MaterializationResponsibility> R) override {
284+
// TODO: Use OptimizedIRRequest for this.
285+
ASTContext &Context = CI.getASTContext();
286+
auto *swiftModule = CI.getMainModule();
287+
const auto PSPs = CI.getPrimarySpecificPathsForAtMostOnePrimary();
288+
const auto &TBDOpts = CI.getInvocation().getTBDGenOptions();
289+
auto GenModule = performIRGeneration(
290+
swiftModule, IRGenOpts, TBDOpts, std::move(SM),
291+
swiftModule->getName().str(), PSPs, ArrayRef<std::string>());
292+
if (Context.hadError()) {
293+
R->failMaterialization();
294+
return;
295+
}
296+
assert(GenModule && "Emitted no diagnostics but IR generation failed?");
297+
auto *Module = GenModule.getModule();
298+
performLLVM(IRGenOpts, Context.Diags, /*diagMutex*/ nullptr,
299+
/*hash*/ nullptr, Module, GenModule.getTargetMachine(),
300+
CI.getPrimarySpecificPathsForAtMostOnePrimary().OutputFilename,
301+
CI.getOutputBackend(), Context.Stats);
302+
303+
switch (IRGenOpts.DumpJIT) {
304+
case JITDebugArtifact::None:
305+
break;
306+
case JITDebugArtifact::LLVMIR:
307+
DumpLLVMIR(*Module);
308+
break;
309+
case JITDebugArtifact::Object:
310+
JIT.getObjTransformLayer().setTransform(llvm::orc::DumpObjects());
311+
break;
312+
}
275313

276-
assert(GenModule && "Emitted no diagnostics but IR generation failed?");
314+
LLVM_DEBUG(llvm::dbgs() << "Module to be executed:\n"; Module->dump());
315+
316+
auto globals = Module->global_objects();
317+
llvm::orc::SymbolFlagsMap Symbols;
318+
for (auto itr = globals.begin(); itr != globals.end(); itr++) {
319+
auto &global = *itr;
320+
if (global.hasLocalLinkage() || global.isDeclaration() ||
321+
global.hasAppendingLinkage())
322+
continue;
323+
auto name = JIT.mangleAndIntern(global.getName());
324+
if (!R->getSymbols().count(name)) {
325+
Symbols[name] = llvm::JITSymbolFlags::fromGlobalValue(global);
326+
}
327+
}
328+
if (auto Err = R->defineMaterializing(Symbols)) {
329+
logAllUnhandledErrors(std::move(Err), llvm::errs(), "");
330+
}
331+
auto TSM = std::move(GenModule).intoThreadSafeContext();
332+
JIT.getIRCompileLayer().emit(std::move(R), std::move(TSM));
333+
}
277334

278-
performLLVM(IRGenOpts, Context.Diags, /*diagMutex*/ nullptr, /*hash*/ nullptr,
279-
GenModule.getModule(), GenModule.getTargetMachine(),
280-
PSPs.OutputFilename, CI.getOutputBackend(),
281-
Context.Stats);
335+
StringRef getName() const override { return "SwiftMU"; }
282336

283-
if (Context.hadError())
284-
return -1;
337+
private:
338+
void discard(const llvm::orc::JITDylib &JD,
339+
const llvm::orc::SymbolStringPtr &Sym) override {}
285340

341+
static MaterializationUnit::Interface getInterface(llvm::orc::LLJIT &JIT,
342+
ModuleDecl *SwiftModule) {
343+
// FIXME: Register symbols with the correct flags
344+
llvm::orc::SymbolFlagsMap Symbols{
345+
{JIT.mangleAndIntern("main"),
346+
llvm::JITSymbolFlags::Callable | llvm::JITSymbolFlags::Exported}};
347+
return {std::move(Symbols), nullptr};
348+
}
349+
std::unique_ptr<llvm::TargetMachine> Target;
350+
llvm::orc::LLJIT &JIT;
351+
CompilerInstance &CI;
352+
const IRGenOptions &IRGenOpts;
353+
std::unique_ptr<SILModule> SM;
354+
};
355+
356+
int swift::RunImmediately(CompilerInstance &CI, const ProcessCmdLine &CmdLine,
357+
const IRGenOptions &IRGenOpts,
358+
const SILOptions &SILOpts,
359+
std::unique_ptr<SILModule> &&SM) {
360+
ASTContext &Context = CI.getASTContext();
286361
// Load libSwiftCore to setup process arguments.
287362
//
288363
// This must be done here, before any library loading has been done, to avoid
289364
// racing with the static initializers in user code.
290365
// Setup interpreted process arguments.
366+
291367
using ArgOverride = void (* SWIFT_CC(swift))(const char **, int);
292368
#if defined(_WIN32)
293369
auto stdlib = loadSwiftRuntime(Context.SearchPathOpts.RuntimeLibraryPaths);
@@ -325,6 +401,8 @@ int swift::RunImmediately(CompilerInstance &CI,
325401
}
326402
#endif
327403

404+
auto *swiftModule = CI.getMainModule();
405+
328406
SmallVector<const char *, 32> argBuf;
329407
for (size_t i = 0; i < CmdLine.size(); ++i) {
330408
argBuf.push_back(CmdLine[i].c_str());
@@ -336,83 +414,36 @@ int swift::RunImmediately(CompilerInstance &CI,
336414
if (autolinkImportedModules(swiftModule, IRGenOpts))
337415
return -1;
338416

339-
llvm::PassManagerBuilder PMBuilder;
340-
PMBuilder.OptLevel = 2;
341-
PMBuilder.Inliner = llvm::createFunctionInliningPass(200);
342-
343-
// Build the ExecutionEngine.
344-
llvm::TargetOptions TargetOpt;
345-
std::string CPU;
346-
std::string Triple;
347-
std::vector<std::string> Features;
348-
std::tie(TargetOpt, CPU, Features, Triple)
349-
= getIRTargetOptions(IRGenOpts, swiftModule->getASTContext());
350-
351-
std::unique_ptr<llvm::orc::LLJIT> JIT;
352-
353-
{
354-
auto JITOrErr =
355-
llvm::orc::LLJITBuilder()
356-
.setJITTargetMachineBuilder(
357-
llvm::orc::JITTargetMachineBuilder(llvm::Triple(Triple))
358-
.setRelocationModel(llvm::Reloc::PIC_)
359-
.setOptions(std::move(TargetOpt))
360-
.setCPU(std::move(CPU))
361-
.addFeatures(Features)
362-
.setCodeGenOptLevel(llvm::CodeGenOpt::Default))
363-
.create();
364-
365-
if (!JITOrErr) {
366-
llvm::logAllUnhandledErrors(JITOrErr.takeError(), llvm::errs(), "");
367-
return -1;
368-
} else
369-
JIT = std::move(*JITOrErr);
370-
}
371-
372-
auto Module = GenModule.getModule();
373-
374-
switch (IRGenOpts.DumpJIT) {
375-
case JITDebugArtifact::None:
376-
break;
377-
case JITDebugArtifact::LLVMIR:
378-
DumpLLVMIR(*Module);
379-
break;
380-
case JITDebugArtifact::Object:
381-
JIT->getObjTransformLayer().setTransform(llvm::orc::DumpObjects());
382-
break;
417+
auto JIT = createJIT(IRGenOpts, swiftModule->getASTContext());
418+
if (auto Err = JIT.takeError()) {
419+
llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(), "");
420+
return -1;
383421
}
384-
385-
{
386-
// Get a generator for the process symbols and attach it to the main
387-
// JITDylib.
388-
if (auto G = llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(Module->getDataLayout().getGlobalPrefix()))
389-
JIT->getMainJITDylib().addGenerator(std::move(*G));
390-
else {
391-
logAllUnhandledErrors(G.takeError(), llvm::errs(), "");
392-
return -1;
393-
}
422+
if (auto G = llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
423+
(*JIT)->getDataLayout().getGlobalPrefix()))
424+
(*JIT)->getMainJITDylib().addGenerator(std::move(*G));
425+
else {
426+
logAllUnhandledErrors(G.takeError(), llvm::errs(), "");
427+
return -1;
394428
}
395-
396-
LLVM_DEBUG(llvm::dbgs() << "Module to be executed:\n";
397-
Module->dump());
398-
399-
{
400-
if (auto Err = JIT->addIRModule(std::move(GenModule).intoThreadSafeContext())) {
401-
llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(), "");
402-
return -1;
403-
}
429+
auto MU = std::make_unique<SILMaterializationUnit>(**JIT, CI, IRGenOpts,
430+
std::move(SM));
431+
if (auto Err = (*JIT)->getMainJITDylib().define(std::move(MU))) {
432+
llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(), "");
433+
return -1;
404434
}
405435

406436
using MainFnTy = int(*)(int, char*[]);
407437

408438
LLVM_DEBUG(llvm::dbgs() << "Running static constructors\n");
409-
if (auto Err = JIT->initialize(JIT->getMainJITDylib())) {
439+
if (auto Err = (*JIT)->initialize((*JIT)->getMainJITDylib())) {
410440
llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(), "");
411441
return -1;
412442
}
413443

414444
MainFnTy JITMain = nullptr;
415-
if (auto MainFnOrErr = JIT->lookup("main"))
445+
auto mangled = (*JIT)->mangleAndIntern("main");
446+
if (auto MainFnOrErr = (*JIT)->lookupLinkerMangled(*mangled))
416447
JITMain = llvm::jitTargetAddressToFunction<MainFnTy>(MainFnOrErr->getValue());
417448
else {
418449
logAllUnhandledErrors(MainFnOrErr.takeError(), llvm::errs(), "");
@@ -423,7 +454,7 @@ int swift::RunImmediately(CompilerInstance &CI,
423454
int Result = llvm::orc::runAsMain(JITMain, CmdLine);
424455

425456
LLVM_DEBUG(llvm::dbgs() << "Running static destructors\n");
426-
if (auto Err = JIT->deinitialize(JIT->getMainJITDylib())) {
457+
if (auto Err = (*JIT)->deinitialize((*JIT)->getMainJITDylib())) {
427458
logAllUnhandledErrors(std::move(Err), llvm::errs(), "");
428459
return -1;
429460
}

0 commit comments

Comments
 (0)