1//===- ExecutionEngineModule.cpp - Python module for execution engine -----===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#include "mlir-c/ExecutionEngine.h"
10#include "mlir/Bindings/Python/NanobindAdaptors.h"
11#include "mlir/Bindings/Python/Nanobind.h"
12
13namespace nb = nanobind;
14using namespace mlir;
15using namespace mlir::python;
16
17namespace {
18
19/// Owning Wrapper around an ExecutionEngine.
20class PyExecutionEngine {
21public:
22 PyExecutionEngine(MlirExecutionEngine executionEngine)
23 : executionEngine(executionEngine) {}
24 PyExecutionEngine(PyExecutionEngine &&other) noexcept
25 : executionEngine(other.executionEngine) {
26 other.executionEngine.ptr = nullptr;
27 }
28 ~PyExecutionEngine() {
29 if (!mlirExecutionEngineIsNull(executionEngine))
30 mlirExecutionEngineDestroy(executionEngine);
31 }
32 MlirExecutionEngine get() { return executionEngine; }
33
34 void release() {
35 executionEngine.ptr = nullptr;
36 referencedObjects.clear();
37 }
38 nb::object getCapsule() {
39 return nb::steal<nb::object>(mlirPythonExecutionEngineToCapsule(get()));
40 }
41
42 // Add an object to the list of referenced objects whose lifetime must exceed
43 // those of the ExecutionEngine.
44 void addReferencedObject(const nb::object &obj) {
45 referencedObjects.push_back(obj);
46 }
47
48 static nb::object createFromCapsule(nb::object capsule) {
49 MlirExecutionEngine rawPm =
50 mlirPythonCapsuleToExecutionEngine(capsule.ptr());
51 if (mlirExecutionEngineIsNull(rawPm))
52 throw nb::python_error();
53 return nb::cast(PyExecutionEngine(rawPm), nb::rv_policy::move);
54 }
55
56private:
57 MlirExecutionEngine executionEngine;
58 // We support Python ctypes closures as callbacks. Keep a list of the objects
59 // so that they don't get garbage collected. (The ExecutionEngine itself
60 // just holds raw pointers with no lifetime semantics).
61 std::vector<nb::object> referencedObjects;
62};
63
64} // namespace
65
66/// Create the `mlir.execution_engine` module here.
67NB_MODULE(_mlirExecutionEngine, m) {
68 m.doc() = "MLIR Execution Engine";
69
70 //----------------------------------------------------------------------------
71 // Mapping of the top-level PassManager
72 //----------------------------------------------------------------------------
73 nb::class_<PyExecutionEngine>(m, "ExecutionEngine")
74 .def(
75 "__init__",
76 [](PyExecutionEngine &self, MlirModule module, int optLevel,
77 const std::vector<std::string> &sharedLibPaths,
78 bool enableObjectDump) {
79 llvm::SmallVector<MlirStringRef, 4> libPaths;
80 for (const std::string &path : sharedLibPaths)
81 libPaths.push_back({path.c_str(), path.length()});
82 MlirExecutionEngine executionEngine =
83 mlirExecutionEngineCreate(module, optLevel, libPaths.size(),
84 libPaths.data(), enableObjectDump);
85 if (mlirExecutionEngineIsNull(executionEngine))
86 throw std::runtime_error(
87 "Failure while creating the ExecutionEngine.");
88 new (&self) PyExecutionEngine(executionEngine);
89 },
90 nb::arg("module"), nb::arg("opt_level") = 2,
91 nb::arg("shared_libs") = nb::list(),
92 nb::arg("enable_object_dump") = true,
93 "Create a new ExecutionEngine instance for the given Module. The "
94 "module must contain only dialects that can be translated to LLVM. "
95 "Perform transformations and code generation at the optimization "
96 "level `opt_level` if specified, or otherwise at the default "
97 "level of two (-O2). Load a list of libraries specified in "
98 "`shared_libs`.")
99 .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyExecutionEngine::getCapsule)
100 .def("_testing_release", &PyExecutionEngine::release,
101 "Releases (leaks) the backing ExecutionEngine (for testing purpose)")
102 .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyExecutionEngine::createFromCapsule)
103 .def(
104 "raw_lookup",
105 [](PyExecutionEngine &executionEngine, const std::string &func) {
106 auto *res = mlirExecutionEngineLookupPacked(
107 executionEngine.get(),
108 mlirStringRefCreate(func.c_str(), func.size()));
109 return reinterpret_cast<uintptr_t>(res);
110 },
111 nb::arg("func_name"),
112 "Lookup function `func` in the ExecutionEngine.")
113 .def(
114 "raw_register_runtime",
115 [](PyExecutionEngine &executionEngine, const std::string &name,
116 nb::object callbackObj) {
117 executionEngine.addReferencedObject(callbackObj);
118 uintptr_t rawSym =
119 nb::cast<uintptr_t>(nb::getattr(callbackObj, "value"));
120 mlirExecutionEngineRegisterSymbol(
121 executionEngine.get(),
122 mlirStringRefCreate(name.c_str(), name.size()),
123 reinterpret_cast<void *>(rawSym));
124 },
125 nb::arg("name"), nb::arg("callback"),
126 "Register `callback` as the runtime symbol `name`.")
127 .def(
128 "dump_to_object_file",
129 [](PyExecutionEngine &executionEngine, const std::string &fileName) {
130 mlirExecutionEngineDumpToObjectFile(
131 executionEngine.get(),
132 mlirStringRefCreate(fileName.c_str(), fileName.size()));
133 },
134 nb::arg("file_name"), "Dump ExecutionEngine to an object file.");
135}
136

source code of mlir/lib/Bindings/Python/ExecutionEngineModule.cpp