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,38 @@ 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 + /// Unmangled symbol names to JIT evaluated symbol mapping. Symbols will be + /// mangled before adding to the jit engine. + using SymbolMap = llvm::DenseMap; + + /// 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 not empty, the underlying JIT-compilation will use it + /// 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, + ArrayRef sharedLibPaths = {}, + SymbolMap symbolMap = 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 @@ -196,8 +196,9 @@ Expected> ExecutionEngine::create( ModuleOp m, std::function transformer, Optional jitCodeGenOptLevel, - ArrayRef sharedLibPaths, bool enableObjectCache, - bool enableGDBNotificationListener, bool enablePerfNotificationListener) { + ArrayRef sharedLibPaths, SymbolMap symbolMap, + bool enableObjectCache, bool enableGDBNotificationListener, + bool enablePerfNotificationListener) { auto engine = std::make_unique( enableObjectCache, enableGDBNotificationListener, enablePerfNotificationListener); @@ -288,6 +289,14 @@ cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess( dataLayout.getGlobalPrefix()))); + // Add additional symbols provided via the symbol map. + llvm::orc::MangleAndInterner Mangle(mainJD.getExecutionSession(), dataLayout); + llvm::orc::SymbolMap mangledSymbolMap; + for (auto &kv : symbolMap) { + mangledSymbolMap[Mangle(kv.first)] = std::move(kv.second); + } + cantFail(mainJD.define(absoluteSymbols(mangledSymbolMap))); + return std::move(engine); }