MLIR: lib/ExecutionEngine/ExecutionEngine.cpp Source File (original) (raw)

1

2

3

4

5

6

7

8

9

10

11

12

18

19#include "llvm/ExecutionEngine/JITEventListener.h"

20#include "llvm/ExecutionEngine/ObjectCache.h"

21#include "llvm/ExecutionEngine/Orc/CompileUtils.h"

22#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"

23#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"

24#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"

25#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"

26#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"

27#include "llvm/IR/IRBuilder.h"

28#include "llvm/MC/TargetRegistry.h"

29#include "llvm/Support/Debug.h"

30#include "llvm/Support/Error.h"

31#include "llvm/Support/ToolOutputFile.h"

32#include "llvm/TargetParser/Host.h"

33#include "llvm/TargetParser/SubtargetFeature.h"

34

35#define DEBUG_TYPE "execution-engine"

36

37using namespace mlir;

38using llvm::dbgs;

39using llvm::Error;

40using llvm::errs;

42using llvm::LLVMContext;

43using llvm::MemoryBuffer;

44using llvm::MemoryBufferRef;

45using llvm::Module;

46using llvm::SectionMemoryManager;

47using llvm::StringError;

48using llvm::Triple;

49using llvm::orc::DynamicLibrarySearchGenerator;

50using llvm::orc::ExecutionSession;

51using llvm::orc::IRCompileLayer;

52using llvm::orc::JITTargetMachineBuilder;

53using llvm::orc::MangleAndInterner;

54using llvm::orc::RTDyldObjectLinkingLayer;

55using llvm::orc::SymbolMap;

56using llvm::orc::ThreadSafeModule;

57using llvm::orc::TMOwningSimpleCompiler;

58

59

61 return llvm::make_error(message.str(),

62 llvm::inconvertibleErrorCode());

63}

64

66 MemoryBufferRef objBuffer) {

67 cachedObjects[m->getModuleIdentifier()] = MemoryBuffer::getMemBufferCopy(

68 objBuffer.getBuffer(), objBuffer.getBufferIdentifier());

69}

70

72 auto i = cachedObjects.find(m->getModuleIdentifier());

73 if (i == cachedObjects.end()) {

74 LLVM_DEBUG(dbgs() << "No object for " << m->getModuleIdentifier()

75 << " in cache. Compiling.\n");

76 return nullptr;

77 }

78 LLVM_DEBUG(dbgs() << "Object for " << m->getModuleIdentifier()

79 << " loaded from cache.\n");

80 return MemoryBuffer::getMemBuffer(i->second->getMemBufferRef());

81}

82

84

85 std::string errorMessage;

86 auto file = openOutputFile(outputFilename, &errorMessage);

87 if (!file) {

88 llvm::errs() << errorMessage << "\n";

89 return;

90 }

91

92

93 assert(cachedObjects.size() == 1 && "Expected only one object entry.");

94 auto &cachedObject = cachedObjects.begin()->second;

95 file->os() << cachedObject->getBuffer();

96 file->keep();

97}

98

100

102 if (cache == nullptr) {

103 llvm::errs() << "cannot dump ExecutionEngine object code to file: "

104 "object cache is disabled\n";

105 return;

106 }

107

108

109

110 if (cache->isEmpty()) {

111 for (std::string &functionName : functionNames) {

114 llvm::errs() << "Could not compile " << functionName << ":\n "

115 << result.takeError() << "\n";

116 return;

117 }

118 }

119 }

120 cache->dumpToObjectFile(filename);

121}

122

125 auto &mainJitDylib = jit->getMainJITDylib();

126 cantFail(mainJitDylib.define(

127 absoluteSymbols(symbolMap(llvm::orc::MangleAndInterner(

128 mainJitDylib.getExecutionSession(), jit->getDataLayout())))));

129}

130

132 llvm::TargetMachine *tm) {

133 llvmModule->setDataLayout(tm->createDataLayout());

134 llvmModule->setTargetTriple(tm->getTargetTriple());

135}

136

138 return "_mlir_" + name.str();

139}

140

141

142

143

145 auto &ctx = module->getContext();

146 llvm::IRBuilder<> builder(ctx);

148 for (auto &func : module->getFunctionList()) {

149 if (func.isDeclaration() || func.hasLocalLinkage())

150 continue;

151 if (interfaceFunctions.count(&func))

152 continue;

153

154

155

156 auto *newType =

157 llvm::FunctionType::get(builder.getVoidTy(), builder.getPtrTy(),

158 false);

160 auto funcCst = module->getOrInsertFunction(newName, newType);

161 llvm::Function *interfaceFunc = castllvm::Function(funcCst.getCallee());

162 interfaceFunctions.insert(interfaceFunc);

163

164

165

166 auto *bb = llvm::BasicBlock::Create(ctx);

167 bb->insertInto(interfaceFunc);

168 builder.SetInsertPoint(bb);

169 llvm::Value *argList = interfaceFunc->arg_begin();

171 args.reserve(llvm::size(func.args()));

172 for (auto [index, arg] : llvm::enumerate(func.args())) {

173 llvm::Value *argIndex = llvm::Constant::getIntegerValue(

174 builder.getInt64Ty(), APInt(64, index));

175 llvm::Value *argPtrPtr =

176 builder.CreateGEP(builder.getPtrTy(), argList, argIndex);

177 llvm::Value *argPtr = builder.CreateLoad(builder.getPtrTy(), argPtrPtr);

178 llvm::Type *argTy = arg.getType();

179 llvm::Value *load = builder.CreateLoad(argTy, argPtr);

180 args.push_back(load);

181 }

182

183

184 llvm::Value *result = builder.CreateCall(&func, args);

185

186

187 if (result->getType()->isVoidTy()) {

188 llvm::Value *retIndex = llvm::Constant::getIntegerValue(

189 builder.getInt64Ty(), APInt(64, llvm::size(func.args())));

190 llvm::Value *retPtrPtr =

191 builder.CreateGEP(builder.getPtrTy(), argList, retIndex);

192 llvm::Value *retPtr = builder.CreateLoad(builder.getPtrTy(), retPtrPtr);

193 builder.CreateStore(result, retPtr);

194 }

195

196

197 builder.CreateRetVoid();

198 }

199}

200

202 bool enableGDBNotificationListener,

203 bool enablePerfNotificationListener)

205 functionNames(),

206 gdbListener(enableGDBNotificationListener

207 ? llvm::JITEventListener::createGDBRegistrationListener()

210 if (enablePerfNotificationListener) {

211 if (auto *listener = llvm::JITEventListener::createPerfJITEventListener())

212 perfListener = listener;

213 else if (auto *listener =

214 llvm::JITEventListener::createIntelJITEventListener())

215 perfListener = listener;

216 }

217}

218

220

221

222

223 if (jit && !jit->getTargetTriple().isAArch64())

224 llvm::consumeError(jit->deinitialize(jit->getMainJITDylib()));

225

227 destroy();

228}

229

232 std::unique_ptrllvm::TargetMachine tm) {

233 auto engine = std::make_unique(

234 options.enableObjectDump, options.enableGDBNotificationListener,

235 options.enablePerfNotificationListener);

236

237

238 if (options.enableObjectDump) {

239 for (auto funcOp : m->getRegion(0).getOpsLLVM::LLVMFuncOp()) {

240 if (funcOp.getBlocks().empty())

241 continue;

242 StringRef funcName = funcOp.getSymName();

243 engine->functionNames.push_back(funcName.str());

244 }

245 }

246

247 std::unique_ptrllvm::LLVMContext ctx(new llvm::LLVMContext);

248 auto llvmModule = options.llvmModuleBuilder

249 ? options.llvmModuleBuilder(m, *ctx)

251 if (!llvmModule)

253

254

255

256 if (!tm) {

257 auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();

258 if (!tmBuilderOrError)

259 return tmBuilderOrError.takeError();

260

261 auto tmOrError = tmBuilderOrError->createTargetMachine();

262 if (!tmOrError)

263 return tmOrError.takeError();

264 tm = std::move(tmOrError.get());

265 }

266

267

268

269

270

271

274

275 auto dataLayout = llvmModule->getDataLayout();

276

277

280 options.sharedLibPaths, std::back_inserter(sharedLibPaths),

281 [](StringRef libPath) {

282 SmallString<256> absPath(libPath.begin(), libPath.end());

283 cantFail(llvm::errorCodeToError(llvm::sys::fs::make_absolute(absPath)));

284 return absPath;

285 });

286

287

288

289

290 llvm::StringMap<void *> exportSymbols;

293

294 for (auto &libPath : sharedLibPaths) {

295 auto lib = llvm::sys::DynamicLibrary::getPermanentLibrary(

296 libPath.str().str().c_str());

299

300

301 if (!initSym || !destroySim) {

302 jitDyLibPaths.push_back(libPath);

303 continue;

304 }

305

306 auto initFn = reinterpret_cast<LibraryInitFn>(initSym);

307 initFn(exportSymbols);

308

309 auto destroyFn = reinterpret_cast<LibraryDestroyFn>(destroySim);

310 destroyFns.push_back(destroyFn);

311 }

312 engine->destroyFns = std::move(destroyFns);

313

314

315

316 auto objectLinkingLayerCreator = [&](ExecutionSession &session) {

317 auto objectLayer = std::make_unique(

318 session, [sectionMemoryMapper =

319 options.sectionMemoryMapper](const MemoryBuffer &) {

320 return std::make_unique(sectionMemoryMapper);

321 });

322

323

324 if (engine->gdbListener)

325 objectLayer->registerJITEventListener(*engine->gdbListener);

326 if (engine->perfListener)

327 objectLayer->registerJITEventListener(*engine->perfListener);

328

329

330

331

332 const llvm::Triple &targetTriple = llvmModule->getTargetTriple();

333 if (targetTriple.isOSBinFormatCOFF()) {

334 objectLayer->setOverrideObjectFlagsWithResponsibilityFlags(true);

335 objectLayer->setAutoClaimResponsibilityForObjectSymbols(true);

336 }

337

338

339 for (auto &libPath : jitDyLibPaths) {

340 auto mb = llvm::MemoryBuffer::getFile(libPath);

341 if (!mb) {

342 errs() << "Failed to create MemoryBuffer for: " << libPath

343 << "\nError: " << mb.getError().message() << "\n";

344 continue;

345 }

346 auto &jd = session.createBareJITDylib(std::string(libPath));

347 auto loaded = DynamicLibrarySearchGenerator::Load(

348 libPath.str().c_str(), dataLayout.getGlobalPrefix());

349 if (!loaded) {

350 errs() << "Could not load " << libPath << ":\n " << loaded.takeError()

351 << "\n";

352 continue;

353 }

354 jd.addGenerator(std::move(*loaded));

355 cantFail(objectLayer->add(jd, std::move(mb.get())));

356 }

357

358 return objectLayer;

359 };

360

361

362

363 auto compileFunctionCreator = [&](JITTargetMachineBuilder jtmb)

364 -> Expected<std::unique_ptrIRCompileLayer::IRCompiler> {

365 if (options.jitCodeGenOptLevel)

366 jtmb.setCodeGenOptLevel(*options.jitCodeGenOptLevel);

367 return std::make_unique(std::move(tm),

368 engine->cache.get());

369 };

370

371

372 auto jit =

373 cantFail(llvm::orc::LLJITBuilder()

374 .setCompileFunctionCreator(compileFunctionCreator)

375 .setObjectLinkingLayerCreator(objectLinkingLayerCreator)

376 .setDataLayout(dataLayout)

378

379

380 ThreadSafeModule tsm(std::move(llvmModule), std::move(ctx));

382 cantFail(tsm.withModuleDo(

383 [&](llvm::Module &module) { return options.transformer(&module); }));

384 cantFail(jit->addIRModule(std::move(tsm)));

385 engine->jit = std::move(jit);

386

387

388 llvm::orc::JITDylib &mainJD = engine->jit->getMainJITDylib();

389 mainJD.addGenerator(

390 cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(

391 dataLayout.getGlobalPrefix())));

392

393

394 auto runtimeSymbolMap = [&](llvm::orc::MangleAndInterner interner) {

395 auto symbolMap = llvm::orc::SymbolMap();

396 for (auto &exportSymbol : exportSymbols)

397 symbolMap[interner(exportSymbol.getKey())] = {

398 llvm::orc::ExecutorAddr::fromPtr(exportSymbol.getValue()),

399 llvm::JITSymbolFlags::Exported};

400 return symbolMap;

401 };

402 engine->registerSymbols(runtimeSymbolMap);

403 return std::move(engine);

404}

405

410 return result.takeError();

411 return reinterpret_cast<void (*)(void **)>(result.get());

412}

413

415 auto expectedSymbol = jit->lookup(name);

416

417

418

419

420

421

422

423 if (!expectedSymbol) {

424 std::string errorMessage;

425 llvm::raw_string_ostream os(errorMessage);

426 llvm::handleAllErrors(expectedSymbol.takeError(),

427 [&os](llvm::ErrorInfoBase &ei) { ei.log(os); });

429 }

430

431 if (void *fptr = expectedSymbol->toPtr<void *>())

432 return fptr;

434}

435

440 if (!expectedFPtr)

441 return expectedFPtr.takeError();

442 auto fptr = *expectedFPtr;

443

444 (*fptr)(args.data());

445

446 return Error::success();

447}

448

450 if (isInitialized)

451 return;

452

453

454 if (!jit->getTargetTriple().isAArch64())

455 cantFail(jit->initialize(jit->getMainJITDylib()));

456 isInitialized = true;

457}

static void packFunctionArguments(Module *module)

Definition ExecutionEngine.cpp:144

static Error makeStringError(const Twine &message)

Wrap a string into an llvm::StringError.

Definition ExecutionEngine.cpp:60

static std::string makePackedFunctionName(StringRef name)

Definition ExecutionEngine.cpp:137

static llvm::ManagedStatic< PassManagerOptions > options

llvm::Expected< void(*)(void **)> lookupPacked(StringRef name) const

Looks up a packed-argument function wrapping the function with the given name and returns a pointer t...

Definition ExecutionEngine.cpp:407

void(*)(llvm::StringMap< void * > &) LibraryInitFn

Function type for init functions of shared libraries.

void(*)() LibraryDestroyFn

Function type for destroy functions of shared libraries.

void registerSymbols(llvm::function_ref< llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)> symbolMap)

Register symbols with this ExecutionEngine.

Definition ExecutionEngine.cpp:123

llvm::Expected< void * > lookup(StringRef name) const

Looks up the original function with the given name and returns a pointer to it.

Definition ExecutionEngine.cpp:414

static llvm::Expected< std::unique_ptr< ExecutionEngine > > create(Operation *op, const ExecutionEngineOptions &options={}, std::unique_ptr< llvm::TargetMachine > tm=nullptr)

Creates an execution engine for the given MLIR IR.

Definition ExecutionEngine.cpp:231

void dumpToObjectFile(StringRef filename)

Dump object code to output file filename.

Definition ExecutionEngine.cpp:101

static Result< T > result(T &t)

Helper function to wrap an output operand when using ExecutionEngine::invoke.

static constexpr const char *const kLibraryInitFnName

Name of init functions of shared libraries.

static constexpr const char *const kLibraryDestroyFnName

Name of destroy functions of shared libraries.

~ExecutionEngine()

Definition ExecutionEngine.cpp:219

llvm::Error invokePacked(StringRef name, MutableArrayRef< void * > args={})

Invokes the function with the given name passing it the list of opaque pointers to the actual argumen...

Definition ExecutionEngine.cpp:436

ExecutionEngine(bool enableObjectDump, bool enableGDBNotificationListener, bool enablePerfNotificationListener)

Definition ExecutionEngine.cpp:201

static void setupTargetTripleAndDataLayout(llvm::Module *llvmModule, llvm::TargetMachine *tm)

Set the target triple and the data layout for the input module based on the input TargetMachine.

Definition ExecutionEngine.cpp:131

void initialize()

Initialize the ExecutionEngine.

Definition ExecutionEngine.cpp:449

Operation is the basic unit of execution within MLIR.

Region & getRegion(unsigned index)

Returns the region held by this operation at position 'index'.

iterator_range< OpIterator > getOps()

A simple object cache following Lang's LLJITWithObjectCache example.

bool isEmpty()

Returns true if cache hasn't been populated yet.

Definition ExecutionEngine.cpp:99

void notifyObjectCompiled(const llvm::Module *m, llvm::MemoryBufferRef objBuffer) override

Definition ExecutionEngine.cpp:65

void dumpToObjectFile(StringRef filename)

Dump cached object to output file filename.

Definition ExecutionEngine.cpp:83

std::unique_ptr< llvm::MemoryBuffer > getObject(const llvm::Module *m) override

Definition ExecutionEngine.cpp:71

The OpAsmOpInterface, see OpAsmInterface.td for more details.

Include the generated interface declarations.

std::unique_ptr< llvm::ToolOutputFile > openOutputFile(llvm::StringRef outputFilename, std::string *errorMessage=nullptr)

Open the file specified by its name for writing.

llvm::DenseSet< ValueT, ValueInfoT > DenseSet

std::unique_ptr< llvm::Module > translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext, llvm::StringRef name="LLVMDialectModule", bool disableVerification=false)

Translates a given LLVM dialect module into an LLVM IR module living in the given context.