Skip to content

Commit 473a18d

Browse files
author
Zak Kent
committed
[Immediate] Implement EagerSwiftMaterializationUnit 
1 parent 6abc5c5 commit 473a18d

File tree

3 files changed

+181
-45
lines changed

3 files changed

+181
-45
lines changed

‎include/swift/Immediate/SwiftMaterializationUnit.h

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,8 @@ class SwiftJIT {
7676

7777
llvm::orc::ObjectTransformLayer &getObjTransformLayer();
7878

79+
llvm::Expected<int> runMain(llvm::ArrayRef<std::string> Args);
80+
7981
private:
8082
static llvm::Expected<std::unique_ptr<llvm::orc::LLJIT>>
8183
CreateLLJIT(CompilerInstance &CI);
@@ -108,16 +110,16 @@ class SwiftJIT {
108110
std::unique_ptr<llvm::orc::IndirectStubsManager> ISM;
109111
};
110112

111-
class SwiftMaterializationUnit : public llvm::orc::MaterializationUnit {
113+
class LazySwiftMaterializationUnit : public llvm::orc::MaterializationUnit {
112114
public:
113-
static std::unique_ptr<SwiftMaterializationUnit> Create(SwiftJIT &JIT,
114-
CompilerInstance &CI);
115+
static std::unique_ptr<LazySwiftMaterializationUnit>
116+
Create(SwiftJIT &JIT, CompilerInstance &CI);
115117
llvm::StringRef getName() const override;
116118

117119
private:
118-
SwiftMaterializationUnit(SwiftJIT &JIT, CompilerInstance &CI,
119-
SymbolSourceMap Sources,
120-
llvm::orc::SymbolFlagsMap Symbols);
120+
LazySwiftMaterializationUnit(SwiftJIT &JIT, CompilerInstance &CI,
121+
SymbolSourceMap Sources,
122+
llvm::orc::SymbolFlagsMap Symbols);
121123
void materialize(
122124
std::unique_ptr<llvm::orc::MaterializationResponsibility> MR) override;
123125
void discard(const llvm::orc::JITDylib &JD,
@@ -126,6 +128,29 @@ class SwiftMaterializationUnit : public llvm::orc::MaterializationUnit {
126128
SwiftJIT &JIT;
127129
CompilerInstance &CI;
128130
};
131+
132+
class EagerSwiftMaterializationUnit : public llvm::orc::MaterializationUnit {
133+
public:
134+
EagerSwiftMaterializationUnit(SwiftJIT &JIT, const CompilerInstance &CI,
135+
const IRGenOptions &IRGenOpts,
136+
std::unique_ptr<SILModule> SM);
137+
138+
StringRef getName() const override;
139+
140+
private:
141+
void materialize(
142+
std::unique_ptr<llvm::orc::MaterializationResponsibility> MR) override;
143+
static MaterializationUnit::Interface
144+
getInterface(SwiftJIT &JIT, const CompilerInstance &CI);
145+
void dumpJIT(const llvm::Module &Module);
146+
void discard(const llvm::orc::JITDylib &JD,
147+
const llvm::orc::SymbolStringPtr &Sym) override;
148+
SwiftJIT &JIT;
149+
const CompilerInstance &CI;
150+
const IRGenOptions &IRGenOpts;
151+
std::unique_ptr<SILModule> SM;
152+
};
153+
129154
} // namespace swift
130155

131156
#endif // SWIFT_IMMEDIATE_SWIFTMATERIALIZATIONUNIT_H

‎lib/Immediate/Immediate.cpp

Lines changed: 21 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -532,39 +532,26 @@ int swift::RunImmediately(CompilerInstance &CI, const ProcessCmdLine &CmdLine,
532532
if (autolinkImportedModules(swiftModule, IRGenOpts))
533533
return -1;
534534

535-
auto &Target = swiftModule->getASTContext().LangOpts.Target;
536-
if (Target.isMacOSX()) {
537-
auto JIT = SwiftJIT::Create(CI);
538-
if (auto Err = JIT.takeError()) {
539-
logError(std::move(Err));
540-
return -1;
541-
}
542-
543-
auto MU = std::make_unique<SILMaterializationUnit>(**JIT, CI, IRGenOpts,
544-
std::move(SM));
545-
if (auto Err = (*JIT)->addSwift((*JIT)->getMainJITDylib(), std::move(MU))) {
546-
logError(std::move(Err));
547-
return -1;
548-
}
549-
550-
return runMain((*JIT)->getJIT(), CmdLine);
551-
}
552-
auto JIT = createLLJIT(IRGenOpts, swiftModule->getASTContext());
535+
auto JIT = SwiftJIT::Create(CI);
553536
if (auto Err = JIT.takeError()) {
554537
logError(std::move(Err));
555538
return -1;
556539
}
557-
auto GenModule = generateModule(CI, std::move(SM));
558-
if (!GenModule)
559-
return -1;
560-
auto *Module = GenModule->getModule();
561-
dumpJIT(**JIT, *Module, IRGenOpts);
562-
auto TSM = std::move(*GenModule).intoThreadSafeContext();
563-
if (auto Err = (*JIT)->addIRModule(std::move(TSM))) {
540+
541+
auto MU = std::make_unique<SILMaterializationUnit>(**JIT, CI, IRGenOpts,
542+
std::move(SM));
543+
if (auto Err = (*JIT)->addSwift((*JIT)->getMainJITDylib(), std::move(MU))) {
564544
logError(std::move(Err));
565545
return -1;
566546
}
567-
return runMain(**JIT, CmdLine);
547+
548+
auto Result = (*JIT)->runMain(CmdLine);
549+
550+
if (!Result) {
551+
logError(Result.takeError());
552+
return -1;
553+
}
554+
return *Result;
568555
}
569556

570557
int swift::RunImmediatelyFromAST(CompilerInstance &CI) {
@@ -642,11 +629,17 @@ int swift::RunImmediatelyFromAST(CompilerInstance &CI) {
642629
return -1;
643630
}
644631

645-
auto MU = SwiftMaterializationUnit::Create(**JIT, CI);
632+
auto MU = LazySwiftMaterializationUnit::Create(**JIT, CI);
646633
if (auto Err = (*JIT)->addSwift((*JIT)->getMainJITDylib(), std::move(MU))) {
647634
logError(std::move(Err));
648635
return -1;
649636
}
650637

651-
return runMain((*JIT)->getJIT(), CmdLine);
638+
auto Result = (*JIT)->runMain(CmdLine);
639+
if (!Result) {
640+
logError(Result.takeError());
641+
return -1;
642+
}
643+
644+
return *Result;
652645
}

‎lib/Immediate/SwiftMaterializationUnit.cpp

Lines changed: 129 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
#include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"
3333
#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
3434
#include "llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h"
35+
#include "llvm/ExecutionEngine/Orc/TargetProcess/TargetExecutionUtils.h"
3536
#include "llvm/Support/Error.h"
3637

3738
#include "swift/AST/IRGenRequests.h"
@@ -43,6 +44,8 @@
4344
#include "swift/SILOptimizer/PassManager/Passes.h"
4445
#include "swift/Subsystems.h"
4546

47+
#define DEBUG_TYPE "swift-immediate"
48+
4649
using namespace swift;
4750

4851
/// The suffix appended to function bodies when creating lazy reexports
@@ -92,6 +95,30 @@ SwiftJIT::~SwiftJIT() {
9295
J->getExecutionSession().reportError(std::move(Err));
9396
}
9497

98+
llvm::Expected<int> SwiftJIT::runMain(llvm::ArrayRef<std::string> Args) {
99+
if (auto Err = J->initialize(J->getMainJITDylib())) {
100+
return Err;
101+
}
102+
103+
auto MainSym = J->lookup("main");
104+
if (!MainSym) {
105+
return MainSym.takeError();
106+
}
107+
108+
using MainFnTy = int (*)(int, char *[]);
109+
MainFnTy JITMain = MainSym->toPtr<MainFnTy>();
110+
111+
LLVM_DEBUG(llvm::dbgs() << "Running main\n");
112+
int Result = llvm::orc::runAsMain(JITMain, Args);
113+
114+
LLVM_DEBUG(llvm::dbgs() << "Running static destructors\n");
115+
if (auto Err = J->deinitialize(J->getMainJITDylib())) {
116+
return Err;
117+
}
118+
119+
return Result;
120+
}
121+
95122
llvm::orc::LLJIT &SwiftJIT::getJIT() { return *J; }
96123

97124
llvm::orc::JITDylib &SwiftJIT::getMainJITDylib() {
@@ -277,8 +304,8 @@ static void logError(llvm::Error Err) {
277304
logAllUnhandledErrors(std::move(Err), llvm::errs(), "");
278305
}
279306

280-
std::unique_ptr<SwiftMaterializationUnit>
281-
SwiftMaterializationUnit::Create(SwiftJIT &JIT, CompilerInstance &CI) {
307+
std::unique_ptr<LazySwiftMaterializationUnit>
308+
LazySwiftMaterializationUnit::Create(SwiftJIT &JIT, CompilerInstance &CI) {
282309
auto *M = CI.getMainModule();
283310
TBDGenOptions Opts;
284311
Opts.PublicSymbolsOnly = false;
@@ -301,21 +328,22 @@ SwiftMaterializationUnit::Create(SwiftJIT &JIT, CompilerInstance &CI) {
301328
auto MangledName = mangle(SymbolName);
302329
PublicInterface[JIT.intern(MangledName)] = Flags;
303330
}
304-
return std::unique_ptr<SwiftMaterializationUnit>(new SwiftMaterializationUnit(
305-
JIT, CI, std::move(Sources), std::move(PublicInterface)));
331+
return std::unique_ptr<LazySwiftMaterializationUnit>(
332+
new LazySwiftMaterializationUnit(JIT, CI, std::move(Sources),
333+
std::move(PublicInterface)));
306334
}
307335

308-
StringRef SwiftMaterializationUnit::getName() const {
336+
StringRef LazySwiftMaterializationUnit::getName() const {
309337
return "SwiftMaterializationUnit";
310338
}
311339

312-
SwiftMaterializationUnit::SwiftMaterializationUnit(
340+
LazySwiftMaterializationUnit::LazySwiftMaterializationUnit(
313341
SwiftJIT &JIT, CompilerInstance &CI, SymbolSourceMap Sources,
314342
llvm::orc::SymbolFlagsMap Symbols)
315343
: MaterializationUnit({std::move(Symbols), nullptr}),
316344
Sources(std::move(Sources)), JIT(JIT), CI(CI) {}
317345

318-
void SwiftMaterializationUnit::materialize(
346+
void LazySwiftMaterializationUnit::materialize(
319347
std::unique_ptr<llvm::orc::MaterializationResponsibility> MR) {
320348
SILRefsToEmit Refs;
321349
const auto &RS = MR->getRequestedSymbols();
@@ -379,8 +407,8 @@ void SwiftMaterializationUnit::materialize(
379407
}
380408
}
381409
std::unique_ptr<MaterializationUnit> UnrequestedMU(
382-
new SwiftMaterializationUnit(JIT, CI, std::move(Sources),
383-
std::move(UnrequestedSymbols)));
410+
new LazySwiftMaterializationUnit(JIT, CI, std::move(Sources),
411+
std::move(UnrequestedSymbols)));
384412
if (auto Err = MR->replace(std::move(UnrequestedMU))) {
385413
logError(std::move(Err));
386414
MR->failMaterialization();
@@ -417,5 +445,95 @@ SwiftJIT::addSwift(llvm::orc::JITDylib &JD,
417445
return JD.define(std::move(MU));
418446
}
419447

420-
void SwiftMaterializationUnit::discard(const llvm::orc::JITDylib &JD,
421-
const llvm::orc::SymbolStringPtr &Sym) {}
448+
void LazySwiftMaterializationUnit::discard(
449+
const llvm::orc::JITDylib &JD, const llvm::orc::SymbolStringPtr &Sym) {}
450+
451+
EagerSwiftMaterializationUnit::EagerSwiftMaterializationUnit(
452+
SwiftJIT &JIT, const CompilerInstance &CI, const IRGenOptions &IRGenOpts,
453+
std::unique_ptr<SILModule> SM)
454+
: MaterializationUnit(getInterface(JIT, CI)), JIT(JIT), CI(CI),
455+
IRGenOpts(IRGenOpts), SM(std::move(SM)) {}
456+
457+
StringRef EagerSwiftMaterializationUnit::getName() const {
458+
return "EagerSwiftMaterializationUnit";
459+
}
460+
461+
void EagerSwiftMaterializationUnit::materialize(
462+
std::unique_ptr<llvm::orc::MaterializationResponsibility> R) {
463+
464+
auto GenModule = generateModule(CI, std::move(SM));
465+
466+
if (!GenModule) {
467+
R->failMaterialization();
468+
return;
469+
}
470+
471+
auto *Module = GenModule->getModule();
472+
473+
// Dump IR if requested
474+
dumpJIT(*Module);
475+
476+
// Now we must register all other public symbols defined by
477+
// the module with the JIT
478+
llvm::orc::SymbolFlagsMap Symbols;
479+
// Register all global values, including global
480+
// variables and functions
481+
for (const auto &GV : Module->global_values()) {
482+
// Ignore all symbols that will not appear in symbol table
483+
if (GV.hasLocalLinkage() || GV.isDeclaration() || GV.hasAppendingLinkage())
484+
continue;
485+
auto Name = GV.getName();
486+
// The entry point is already registered up front with the
487+
// interface, so ignore it as well
488+
if (Name == CI.getASTContext().getEntryPointFunctionName())
489+
continue;
490+
auto MangledName = JIT.mangleAndIntern(Name);
491+
// Register this symbol with the proper flags
492+
Symbols[MangledName] = llvm::JITSymbolFlags::fromGlobalValue(GV);
493+
}
494+
// Register the symbols we have discovered with the JIT
495+
if (auto Err = R->defineMaterializing(Symbols)) {
496+
logError(std::move(Err));
497+
}
498+
auto TSM = std::move(*GenModule).intoThreadSafeContext();
499+
JIT.getIRCompileLayer().emit(std::move(R), std::move(TSM));
500+
}
501+
502+
llvm::orc::MaterializationUnit::Interface
503+
EagerSwiftMaterializationUnit::getInterface(SwiftJIT &JIT,
504+
const CompilerInstance &CI) {
505+
const auto &EntryPoint = CI.getASTContext().getEntryPointFunctionName();
506+
auto MangledEntryPoint = JIT.mangleAndIntern(EntryPoint);
507+
auto Flags = llvm::JITSymbolFlags::Callable | llvm::JITSymbolFlags::Exported;
508+
llvm::orc::SymbolFlagsMap Symbols{{MangledEntryPoint, Flags}};
509+
return {std::move(Symbols), nullptr};
510+
}
511+
512+
void EagerSwiftMaterializationUnit::discard(
513+
const llvm::orc::JITDylib &JD, const llvm::orc::SymbolStringPtr &Sym) {}
514+
515+
static void DumpLLVMIR(const llvm::Module &M) {
516+
std::string path = (M.getName() + ".ll").str();
517+
for (size_t count = 0; llvm::sys::fs::exists(path);)
518+
path = (M.getName() + llvm::utostr(count++) + ".ll").str();
519+
520+
std::error_code error;
521+
llvm::raw_fd_ostream stream(path, error);
522+
if (error)
523+
return;
524+
M.print(stream, /*AssemblyAnnotationWriter=*/nullptr);
525+
}
526+
527+
void EagerSwiftMaterializationUnit::dumpJIT(const llvm::Module &M) {
528+
LLVM_DEBUG(llvm::dbgs() << "Module to be executed:\n"; M.dump());
529+
switch (IRGenOpts.DumpJIT) {
530+
case JITDebugArtifact::None:
531+
break;
532+
case JITDebugArtifact::LLVMIR:
533+
DumpLLVMIR(M);
534+
break;
535+
case JITDebugArtifact::Object:
536+
JIT.getObjTransformLayer().setTransform(llvm::orc::DumpObjects());
537+
break;
538+
}
539+
}

0 commit comments

Comments
 (0)