diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h --- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h +++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h @@ -63,23 +63,37 @@ ExecutionEngine(bool enableObjectCache, bool enableGDBNotificationListener, bool enablePerfNotificationListener); - /// Creates an execution engine for the given module. If `transformer` is - /// provided, it will be called on the LLVM module during JIT-compilation and - /// can be used, e.g., for reporting or optimization. `jitCodeGenOptLevel`, - /// when provided, is used as the optimization level for target code - /// generation. If `sharedLibPaths` are provided, the underlying - /// JIT-compilation will open and link the shared libraries for symbol - /// resolution. If `enableObjectCache` is set, the JIT compiler will create - /// one to store the object generated for the given module. If enable - /// `enableGDBNotificationListener` is set, the JIT compiler will notify - /// the llvm's global GDB notification listener. If - /// `enablePerfNotificationListener` is set, the JIT compiler will notify + /// Creates an execution engine for the given module. + /// + /// If `transformer` is provided, it will be called on the LLVM module during + /// JIT-compilation and can be used, e.g., for reporting or optimization. + /// + /// `jitCodeGenOptLevel`, when provided, is used as the optimization level for + /// target code generation. + /// + /// If `sharedLibPaths` are provided, the underlying JIT-compilation will + /// open and link the shared libraries for symbol resolution. + /// + /// If `symbolMap` is provided, it will be called with a name mangler bound + /// to the underlying JIT-compilation engine. All symbols in the returned + /// symbol map will be added for symbol resolution. + /// + /// If `enableObjectCache` is set, the JIT compiler will create one to store + /// the object generated for the given module. + /// + /// If enable `enableGDBNotificationListener` is set, the JIT compiler will + /// notify the llvm's global GDB notification listener. + /// + /// If `enablePerfNotificationListener` is set, the JIT compiler will notify /// the llvm's global Perf notification listener. static llvm::Expected> create(ModuleOp m, - std::function transformer = {}, + llvm::function_ref transformer = {}, Optional jitCodeGenOptLevel = llvm::None, - ArrayRef sharedLibPaths = {}, bool enableObjectCache = true, + ArrayRef sharedLibPaths = {}, + llvm::function_ref + symbolMap = {}, + bool enableObjectCache = true, bool enableGDBNotificationListener = true, bool enablePerfNotificationListener = true); diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -194,10 +194,13 @@ : nullptr) {} Expected> ExecutionEngine::create( - ModuleOp m, std::function transformer, + ModuleOp m, llvm::function_ref transformer, Optional jitCodeGenOptLevel, - ArrayRef sharedLibPaths, bool enableObjectCache, - bool enableGDBNotificationListener, bool enablePerfNotificationListener) { + ArrayRef sharedLibPaths, + llvm::function_ref + symbolMap, + bool enableObjectCache, bool enableGDBNotificationListener, + bool enablePerfNotificationListener) { auto engine = std::make_unique( enableObjectCache, enableGDBNotificationListener, enablePerfNotificationListener); @@ -288,6 +291,13 @@ cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess( dataLayout.getGlobalPrefix()))); + // Add additional symbols provided via the symbol map. + if (symbolMap) { + auto symbols = symbolMap( + llvm::orc::MangleAndInterner(mainJD.getExecutionSession(), dataLayout)); + cantFail(mainJD.define(absoluteSymbols(std::move(symbols)))); + } + return std::move(engine); }