ExecutionEngine: allow for running MLIR passes during JIT-compilation

    The existing implementation of the ExecutionEngine unconditionally runs a list
    of "default" MLIR passes on the module upon creation.  These passes include,
    among others, dialect conversions from affine to standard and from standard to
    LLVM IR dialects.  In some cases, these conversions might have been performed
    before ExecutionEngine is created.  More advanced use cases may be performing
    additional transformations that the "default" passes will conflict with.
    Provide an overload for ExecutionEngine::create that takes a PassManager
    configured with the passes to run on the module.  If it is not provided, do not
    run any passes.  The engine will not be created if the input module, after the
    pass manager, has any other dialect than the LLVM IR dialect.

--

PiperOrigin-RevId: 242127393

1 files changed, 11 insertions, 7 deletions

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index e63a3e8569c..e9f99848e37 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -273,19 +273,15 @@ void packFunctionArguments(llvm::Module *module) {
 // Out of line for PIMPL unique_ptr.
 ExecutionEngine::~ExecutionEngine() = default;
 
-std::unique_ptr<llvm::Module> translateModuleToLLVMIR(Module &m);
-
 Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
-    Module *m, std::function<llvm::Error(llvm::Module *)> transformer) {
+    Module *m, PassManager *pm,
+    std::function<llvm::Error(llvm::Module *)> transformer) {
   auto engine = llvm::make_unique<ExecutionEngine>();
   auto expectedJIT = impl::OrcJIT::createDefault(transformer);
   if (!expectedJIT)
     return expectedJIT.takeError();
 
-  // Construct and run the default MLIR pipeline.
-  PassManager manager;
-  getDefaultPasses(manager, {});
-  if (failed(manager.run(m)))
+  if (pm && failed(pm->run(m)))
     return make_string_error("passes failed");
 
   auto llvmModule = translateModuleToLLVMIR(*m);
@@ -304,6 +300,14 @@ Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
   return std::move(engine);
 }
 
+Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
+    Module *m, std::function<llvm::Error(llvm::Module *)> transformer) {
+  // Construct and run the default MLIR pipeline.
+  PassManager manager;
+  getDefaultPasses(manager, {});
+  return create(m, &manager, transformer);
+}
+
 Expected<void (*)(void **)> ExecutionEngine::lookup(StringRef name) const {
   auto expectedSymbol = jit->lookup(makePackedFunctionName(name));
   if (!expectedSymbol)