258 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8 9#include "Pass.h"10 11#include "Globals.h"12#include "IRModule.h"13#include "mlir-c/Pass.h"14// clang-format off15#include "mlir/Bindings/Python/Nanobind.h"16#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.17// clang-format on18 19namespace nb = nanobind;20using namespace nb::literals;21using namespace mlir;22using namespace mlir::python;23 24namespace {25 26/// Owning Wrapper around a PassManager.27class PyPassManager {28public:29 PyPassManager(MlirPassManager passManager) : passManager(passManager) {}30 PyPassManager(PyPassManager &&other) noexcept31 : passManager(other.passManager) {32 other.passManager.ptr = nullptr;33 }34 ~PyPassManager() {35 if (!mlirPassManagerIsNull(passManager))36 mlirPassManagerDestroy(passManager);37 }38 MlirPassManager get() { return passManager; }39 40 void release() { passManager.ptr = nullptr; }41 nb::object getCapsule() {42 return nb::steal<nb::object>(mlirPythonPassManagerToCapsule(get()));43 }44 45 static nb::object createFromCapsule(const nb::object &capsule) {46 MlirPassManager rawPm = mlirPythonCapsuleToPassManager(capsule.ptr());47 if (mlirPassManagerIsNull(rawPm))48 throw nb::python_error();49 return nb::cast(PyPassManager(rawPm), nb::rv_policy::move);50 }51 52private:53 MlirPassManager passManager;54};55 56} // namespace57 58/// Create the `mlir.passmanager` here.59void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {60 //----------------------------------------------------------------------------61 // Mapping of enumerated types62 //----------------------------------------------------------------------------63 nb::enum_<MlirPassDisplayMode>(m, "PassDisplayMode")64 .value("LIST", MLIR_PASS_DISPLAY_MODE_LIST)65 .value("PIPELINE", MLIR_PASS_DISPLAY_MODE_PIPELINE);66 67 //----------------------------------------------------------------------------68 // Mapping of MlirExternalPass69 //----------------------------------------------------------------------------70 nb::class_<MlirExternalPass>(m, "ExternalPass")71 .def("signal_pass_failure",72 [](MlirExternalPass pass) { mlirExternalPassSignalFailure(pass); });73 74 //----------------------------------------------------------------------------75 // Mapping of the top-level PassManager76 //----------------------------------------------------------------------------77 nb::class_<PyPassManager>(m, "PassManager")78 .def(79 "__init__",80 [](PyPassManager &self, const std::string &anchorOp,81 DefaultingPyMlirContext context) {82 MlirPassManager passManager = mlirPassManagerCreateOnOperation(83 context->get(),84 mlirStringRefCreate(anchorOp.data(), anchorOp.size()));85 new (&self) PyPassManager(passManager);86 },87 "anchor_op"_a = nb::str("any"), "context"_a = nb::none(),88 // clang-format off89 nb::sig("def __init__(self, anchor_op: str = 'any', context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> None"),90 // clang-format on91 "Create a new PassManager for the current (or provided) Context.")92 .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyPassManager::getCapsule)93 .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyPassManager::createFromCapsule)94 .def("_testing_release", &PyPassManager::release,95 "Releases (leaks) the backing pass manager (testing)")96 .def(97 "enable_ir_printing",98 [](PyPassManager &passManager, bool printBeforeAll,99 bool printAfterAll, bool printModuleScope, bool printAfterChange,100 bool printAfterFailure, std::optional<int64_t> largeElementsLimit,101 std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,102 bool printGenericOpForm,103 std::optional<std::string> optionalTreePrintingPath) {104 MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();105 if (largeElementsLimit) {106 mlirOpPrintingFlagsElideLargeElementsAttrs(flags,107 *largeElementsLimit);108 mlirOpPrintingFlagsElideLargeResourceString(flags,109 *largeElementsLimit);110 }111 if (largeResourceLimit)112 mlirOpPrintingFlagsElideLargeResourceString(flags,113 *largeResourceLimit);114 if (enableDebugInfo)115 mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,116 /*prettyForm=*/false);117 if (printGenericOpForm)118 mlirOpPrintingFlagsPrintGenericOpForm(flags);119 std::string treePrintingPath = "";120 if (optionalTreePrintingPath.has_value())121 treePrintingPath = optionalTreePrintingPath.value();122 mlirPassManagerEnableIRPrinting(123 passManager.get(), printBeforeAll, printAfterAll,124 printModuleScope, printAfterChange, printAfterFailure, flags,125 mlirStringRefCreate(treePrintingPath.data(),126 treePrintingPath.size()));127 mlirOpPrintingFlagsDestroy(flags);128 },129 "print_before_all"_a = false, "print_after_all"_a = true,130 "print_module_scope"_a = false, "print_after_change"_a = false,131 "print_after_failure"_a = false,132 "large_elements_limit"_a = nb::none(),133 "large_resource_limit"_a = nb::none(), "enable_debug_info"_a = false,134 "print_generic_op_form"_a = false,135 "tree_printing_dir_path"_a = nb::none(),136 "Enable IR printing, default as mlir-print-ir-after-all.")137 .def(138 "enable_verifier",139 [](PyPassManager &passManager, bool enable) {140 mlirPassManagerEnableVerifier(passManager.get(), enable);141 },142 "enable"_a, "Enable / disable verify-each.")143 .def(144 "enable_timing",145 [](PyPassManager &passManager) {146 mlirPassManagerEnableTiming(passManager.get());147 },148 "Enable pass timing.")149 .def(150 "enable_statistics",151 [](PyPassManager &passManager, MlirPassDisplayMode displayMode) {152 mlirPassManagerEnableStatistics(passManager.get(), displayMode);153 },154 "displayMode"_a =155 MlirPassDisplayMode::MLIR_PASS_DISPLAY_MODE_PIPELINE,156 "Enable pass statistics.")157 .def_static(158 "parse",159 [](const std::string &pipeline, DefaultingPyMlirContext context) {160 MlirPassManager passManager = mlirPassManagerCreate(context->get());161 PyPrintAccumulator errorMsg;162 MlirLogicalResult status = mlirParsePassPipeline(163 mlirPassManagerGetAsOpPassManager(passManager),164 mlirStringRefCreate(pipeline.data(), pipeline.size()),165 errorMsg.getCallback(), errorMsg.getUserData());166 if (mlirLogicalResultIsFailure(status))167 throw nb::value_error(errorMsg.join().c_str());168 return new PyPassManager(passManager);169 },170 "pipeline"_a, "context"_a = nb::none(),171 // clang-format off172 nb::sig("def parse(pipeline: str, context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> PassManager"),173 // clang-format on174 "Parse a textual pass-pipeline and return a top-level PassManager "175 "that can be applied on a Module. Throw a ValueError if the pipeline "176 "can't be parsed")177 .def(178 "add",179 [](PyPassManager &passManager, const std::string &pipeline) {180 PyPrintAccumulator errorMsg;181 MlirLogicalResult status = mlirOpPassManagerAddPipeline(182 mlirPassManagerGetAsOpPassManager(passManager.get()),183 mlirStringRefCreate(pipeline.data(), pipeline.size()),184 errorMsg.getCallback(), errorMsg.getUserData());185 if (mlirLogicalResultIsFailure(status))186 throw nb::value_error(errorMsg.join().c_str());187 },188 "pipeline"_a,189 "Add textual pipeline elements to the pass manager. Throws a "190 "ValueError if the pipeline can't be parsed.")191 .def(192 "add",193 [](PyPassManager &passManager, const nb::callable &run,194 std::optional<std::string> &name, const std::string &argument,195 const std::string &description, const std::string &opName) {196 if (!name.has_value()) {197 name = nb::cast<std::string>(198 nb::borrow<nb::str>(run.attr("__name__")));199 }200 MlirTypeID passID = PyGlobals::get().allocateTypeID();201 MlirExternalPassCallbacks callbacks;202 callbacks.construct = [](void *obj) {203 (void)nb::handle(static_cast<PyObject *>(obj)).inc_ref();204 };205 callbacks.destruct = [](void *obj) {206 (void)nb::handle(static_cast<PyObject *>(obj)).dec_ref();207 };208 callbacks.initialize = nullptr;209 callbacks.clone = [](void *) -> void * {210 throw std::runtime_error("Cloning Python passes not supported");211 };212 callbacks.run = [](MlirOperation op, MlirExternalPass pass,213 void *userData) {214 nb::handle(static_cast<PyObject *>(userData))(op, pass);215 };216 auto externalPass = mlirCreateExternalPass(217 passID, mlirStringRefCreate(name->data(), name->length()),218 mlirStringRefCreate(argument.data(), argument.length()),219 mlirStringRefCreate(description.data(), description.length()),220 mlirStringRefCreate(opName.data(), opName.size()),221 /*nDependentDialects*/ 0, /*dependentDialects*/ nullptr,222 callbacks, /*userData*/ run.ptr());223 mlirPassManagerAddOwnedPass(passManager.get(), externalPass);224 },225 "run"_a, "name"_a.none() = nb::none(), "argument"_a.none() = "",226 "description"_a.none() = "", "op_name"_a.none() = "",227 "Add a python-defined pass to the pass manager.")228 .def(229 "run",230 [](PyPassManager &passManager, PyOperationBase &op) {231 // Actually run the pass manager.232 PyMlirContext::ErrorCapture errors(op.getOperation().getContext());233 MlirLogicalResult status = mlirPassManagerRunOnOp(234 passManager.get(), op.getOperation().get());235 if (mlirLogicalResultIsFailure(status))236 throw MLIRError("Failure while executing pass pipeline",237 errors.take());238 },239 "operation"_a,240 // clang-format off241 nb::sig("def run(self, operation: " MAKE_MLIR_PYTHON_QUALNAME("ir._OperationBase") ") -> None"),242 // clang-format on243 "Run the pass manager on the provided operation, raising an "244 "MLIRError on failure.")245 .def(246 "__str__",247 [](PyPassManager &self) {248 MlirPassManager passManager = self.get();249 PyPrintAccumulator printAccum;250 mlirPrintPassPipeline(251 mlirPassManagerGetAsOpPassManager(passManager),252 printAccum.getCallback(), printAccum.getUserData());253 return printAccum.join();254 },255 "Print the textual representation for this PassManager, suitable to "256 "be passed to `parse` for round-tripping.");257}258