Break apart the MLIR ExecutionEngine from core python module.

[stellaraccident]

Summary

• For python projects that don't need JIT/ExecutionEngine, cuts the number of files to compile roughly in half (with similar reduction in end binary size).

[stellaraccident]

fyi - this has some failures. I'll need to triage more. will get to it soon.

[stellaraccident]

Update test class name that changed.

mlir/lib/Bindings/Python/ExecutionEngine.h

-//===- ExecutionEngine.h - ExecutionEngine submodule of pybind module -----===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_BINDINGS_PYTHON_EXECUTIONENGINE_H
-#define MLIR_BINDINGS_PYTHON_EXECUTIONENGINE_H
-
-#include "PybindUtils.h"
-
-namespace mlir {
-namespace python {
-
-void populateExecutionEngineSubmodule(pybind11::module &m);
-
-} // namespace python
-} // namespace mlir
-
-#endif // MLIR_BINDINGS_PYTHON_EXECUTIONENGINE_H

mlir/lib/Bindings/Python/ExecutionEngineModule.cpp

-//===- ExecutionEngine.cpp - Python MLIR ExecutionEngine Bindings ---------===//
+//===- ExecutionEngineModule.cpp - Python module for execution engine -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "ExecutionEngine.h"
-
-#include "IRModule.h"
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/ExecutionEngine.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 namespace {
 
 /// Owning Wrapper around an ExecutionEngine.

 
 private:
   MlirExecutionEngine executionEngine;
 };
 
 } // anonymous namespace
 
 /// Create the `mlir.execution_engine` module here.
-void mlir::python::populateExecutionEngineSubmodule(py::module &m) {
+PYBIND11_MODULE(_mlirExecutionEngine, m) {
+  m.doc() = "MLIR Execution Engine";
+
   //----------------------------------------------------------------------------
   // Mapping of the top-level PassManager
   //----------------------------------------------------------------------------
   py::class_<PyExecutionEngine>(m, "ExecutionEngine")
-      .def(py::init<>([](PyModule &module, int optLevel,
+      .def(py::init<>([](MlirModule module, int optLevel,
                          const std::vector<std::string> &sharedLibPaths) {
              llvm::SmallVector<MlirStringRef, 4> libPaths;
              for (const std::string &path : sharedLibPaths)
                libPaths.push_back({path.c_str(), path.length()});
              MlirExecutionEngine executionEngine = mlirExecutionEngineCreate(
-                 module.get(), optLevel, libPaths.size(), libPaths.data());
+                 module, optLevel, libPaths.size(), libPaths.data());
              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(),
            "Create a new ExecutionEngine instance for the given Module. The "

mlir/lib/Bindings/Python/MainModule.cpp

 //===- MainModule.cpp - Main pybind module --------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include <tuple>
 
 #include "PybindUtils.h"
 
 #include "Dialects.h"
-#include "ExecutionEngine.h"
 #include "Globals.h"
 #include "IRModule.h"
 #include "Pass.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
PYBIND11_MODULE(_mlir, m) {
   populateIRAttributes(irModule);
   populateIRTypes(irModule);
 
   // Define and populate PassManager submodule.
   auto passModule =
       m.def_submodule("passmanager", "MLIR Pass Management Bindings");
   populatePassManagerSubmodule(passModule);
 
-  // Define and populate ExecutionEngine submodule.
-  auto executionEngineModule =
-      m.def_submodule("execution_engine", "MLIR JIT Execution Engine");
-  populateExecutionEngineSubmodule(executionEngineModule);
-
   // Define and populate dialect submodules.
   auto dialectsModule = m.def_submodule("dialects");
   auto linalgModule = dialectsModule.def_submodule("linalg");
   populateDialectLinalgSubmodule(linalgModule);
   populateDialectSparseTensorSubmodule(
       dialectsModule.def_submodule("sparse_tensor"), irModule);
 }

mlir/python/CMakeLists.txt

SOURCES
     ${PYTHON_SOURCE_DIR}/MainModule.cpp
     ${PYTHON_SOURCE_DIR}/IRAffine.cpp
     ${PYTHON_SOURCE_DIR}/IRAttributes.cpp
     ${PYTHON_SOURCE_DIR}/IRCore.cpp
     ${PYTHON_SOURCE_DIR}/IRModule.cpp
     ${PYTHON_SOURCE_DIR}/IRTypes.cpp
     ${PYTHON_SOURCE_DIR}/PybindUtils.cpp
     ${PYTHON_SOURCE_DIR}/Pass.cpp
-    ${PYTHON_SOURCE_DIR}/ExecutionEngine.cpp # TODO: Break this out.
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIDebug
     MLIRCAPIIR
     MLIRCAPIRegistration  # TODO: See about dis-aggregating
 
     # Dialects
     MLIRCAPILinalg  # TODO: Remove when above is removed.
     MLIRCAPISparseTensor  # TODO: Remove when above is removed.
     MLIRCAPIStandard
-
-    # Execution engine (remove once disaggregated).
-    MLIRCEXECUTIONENGINE
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
   MODULE_NAME _mlirAllPassesRegistration
   SOURCES
     ${PYTHON_SOURCE_DIR}/AllPassesRegistration.cpp
   PRIVATE_LINK_LIBS
     LLVMSupport
declare_mlir_python_extension(MLIRPythonExtension.Conversions
   SOURCES
     ${PYTHON_SOURCE_DIR}/Conversions/Conversions.cpp
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
   MLIRCAPIConversion
 )
 
+declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
+  MODULE_NAME _mlirExecutionEngine
+  ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/ExecutionEngineModule.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCEXECUTIONENGINE
+)
+
 declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
   MODULE_NAME _mlirGPUPasses
   ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
   SOURCES
     ${PYTHON_SOURCE_DIR}/GPUPasses.cpp
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS

mlir/python/mlir/execution_engine.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Simply a wrapper around the extension module of the same name.
-from ._cext_loader import _cext
+from ._cext_loader import load_extension
+_execution_engine = load_extension("_mlirExecutionEngine")
 import ctypes
 
-class ExecutionEngine(_cext.execution_engine.ExecutionEngine):
+__all__ = [
+  "ExecutionEngine",
+]
+
+class ExecutionEngine(_execution_engine.ExecutionEngine):
 
   def lookup(self, name):
     """Lookup a function emitted with the `llvm.emit_c_interface`
     attribute and returns a ctype callable.
     Raise a RuntimeError if the function isn't found.
     """
     func = self.raw_lookup("_mlir_ciface_" + name)
     if not func:

mlir/test/python/execution_engine.py

 }
     """)
     execution_engine = ExecutionEngine(module)
     execution_engine_capsule = execution_engine._CAPIPtr
     # CHECK: mlir.execution_engine.ExecutionEngine._CAPIPtr
     log(repr(execution_engine_capsule))
     execution_engine._testing_release()
     execution_engine1 = ExecutionEngine._CAPICreate(execution_engine_capsule)
-    # CHECK: _mlir.execution_engine.ExecutionEngine
+    # CHECK: _mlirExecutionEngine.ExecutionEngine
     log(repr(execution_engine1))
 
 run(testCapsule)
 
 # Test invalid ExecutionEngine creation
 # CHECK-LABEL: TEST: testInvalidModule
 def testInvalidModule():
   with Context():

 # Test simple ExecutionEngine execution
 # CHECK-LABEL: TEST: testInvokeVoid
 def testInvokeVoid():
   with Context():
     module = Module.parse(r"""
 func @void() attributes { llvm.emit_c_interface } {
   return
 }
-    """)    
+    """)
     execution_engine = ExecutionEngine(lowerToLLVM(module))
     # Nothing to check other than no exception thrown here.
     execution_engine.invoke("void")
 
 run(testInvokeVoid)
 
 
 # Test argument passing and result with a simple float addition.
func @callback_memref(%arg0: memref<*xf32>) attributes { llvm.emit_c_interface } {
   return
 }
 func private @some_callback_into_python(memref<*xf32>) -> () attributes { llvm.emit_c_interface }
 """
         )
         execution_engine = ExecutionEngine(lowerToLLVM(module))
         execution_engine.register_runtime("some_callback_into_python", callback)
         inp_arr = np.array([[1.0, 2.0], [3.0, 4.0]], np.float32)
-        # CHECK: Inside callback: 
+        # CHECK: Inside callback:
         # CHECK{LITERAL}: [[1. 2.]
         # CHECK{LITERAL}:  [3. 4.]]
         execution_engine.invoke(
             "callback_memref",
             ctypes.pointer(ctypes.pointer(get_unranked_memref_descriptor(inp_arr))),
         )
         inp_arr_1 = np.array([5, 6, 7], dtype=np.float32)
         strided_arr = np.lib.stride_tricks.as_strided(
             inp_arr_1, strides=(4, 0), shape=(3, 4)
         )
-        # CHECK: Inside callback: 
+        # CHECK: Inside callback:
         # CHECK{LITERAL}: [[5. 5. 5. 5.]
         # CHECK{LITERAL}:  [6. 6. 6. 6.]
         # CHECK{LITERAL}:  [7. 7. 7. 7.]]
         execution_engine.invoke(
             "callback_memref",
             ctypes.pointer(
                 ctypes.pointer(get_unranked_memref_descriptor(strided_arr))
             ),