[mlir] Make C/Python ExecutionEngine constructible with an Operation.

mlir/include/mlir-c/ExecutionEngine.h

@@ -42,8 +42,15 @@ DEFINE_C_API_STRUCT(MlirExecutionEngine, void);
 /// that will be loaded are specified via `numPaths` and `sharedLibPaths`
 /// respectively.
 /// TODO: figure out other options.
+MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreateFromOp(
+    MlirOperation op, int optLevel, int numPaths,
+    const MlirStringRef *sharedLibPaths, bool enableObjectDump);
+
+// Deprecated variant which takes an MlirModule instead of an operation.
+// This is being preserved as of 2024-Mar for short term consistency and should
+// be dropped soon.
 MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate(
-    MlirModule op, int optLevel, int numPaths,
+    MlirModule module, int optLevel, int numPaths,
     const MlirStringRef *sharedLibPaths, bool enableObjectDump);
 
 /// Destroy an ExecutionEngine instance.

mlir/lib/Bindings/Python/ExecutionEngineModule.cpp

@@ -71,15 +71,34 @@ PYBIND11_MODULE(_mlirExecutionEngine, m) {
   // Mapping of the top-level PassManager
   //----------------------------------------------------------------------------
   py::class_<PyExecutionEngine>(m, "ExecutionEngine", py::module_local())
-      .def(py::init<>([](MlirModule module, int optLevel,
+      .def(py::init<>([](py::object operation_or_module, int optLevel,
                          const std::vector<std::string> &sharedLibPaths,
                          bool enableObjectDump) {
+             // Manually type cast from either a Module or Operation. The
+             // automatic type casters do not handle such cascades well,
+             // so be explicit.
+             py::object capsule = mlirApiObjectToCapsule(operation_or_module);
+             MlirOperation module_op =
+                 mlirPythonCapsuleToOperation(capsule.ptr());
+             if (mlirOperationIsNull(module_op)) {
+               // If null, then a PyErr_Set has set an exception, which we must
+               // clear.
+               PyErr_Clear();
+               MlirModule mod = mlirPythonCapsuleToModule(capsule.ptr());
+               if (mlirModuleIsNull(mod)) {
+                 throw py::type_error(
+                     "ExecutionEngine expects a Module or Operation");
+               }
+               module_op = mlirModuleGetOperation(mod);
+             }
+
              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(), enableObjectDump);
+                 mlirExecutionEngineCreateFromOp(
+                     module_op, optLevel, libPaths.size(), libPaths.data(),
+                     enableObjectDump);
              if (mlirExecutionEngineIsNull(executionEngine))
                throw std::runtime_error(
                    "Failure while creating the ExecutionEngine.");

mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp

@@ -20,9 +20,18 @@
 using namespace mlir;
 
 extern "C" MlirExecutionEngine
-mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths,
+mlirExecutionEngineCreate(MlirModule module, int optLevel, int numPaths,
                           const MlirStringRef *sharedLibPaths,
                           bool enableObjectDump) {
+  return mlirExecutionEngineCreateFromOp(mlirModuleGetOperation(module),
+                                         optLevel, numPaths, sharedLibPaths,
+                                         enableObjectDump);
+}
+
+extern "C" MlirExecutionEngine
+mlirExecutionEngineCreateFromOp(MlirOperation op, int optLevel, int numPaths,
+                                const MlirStringRef *sharedLibPaths,
+                                bool enableObjectDump) {
   static bool initOnce = [] {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmParser(); // needed for inline_asm
@@ -104,9 +113,8 @@ extern "C" void mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit,
                                                   void *sym) {
   unwrap(jit)->registerSymbols([&](llvm::orc::MangleAndInterner interner) {
     llvm::orc::SymbolMap symbolMap;
-    symbolMap[interner(unwrap(name))] =
-        { llvm::orc::ExecutorAddr::fromPtr(sym),
-          llvm::JITSymbolFlags::Exported };
+    symbolMap[interner(unwrap(name))] = {llvm::orc::ExecutorAddr::fromPtr(sym),
+                                         llvm::JITSymbolFlags::Exported};
     return symbolMap;
   });
 }

mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi

@@ -4,7 +4,7 @@
 #   * Relative imports for cross-module references.
 #   * Add __all__
 
-from typing import List, Sequence
+from typing import List, Sequence,Union
 
 from ._mlir import ir as _ir
 
@@ -13,7 +13,7 @@ __all__ = [
 ]
 
 class ExecutionEngine:
-    def __init__(self, module: _ir.Module, opt_level: int = 2, shared_libs: Sequence[str] = ...) -> None: ...
+    def __init__(self, module: Union[_ir.Operation, _ir.Module], opt_level: int = 2, shared_libs: Sequence[str] = ...) -> None: ...
     def _CAPICreate(self) -> object: ...
     def _testing_release(self) -> None: ...
     def dump_to_object_file(self, file_name: str) -> None: ...