1 //===- Pass.cpp - Pass Management -----------------------------------------===//
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 "Pass.h"
10
11 #include "IRModules.h"
12 #include "mlir-c/Bindings/Python/Interop.h"
13 #include "mlir-c/Pass.h"
14
15 namespace py = pybind11;
16 using namespace mlir;
17 using namespace mlir::python;
18
19 namespace {
20
21 /// Owning Wrapper around a PassManager.
22 class PyPassManager {
23 public:
PyPassManager(MlirPassManager passManager)24 PyPassManager(MlirPassManager passManager) : passManager(passManager) {}
PyPassManager(PyPassManager && other)25 PyPassManager(PyPassManager &&other) : passManager(other.passManager) {
26 other.passManager.ptr = nullptr;
27 }
~PyPassManager()28 ~PyPassManager() {
29 if (!mlirPassManagerIsNull(passManager))
30 mlirPassManagerDestroy(passManager);
31 }
get()32 MlirPassManager get() { return passManager; }
33
release()34 void release() { passManager.ptr = nullptr; }
getCapsule()35 pybind11::object getCapsule() {
36 return py::reinterpret_steal<py::object>(
37 mlirPythonPassManagerToCapsule(get()));
38 }
39
createFromCapsule(pybind11::object capsule)40 static pybind11::object createFromCapsule(pybind11::object capsule) {
41 MlirPassManager rawPm = mlirPythonCapsuleToPassManager(capsule.ptr());
42 if (mlirPassManagerIsNull(rawPm))
43 throw py::error_already_set();
44 return py::cast(PyPassManager(rawPm), py::return_value_policy::move);
45 }
46
47 private:
48 MlirPassManager passManager;
49 };
50
51 } // anonymous namespace
52
53 /// Create the `mlir.passmanager` here.
populatePassManagerSubmodule(py::module & m)54 void mlir::python::populatePassManagerSubmodule(py::module &m) {
55 //----------------------------------------------------------------------------
56 // Mapping of the top-level PassManager
57 //----------------------------------------------------------------------------
58 py::class_<PyPassManager>(m, "PassManager")
59 .def(py::init<>([](DefaultingPyMlirContext context) {
60 MlirPassManager passManager =
61 mlirPassManagerCreate(context->get());
62 return new PyPassManager(passManager);
63 }),
64 py::arg("context") = py::none(),
65 "Create a new PassManager for the current (or provided) Context.")
66 .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
67 &PyPassManager::getCapsule)
68 .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyPassManager::createFromCapsule)
69 .def("_testing_release", &PyPassManager::release,
70 "Releases (leaks) the backing pass manager (testing)")
71 .def_static(
72 "parse",
73 [](const std::string pipeline, DefaultingPyMlirContext context) {
74 MlirPassManager passManager = mlirPassManagerCreate(context->get());
75 MlirLogicalResult status = mlirParsePassPipeline(
76 mlirPassManagerGetAsOpPassManager(passManager),
77 mlirStringRefCreate(pipeline.data(), pipeline.size()));
78 if (mlirLogicalResultIsFailure(status))
79 throw SetPyError(PyExc_ValueError,
80 llvm::Twine("invalid pass pipeline '") +
81 pipeline + "'.");
82 return new PyPassManager(passManager);
83 },
84 py::arg("pipeline"), py::arg("context") = py::none(),
85 "Parse a textual pass-pipeline and return a top-level PassManager "
86 "that can be applied on a Module. Throw a ValueError if the pipeline "
87 "can't be parsed")
88 .def(
89 "run",
90 [](PyPassManager &passManager, PyModule &module) {
91 MlirLogicalResult status =
92 mlirPassManagerRun(passManager.get(), module.get());
93 if (mlirLogicalResultIsFailure(status))
94 throw SetPyError(PyExc_RuntimeError,
95 "Failure while executing pass pipeline.");
96 },
97 "Run the pass manager on the provided module, throw a RuntimeError "
98 "on failure.")
99 .def(
100 "__str__",
101 [](PyPassManager &self) {
102 MlirPassManager passManager = self.get();
103 PyPrintAccumulator printAccum;
104 mlirPrintPassPipeline(
105 mlirPassManagerGetAsOpPassManager(passManager),
106 printAccum.getCallback(), printAccum.getUserData());
107 return printAccum.join();
108 },
109 "Print the textual representation for this PassManager, suitable to "
110 "be passed to `parse` for round-tripping.");
111 }
112