xref: /llvm-project/mlir/lib/Bindings/Python/Pass.cpp (revision 436c6c9c20cc522c92a923440a5fc509c342a7db)
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 "IRModule.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:
24   PyPassManager(MlirPassManager passManager) : passManager(passManager) {}
25   PyPassManager(PyPassManager &&other) : passManager(other.passManager) {
26     other.passManager.ptr = nullptr;
27   }
28   ~PyPassManager() {
29     if (!mlirPassManagerIsNull(passManager))
30       mlirPassManagerDestroy(passManager);
31   }
32   MlirPassManager get() { return passManager; }
33 
34   void release() { passManager.ptr = nullptr; }
35   pybind11::object getCapsule() {
36     return py::reinterpret_steal<py::object>(
37         mlirPythonPassManagerToCapsule(get()));
38   }
39 
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.
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