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 @@ -46,6 +46,8 @@ /// Dump cached object to output file `filename`. void dumpToObjectFile(StringRef filename); + bool isEmpty(); + private: llvm::StringMap> cachedObjects; }; @@ -77,8 +79,8 @@ /// If `enableObjectCache` is set, the JIT compiler will create one to store /// the object generated for the given module. The contents of the cache can - /// be dumped to a file via the `dumpToObjectfile` method. - bool enableObjectCache = false; + /// be dumped to a file via the `dumpToObjectFile` method. + bool enableObjectDump = false; /// If enable `enableGDBNotificationListener` is set, the JIT compiler will /// notify the llvm's global GDB notification listener. @@ -101,7 +103,7 @@ /// be used to invoke the JIT-compiled function. class ExecutionEngine { public: - ExecutionEngine(bool enableObjectCache, bool enableGDBNotificationListener, + ExecutionEngine(bool enableObjectDump, bool enableGDBNotificationListener, bool enablePerfNotificationListener); /// Creates an execution engine for the given MLIR IR. @@ -199,6 +201,9 @@ /// Underlying cache. std::unique_ptr cache; + /// Names of functions that may be looked up. + std::vector functionNames; + /// GDB notification listener. llvm::JITEventListener *gdbListener; diff --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp --- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp @@ -54,6 +54,7 @@ jitOptions.transformer = transformer; jitOptions.jitCodeGenOptLevel = llvmOptLevel; jitOptions.sharedLibPaths = libPaths; + jitOptions.enableObjectDump = true; auto jitOrError = ExecutionEngine::create(unwrap(op), jitOptions); if (!jitOrError) { consumeError(jitOrError.takeError()); 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 @@ -96,12 +96,26 @@ file->keep(); } +bool SimpleObjectCache::isEmpty() { return cachedObjects.size() == 0; } + void ExecutionEngine::dumpToObjectFile(StringRef filename) { if (cache == nullptr) { llvm::errs() << "cannot dump ExecutionEngine object code to file: " "object cache is disabled\n"; return; } + // Compilation is lazy and it doesn't populate object cache unless requested. + // In case object dump is requested before cache is populated, we need to + // force compilation manually. + if (cache->isEmpty()) { + for (auto &functionName : functionNames) { + auto result = lookupPacked(functionName); + if (!result) { + llvm::handleAllErrors(result.takeError(), + [](llvm::ErrorInfoBase &ei) {}); + } + } + } cache->dumpToObjectFile(filename); } @@ -214,10 +228,11 @@ } } -ExecutionEngine::ExecutionEngine(bool enableObjectCache, +ExecutionEngine::ExecutionEngine(bool enableObjectDump, bool enableGDBNotificationListener, bool enablePerfNotificationListener) - : cache(enableObjectCache ? new SimpleObjectCache() : nullptr), + : cache(enableObjectDump ? new SimpleObjectCache() : nullptr), + functionNames(), gdbListener(enableGDBNotificationListener ? llvm::JITEventListener::createGDBRegistrationListener() : nullptr), @@ -234,9 +249,17 @@ Expected> ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options) { auto engine = std::make_unique( - options.enableObjectCache, options.enableGDBNotificationListener, + options.enableObjectDump, options.enableGDBNotificationListener, options.enablePerfNotificationListener); + // Remember all entry-points if object dumping is enabled. + if (options.enableObjectDump) { + m->walk([&](LLVM::LLVMFuncOp funcOp) { + StringRef funcName = funcOp.getSymName(); + engine->functionNames.push_back(funcName); + }); + } + std::unique_ptr ctx(new llvm::LLVMContext); auto llvmModule = options.llvmModuleBuilder ? options.llvmModuleBuilder(m, *ctx) diff --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp --- a/mlir/lib/ExecutionEngine/JitRunner.cpp +++ b/mlir/lib/ExecutionEngine/JitRunner.cpp @@ -229,7 +229,7 @@ engineOptions.transformer = config.transformer; engineOptions.jitCodeGenOptLevel = jitCodeGenOptLevel; engineOptions.sharedLibPaths = executionEngineLibs; - engineOptions.enableObjectCache = true; + engineOptions.enableObjectDump = true; auto expectedEngine = mlir::ExecutionEngine::create(module, engineOptions); if (!expectedEngine) return expectedEngine.takeError(); diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py --- a/mlir/test/python/execution_engine.py +++ b/mlir/test/python/execution_engine.py @@ -1,6 +1,6 @@ # RUN: %PYTHON %s 2>&1 | FileCheck %s # REQUIRES: host-supports-jit -import gc, sys +import gc, sys, os, tempfile from mlir.ir import * from mlir.passmanager import * from mlir.execution_engine import * @@ -552,3 +552,40 @@ run(testNanoTime) + + +# Test that nano time clock is available. +# CHECK-LABEL: TEST: testDumpToObjectFile +def testDumpToObjectFile(): + _, object_path = tempfile.mkstemp(suffix=".o") + + try: + with Context(): + module = Module.parse(""" + module { + func.func @main() attributes { llvm.emit_c_interface } { + return + } + }""") + + execution_engine = ExecutionEngine( + lowerToLLVM(module), + opt_level=3) + + # CHECK: Object file exists: True + print(f"Object file exists: {os.path.exists(object_path)}") + # CHECK: Object file is empty: True + print(f"Object file is empty: {os.path.getsize(object_path) == 0}") + + execution_engine.dump_to_object_file(object_path) + + # CHECK: Object file exists: True + print(f"Object file exists: {os.path.exists(object_path)}") + # CHECK: Object file is empty: False + print(f"Object file is empty: {os.path.getsize(object_path) == 0}") + + finally: + os.remove(object_path) + + +run(testDumpToObjectFile)