[mlir] Fix and test python bindings for dump_to_object_file

[ftynse]

mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp

58

Are you sure there's no unintended effects due to this being always-on? I'd rather take the flag as argument and thread that to Python.

[shabalin]

Make dump_to_object_file work even if no functions were invoked previously.

[shabalin]

Make sure temporary file is always deleted, even if test fails.

[ftynse]

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

49

Please add documentation.

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

111

Nit: please expand auto. MLIR only uses auto if the type is clear from immediate context.

113–116

So this just suppresses all errors? This must be at least documented. Even better if we propagate failure to avoid the surprise for the caller when the dumped object is silently missing some functions that couldn't be compiled.

258

I'd rather do m->getOps<LLVM::FuncOp> or at least indicate that traversal _inside_ functions isn't necessary.

[shabalin]

Address another round of comments.

[shabalin]

Expose enableObjectDump toggle to python & c bindings.

[shabalin]

Report failures on dump-triggered lazy compilation failures.

[ftynse]

LGTM with the last batch of comments addressed.

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

206

I think we need to use owning strings here. There doesn't seem to be a guarantee that the owner of the function name (attribute of the MLIR operation, so the content is presumably owned by the MLIRContext) outlives the execution engine. It looks unlikely, but someone may decide to free up the memory after compilation and before execution.

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

111

There's no need to have a reference to StringRef which itself a reference to a string.

[shabalin]

Address another round of review comments.

[NathanielMcVicar]

This change has broken the Windows buildbot, in an unfortunately opaque way. https://lab.llvm.org/buildbot/#/builders/13/builds/27291

******************** TEST 'MLIR :: python/execution_engine.py' FAILED ********************
Script:
--
: 'RUN: at line 1';   "C:/Program Files/Python39/python.exe" C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\test\python\execution_engine.py 2>&1 | c:\buildbot\mlir-x64-windows-ninja\build\bin\filecheck.exe C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\test\python\execution_engine.py
--
Exit Code: 1
Command Output (stdout):
--
$ ":" "RUN: at line 1"
$ "C:/Program Files/Python39/python.exe" "C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\test\python\execution_engine.py"
note: command had no output on stdout or stderr
error: command failed with exit status: 1
$ "c:\buildbot\mlir-x64-windows-ninja\build\bin\filecheck.exe" "C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\test\python\execution_engine.py"
--
********************

mlir/include/mlir-c/ExecutionEngine.h

 /// dialects that implement the `LLVMTranslationDialectInterface`). The module
 /// ownership stays with the client and can be destroyed as soon as the call
 /// returns. `optLevel` is the optimization level to be used for transformation
 /// and code generation. LLVM passes at `optLevel` are run before code
 /// generation. The number and array of paths corresponding to shared libraries
 /// that will be loaded are specified via `numPaths` and `sharedLibPaths`
 /// respectively.
 /// TODO: figure out other options.
-MLIR_CAPI_EXPORTED MlirExecutionEngine
-mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths,
-                          const MlirStringRef *sharedLibPaths);
+MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate(
+    MlirModule op, int optLevel, int numPaths,
+    const MlirStringRef *sharedLibPaths, bool enableObjectDump);
 
 /// Destroy an ExecutionEngine instance.
 MLIR_CAPI_EXPORTED void mlirExecutionEngineDestroy(MlirExecutionEngine jit);
 
 /// Checks whether an execution engine is null.
 static inline bool mlirExecutionEngineIsNull(MlirExecutionEngine jit) {
   return !jit.ptr;
 }

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

 public:
   void notifyObjectCompiled(const llvm::Module *m,
                             llvm::MemoryBufferRef objBuffer) override;
   std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *m) override;
 
   /// Dump cached object to output file `filename`.
   void dumpToObjectFile(StringRef filename);
 
+  /// Returns `true` if cache hasn't been populated yet.
+  bool isEmpty();
+
 private:
   llvm::StringMap<std::unique_ptr<llvm::MemoryBuffer>> cachedObjects;
 };
 
 struct ExecutionEngineOptions {
   /// If `llvmModuleBuilder` is provided, it will be used to create an LLVM
   /// module from the given MLIR IR. Otherwise, a default
   /// `translateModuleToLLVMIR` function will be used to translate to LLVM IR.
struct ExecutionEngineOptions {
 
   /// Specifies an existing `sectionMemoryMapper` to be associated with the
   /// compiled code. If none is provided, a default memory mapper that directly
   /// calls into the operating system is used.
   llvm::SectionMemoryManager::MemoryMapper *sectionMemoryMapper = nullptr;
 
   /// 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.
   bool enableGDBNotificationListener = true;
 
   /// If `enablePerfNotificationListener` is set, the JIT compiler will notify
   /// the llvm's global Perf notification listener.
   bool enablePerfNotificationListener = true;
 };
 
 /// JIT-backed execution engine for MLIR. Assumes the IR can be converted to
 /// LLVM IR. For each function, creates a wrapper function with the fixed
 /// interface
 ///
 ///     void _mlir_funcName(void **)
 ///
 /// where the only argument is interpreted as a list of pointers to the actual
 /// arguments of the function, followed by a pointer to the result.  This allows
 /// the engine to provide the caller with a generic function pointer that can
 /// 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.
   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
   create(Operation *op, const ExecutionEngineOptions &options = {});
 
   /// Looks up a packed-argument function wrapping the function with the given
   /// name and returns a pointer to it. Propagates errors in case of failure.
private:
   llvm::LLVMContext llvmContext;
 
   /// Underlying LLJIT.
   std::unique_ptr<llvm::orc::LLJIT> jit;
 
   /// Underlying cache.
   std::unique_ptr<SimpleObjectCache> cache;
 
+  /// Names of functions that may be looked up.
+  std::vector<std::string> functionNames;
+
   /// GDB notification listener.
   llvm::JITEventListener *gdbListener;
 
   /// Perf notification listener.
   llvm::JITEventListener *perfListener;
 };
 
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_EXECUTIONENGINE_H_

mlir/lib/Bindings/Python/ExecutionEngineModule.cpp

 PYBIND11_MODULE(_mlirExecutionEngine, m) {
   m.doc() = "MLIR Execution Engine";
 
   //----------------------------------------------------------------------------
   // Mapping of the top-level PassManager
   //----------------------------------------------------------------------------
   py::class_<PyExecutionEngine>(m, "ExecutionEngine", py::module_local())
       .def(py::init<>([](MlirModule module, int optLevel,
-                         const std::vector<std::string> &sharedLibPaths) {
+                         const std::vector<std::string> &sharedLibPaths,
+                         bool enableObjectDump) {
              llvm::SmallVector<MlirStringRef, 4> libPaths;
              for (const std::string &path : sharedLibPaths)
                libPaths.push_back({path.c_str(), path.length()});
-             MlirExecutionEngine executionEngine = mlirExecutionEngineCreate(
-                 module, optLevel, libPaths.size(), libPaths.data());
+             MlirExecutionEngine executionEngine =
+                 mlirExecutionEngineCreate(module, optLevel, libPaths.size(),
+                                           libPaths.data(), enableObjectDump);
              if (mlirExecutionEngineIsNull(executionEngine))
                throw std::runtime_error(
                    "Failure while creating the ExecutionEngine.");
              return new PyExecutionEngine(executionEngine);
            }),
            py::arg("module"), py::arg("opt_level") = 2,
            py::arg("shared_libs") = py::list(),
+           py::arg("enable_object_dump") = true,
            "Create a new ExecutionEngine instance for the given Module. The "
            "module must contain only dialects that can be translated to LLVM. "
            "Perform transformations and code generation at the optimization "
            "level `opt_level` if specified, or otherwise at the default "
            "level of two (-O2). Load a list of libraries specified in "
            "`shared_libs`.")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyExecutionEngine::getCapsule)

mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp

 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "llvm/ExecutionEngine/Orc/Mangling.h"
 #include "llvm/Support/TargetSelect.h"
 
 using namespace mlir;
 
 extern "C" MlirExecutionEngine
 mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths,
-                          const MlirStringRef *sharedLibPaths) {
+                          const MlirStringRef *sharedLibPaths,
+                          bool enableObjectDump) {
   static bool initOnce = [] {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmParser(); // needed for inline_asm
     llvm::InitializeNativeTargetAsmPrinter();
     return true;
   }();
   (void)initOnce;
 
mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths,
   // specified optimization level.
   auto llvmOptLevel = static_cast<llvm::CodeGenOpt::Level>(optLevel);
   auto transformer = mlir::makeOptimizingTransformer(
       llvmOptLevel, /*sizeLevel=*/0, /*targetMachine=*/tmOrError->get());
   ExecutionEngineOptions jitOptions;
   jitOptions.transformer = transformer;
   jitOptions.jitCodeGenOptLevel = llvmOptLevel;
   jitOptions.sharedLibPaths = libPaths;
+  jitOptions.enableObjectDump = enableObjectDump;
   auto jitOrError = ExecutionEngine::create(unwrap(op), jitOptions);
   if (!jitOrError) {
     consumeError(jitOrError.takeError());
     return MlirExecutionEngine{nullptr};
   }
   return wrap(jitOrError->release());
 }
 

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

void SimpleObjectCache::dumpToObjectFile(StringRef outputFilename) {
 
   // Dump the object generated for a single module to the output file.
   assert(cachedObjects.size() == 1 && "Expected only one object entry.");
   auto &cachedObject = cachedObjects.begin()->second;
   file->os() << cachedObject->getBuffer();
   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 (std::string &functionName : functionNames) {
+      auto result = lookupPacked(functionName);
+      if (!result) {
+        llvm::errs() << "Could not compile " << functionName << ":\n  "
+                     << result.takeError() << "\n";
+        return;
+      }
+    }
+  }
   cache->dumpToObjectFile(filename);
 }
 
 void ExecutionEngine::registerSymbols(
     llvm::function_ref<SymbolMap(MangleAndInterner)> symbolMap) {
   auto &mainJitDylib = jit->getMainJITDylib();
   cantFail(mainJitDylib.define(
       absoluteSymbols(symbolMap(llvm::orc::MangleAndInterner(
if (!result->getType()->isVoidTy()) {
       builder.CreateStore(result, retPtr);
     }
 
     // The interface function returns void.
     builder.CreateRetVoid();
   }
 }
 
-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),
       perfListener(nullptr) {
   if (enablePerfNotificationListener) {
     if (auto *listener = llvm::JITEventListener::createPerfJITEventListener())
       perfListener = listener;
     else if (auto *listener =
                  llvm::JITEventListener::createIntelJITEventListener())
       perfListener = listener;
   }
 }
 
 Expected<std::unique_ptr<ExecutionEngine>>
 ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options) {
   auto engine = std::make_unique<ExecutionEngine>(
-      options.enableObjectCache, options.enableGDBNotificationListener,
+      options.enableObjectDump, options.enableGDBNotificationListener,
       options.enablePerfNotificationListener);
 
+  // Remember all entry-points if object dumping is enabled.
+  if (options.enableObjectDump) {
+    for (auto funcOp : m->getRegion(0).getOps<LLVM::LLVMFuncOp>()) {
+      StringRef funcName = funcOp.getSymName();
+      engine->functionNames.push_back(funcName.str());
+    }
+  }
+
   std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
   auto llvmModule = options.llvmModuleBuilder
                         ? options.llvmModuleBuilder(m, *ctx)
                         : translateModuleToLLVMIR(m, *ctx);
   if (!llvmModule)
     return makeStringError("could not convert to LLVM IR");
   // FIXME: the triple should be passed to the translation or dialect conversion
   // instead of this.  Currently, the LLVM module created above has no triple

mlir/lib/ExecutionEngine/JitRunner.cpp

static Error compileAndExecute(Options &options, Operation *module,
   };
 
   mlir::ExecutionEngineOptions engineOptions;
   engineOptions.llvmModuleBuilder = config.llvmModuleBuilder;
   if (config.transformer)
     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();
 
   auto engine = std::move(*expectedEngine);
   engine->registerSymbols(runtimeSymbolMap);
 
   auto expectedFPtr = engine->lookupPacked(entryPoint);

mlir/test/CAPI/execution_engine.c

 "    %res = arith.addi %arg0, %arg0 : i32                                        \n"
 "    return %res : i32                                                           \n"
 "  }                                                                             \n"
 "}"));
   // clang-format on
   lowerModuleToLLVM(ctx, module);
   mlirRegisterAllLLVMTranslations(ctx);
   MlirExecutionEngine jit = mlirExecutionEngineCreate(
-      module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL);
+      module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL,
+      /*enableObjectDump=*/false);
   if (mlirExecutionEngineIsNull(jit)) {
     fprintf(stderr, "Execution engine creation failed");
     exit(2);
   }
   int input = 42;
   int result = -1;
   void *args[2] = {&input, &result};
   if (mlirLogicalResultIsFailure(mlirExecutionEngineInvokePacked(

mlir/test/python/execution_engine.py

 # 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 *
 from mlir.runtime import *
 
 
 # Log everything to stderr and flush so that we have a unified stream to match
 # errors/info emitted by MLIR to stderr.
execution_engine = ExecutionEngine(
         opt_level=3,
         shared_libs=shared_libs)
     execution_engine.invoke("main")
     # CHECK: Unranked Memref
     # CHECK: [{{.*}}]
 
 
 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)