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,25 +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 = {}, - Optional jitCodeGenOptLevel = llvm::None, - ArrayRef sharedLibPaths = {}, bool enableObjectCache = true, - bool enableGDBNotificationListener = true, - bool enablePerfNotificationListener = true); + static llvm::Expected> create( + ModuleOp m, std::function transformer = {}, + Optional jitCodeGenOptLevel = llvm::None, + ArrayRef sharedLibPaths = {}, + std::function + symbolMap = {}, + bool enableObjectCache = true, bool enableGDBNotificationListener = true, + bool enablePerfNotificationListener = true); /// Looks up a packed-argument function with the given name and returns a /// pointer to it. Propagates errors in case of failure. 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 @@ -196,8 +196,10 @@ Expected> ExecutionEngine::create( ModuleOp m, std::function transformer, Optional jitCodeGenOptLevel, - ArrayRef sharedLibPaths, bool enableObjectCache, - bool enableGDBNotificationListener, bool enablePerfNotificationListener) { + ArrayRef sharedLibPaths, + std::function symbolMap, + bool enableObjectCache, bool enableGDBNotificationListener, + bool enablePerfNotificationListener) { auto engine = std::make_unique( enableObjectCache, enableGDBNotificationListener, enablePerfNotificationListener); @@ -288,6 +290,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); }