[Mlir-commits] [mlir] 0cdf491 - Break apart the MLIR ExecutionEngine from core python module.
Stella Laurenzo
llvmlistbot at llvm.org
Wed Jul 28 17:01:49 PDT 2021
Author: Stella Laurenzo
Date: 2021-07-28T23:59:32Z
New Revision: 0cdf4915019a8ebc6570229cf140ad879dfaef56
URL: https://github.com/llvm/llvm-project/commit/0cdf4915019a8ebc6570229cf140ad879dfaef56
DIFF: https://github.com/llvm/llvm-project/commit/0cdf4915019a8ebc6570229cf140ad879dfaef56.diff
LOG: Break apart the MLIR ExecutionEngine from core python module.
* 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).
Differential Revision: https://reviews.llvm.org/D106992
Added:
mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
Modified:
mlir/lib/Bindings/Python/MainModule.cpp
mlir/python/CMakeLists.txt
mlir/python/mlir/execution_engine.py
mlir/test/python/execution_engine.py
Removed:
mlir/lib/Bindings/Python/ExecutionEngine.cpp
mlir/lib/Bindings/Python/ExecutionEngine.h
################################################################################
diff --git a/mlir/lib/Bindings/Python/ExecutionEngine.h b/mlir/lib/Bindings/Python/ExecutionEngine.h
deleted file mode 100644
index cc61648b569f3..0000000000000
--- a/mlir/lib/Bindings/Python/ExecutionEngine.h
+++ /dev/null
@@ -1,22 +0,0 @@
-//===- 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
diff --git a/mlir/lib/Bindings/Python/ExecutionEngine.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
similarity index 92%
rename from mlir/lib/Bindings/Python/ExecutionEngine.cpp
rename to mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
index 089c29507c797..510e3f8dd7929 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngine.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
@@ -1,4 +1,4 @@
-//===- 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.
@@ -6,11 +6,9 @@
//
//===----------------------------------------------------------------------===//
-#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;
@@ -54,18 +52,20 @@ class PyExecutionEngine {
} // 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.");
diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index 6e861c2f2f761..073ac90375a30 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/mlir/lib/Bindings/Python/MainModule.cpp
@@ -11,7 +11,6 @@
#include "PybindUtils.h"
#include "Dialects.h"
-#include "ExecutionEngine.h"
#include "Globals.h"
#include "IRModule.h"
#include "Pass.h"
@@ -93,11 +92,6 @@ PYBIND11_MODULE(_mlir, m) {
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");
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 15b1813894a7b..f5b261e8a8d50 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -169,7 +169,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
${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
@@ -181,9 +180,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
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
@@ -219,6 +215,17 @@ declare_mlir_python_extension(MLIRPythonExtension.Conversions
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
diff --git a/mlir/python/mlir/execution_engine.py b/mlir/python/mlir/execution_engine.py
index 39d9501d9c8bc..f3bcd0e0d78a7 100644
--- a/mlir/python/mlir/execution_engine.py
+++ b/mlir/python/mlir/execution_engine.py
@@ -3,10 +3,15 @@
# 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`
diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py
index e8a6fe1cd2a5f..1e9b75fef8d38 100644
--- a/mlir/test/python/execution_engine.py
+++ b/mlir/test/python/execution_engine.py
@@ -33,7 +33,7 @@ def testCapsule():
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)
@@ -68,7 +68,7 @@ def testInvokeVoid():
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")
@@ -157,7 +157,7 @@ def callback(a):
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(
@@ -168,7 +168,7 @@ def callback(a):
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.]]
More information about the Mlir-commits
mailing list