MLIR  16.0.0git
ExecutionEngineModule.cpp
Go to the documentation of this file.
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 
11 
12 namespace py = pybind11;
13 using namespace mlir;
14 using namespace mlir::python;
15 
16 namespace {
17 
18 /// Owning Wrapper around an ExecutionEngine.
19 class PyExecutionEngine {
20 public:
21  PyExecutionEngine(MlirExecutionEngine executionEngine)
22  : executionEngine(executionEngine) {}
23  PyExecutionEngine(PyExecutionEngine &&other)
24  : executionEngine(other.executionEngine) {
25  other.executionEngine.ptr = nullptr;
26  }
27  ~PyExecutionEngine() {
28  if (!mlirExecutionEngineIsNull(executionEngine))
29  mlirExecutionEngineDestroy(executionEngine);
30  }
31  MlirExecutionEngine get() { return executionEngine; }
32 
33  void release() {
34  executionEngine.ptr = nullptr;
35  referencedObjects.clear();
36  }
37  pybind11::object getCapsule() {
38  return py::reinterpret_steal<py::object>(
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 pybind11::object &obj) {
45  referencedObjects.push_back(obj);
46  }
47 
48  static pybind11::object createFromCapsule(pybind11::object capsule) {
49  MlirExecutionEngine rawPm =
51  if (mlirExecutionEngineIsNull(rawPm))
52  throw py::error_already_set();
53  return py::cast(PyExecutionEngine(rawPm), py::return_value_policy::move);
54  }
55 
56 private:
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<py::object> referencedObjects;
62 };
63 
64 } // namespace
65 
66 /// Create the `mlir.execution_engine` module here.
67 PYBIND11_MODULE(_mlirExecutionEngine, m) {
68  m.doc() = "MLIR Execution Engine";
69 
70  //----------------------------------------------------------------------------
71  // Mapping of the top-level PassManager
72  //----------------------------------------------------------------------------
73  py::class_<PyExecutionEngine>(m, "ExecutionEngine", py::module_local())
74  .def(py::init<>([](MlirModule module, int optLevel,
75  const std::vector<std::string> &sharedLibPaths) {
77  for (const std::string &path : sharedLibPaths)
78  libPaths.push_back({path.c_str(), path.length()});
79  MlirExecutionEngine executionEngine = mlirExecutionEngineCreate(
80  module, optLevel, libPaths.size(), libPaths.data());
81  if (mlirExecutionEngineIsNull(executionEngine))
82  throw std::runtime_error(
83  "Failure while creating the ExecutionEngine.");
84  return new PyExecutionEngine(executionEngine);
85  }),
86  py::arg("module"), py::arg("opt_level") = 2,
87  py::arg("shared_libs") = py::list(),
88  "Create a new ExecutionEngine instance for the given Module. The "
89  "module must contain only dialects that can be translated to LLVM. "
90  "Perform transformations and code generation at the optimization "
91  "level `opt_level` if specified, or otherwise at the default "
92  "level of two (-O2). Load a list of libraries specified in "
93  "`shared_libs`.")
94  .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
95  &PyExecutionEngine::getCapsule)
96  .def("_testing_release", &PyExecutionEngine::release,
97  "Releases (leaks) the backing ExecutionEngine (for testing purpose)")
98  .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyExecutionEngine::createFromCapsule)
99  .def(
100  "raw_lookup",
101  [](PyExecutionEngine &executionEngine, const std::string &func) {
103  executionEngine.get(),
104  mlirStringRefCreate(func.c_str(), func.size()));
105  return reinterpret_cast<uintptr_t>(res);
106  },
107  py::arg("func_name"),
108  "Lookup function `func` in the ExecutionEngine.")
109  .def(
110  "raw_register_runtime",
111  [](PyExecutionEngine &executionEngine, const std::string &name,
112  py::object callbackObj) {
113  executionEngine.addReferencedObject(callbackObj);
114  uintptr_t rawSym =
115  py::cast<uintptr_t>(py::getattr(callbackObj, "value"));
117  executionEngine.get(),
118  mlirStringRefCreate(name.c_str(), name.size()),
119  reinterpret_cast<void *>(rawSym));
120  },
121  py::arg("name"), py::arg("callback"),
122  "Register `callback` as the runtime symbol `name`.")
123  .def(
124  "dump_to_object_file",
125  [](PyExecutionEngine &executionEngine, const std::string &fileName) {
127  executionEngine.get(),
128  mlirStringRefCreate(fileName.c_str(), fileName.size()));
129  },
130  py::arg("file_name"), "Dump ExecutionEngine to an object file.");
131 }
Include the generated interface declarations.
static bool mlirExecutionEngineIsNull(MlirExecutionEngine jit)
Checks whether an execution engine is null.
MLIR_CAPI_EXPORTED void mlirExecutionEngineDumpToObjectFile(MlirExecutionEngine jit, MlirStringRef fileName)
Dump as an object in fileName.
MLIR_CAPI_EXPORTED void * mlirExecutionEngineLookupPacked(MlirExecutionEngine jit, MlirStringRef name)
Lookup the wrapper of the native function in the execution engine with the given name, returns nullptr if the function can&#39;t be looked-up.
PYBIND11_MODULE(_mlirExecutionEngine, m)
Create the mlir.execution_engine module here.
MLIR_CAPI_EXPORTED void mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit, MlirStringRef name, void *sym)
Register a symbol with the jit: this symbol will be accessible to the jitted code.
MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths, const MlirStringRef *sharedLibPaths)
Creates an ExecutionEngine for the provided ModuleOp.
MLIR_CAPI_EXPORTED void mlirExecutionEngineDestroy(MlirExecutionEngine jit)
Destroy an ExecutionEngine instance.
static MlirExecutionEngine mlirPythonCapsuleToExecutionEngine(PyObject *capsule)
Extracts an MlirExecutionEngine from a capsule as produced from mlirPythonIntegerSetToCapsule.
Definition: Interop.h:342
static MlirStringRef mlirStringRefCreate(const char *str, size_t length)
Constructs a string reference from the pointer and length.
Definition: Support.h:80
#define MLIR_PYTHON_CAPI_PTR_ATTR
Attribute on MLIR Python objects that expose their C-API pointer.
Definition: Interop.h:93
static PyObject * mlirPythonExecutionEngineToCapsule(MlirExecutionEngine jit)
Creates a capsule object encapsulating the raw C-API MlirExecutionEngine.
Definition: Interop.h:332
#define MLIR_PYTHON_CAPI_FACTORY_ATTR
Attribute on MLIR Python objects that exposes a factory function for constructing the corresponding P...
Definition: Interop.h:106