147 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir-c/ExecutionEngine.h"10#include "mlir/Bindings/Python/Nanobind.h"11#include "mlir/Bindings/Python/NanobindAdaptors.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) noexcept25 : 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 exceed43 // those of the ExecutionEngine.44 void addReferencedObject(const nb::object &obj) {45 referencedObjects.push_back(obj);46 }47 48 static nb::object createFromCapsule(const 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 objects59 // so that they don't get garbage collected. (The ExecutionEngine itself60 // just holds raw pointers with no lifetime semantics).61 std::vector<nb::object> referencedObjects;62};63 64} // namespace65 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 PassManager72 //----------------------------------------------------------------------------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 const 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 "initialize",129 [](PyExecutionEngine &executionEngine) {130 mlirExecutionEngineInitialize(executionEngine.get());131 },132 "Initialize the ExecutionEngine. Global constructors specified by "133 "`llvm.mlir.global_ctors` will be run. One common scenario is that "134 "kernel binary compiled from `gpu.module` gets loaded during "135 "initialization. Make sure all symbols are resolvable before "136 "initialization by calling `register_runtime` or including "137 "shared libraries.")138 .def(139 "dump_to_object_file",140 [](PyExecutionEngine &executionEngine, const std::string &fileName) {141 mlirExecutionEngineDumpToObjectFile(142 executionEngine.get(),143 mlirStringRefCreate(fileName.c_str(), fileName.size()));144 },145 nb::arg("file_name"), "Dump ExecutionEngine to an object file.");146}147