442 lines · cpp
1//===- Rewrite.cpp - Rewrite ----------------------------------------------===//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 "Rewrite.h"10 11#include "IRModule.h"12#include "mlir-c/IR.h"13#include "mlir-c/Rewrite.h"14#include "mlir-c/Support.h"15// clang-format off16#include "mlir/Bindings/Python/Nanobind.h"17#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.18// clang-format on19#include "mlir/Config/mlir-config.h"20#include "nanobind/nanobind.h"21 22namespace nb = nanobind;23using namespace mlir;24using namespace nb::literals;25using namespace mlir::python;26 27namespace {28 29class PyPatternRewriter {30public:31 PyPatternRewriter(MlirPatternRewriter rewriter)32 : base(mlirPatternRewriterAsBase(rewriter)),33 ctx(PyMlirContext::forContext(mlirRewriterBaseGetContext(base))) {}34 35 PyInsertionPoint getInsertionPoint() const {36 MlirBlock block = mlirRewriterBaseGetInsertionBlock(base);37 MlirOperation op = mlirRewriterBaseGetOperationAfterInsertion(base);38 39 if (mlirOperationIsNull(op)) {40 MlirOperation owner = mlirBlockGetParentOperation(block);41 auto parent = PyOperation::forOperation(ctx, owner);42 return PyInsertionPoint(PyBlock(parent, block));43 }44 45 return PyInsertionPoint(PyOperation::forOperation(ctx, op));46 }47 48 void replaceOp(MlirOperation op, MlirOperation newOp) {49 mlirRewriterBaseReplaceOpWithOperation(base, op, newOp);50 }51 52 void replaceOp(MlirOperation op, const std::vector<MlirValue> &values) {53 mlirRewriterBaseReplaceOpWithValues(base, op, values.size(), values.data());54 }55 56 void eraseOp(MlirOperation op) { mlirRewriterBaseEraseOp(base, op); }57 58private:59 MlirRewriterBase base;60 PyMlirContextRef ctx;61};62 63#if MLIR_ENABLE_PDL_IN_PATTERNMATCH64static nb::object objectFromPDLValue(MlirPDLValue value) {65 if (MlirValue v = mlirPDLValueAsValue(value); !mlirValueIsNull(v))66 return nb::cast(v);67 if (MlirOperation v = mlirPDLValueAsOperation(value); !mlirOperationIsNull(v))68 return nb::cast(v);69 if (MlirAttribute v = mlirPDLValueAsAttribute(value); !mlirAttributeIsNull(v))70 return nb::cast(v);71 if (MlirType v = mlirPDLValueAsType(value); !mlirTypeIsNull(v))72 return nb::cast(v);73 74 throw std::runtime_error("unsupported PDL value type");75}76 77static std::vector<nb::object> objectsFromPDLValues(size_t nValues,78 MlirPDLValue *values) {79 std::vector<nb::object> args;80 args.reserve(nValues);81 for (size_t i = 0; i < nValues; ++i)82 args.push_back(objectFromPDLValue(values[i]));83 return args;84}85 86// Convert the Python object to a boolean.87// If it evaluates to False, treat it as success;88// otherwise, treat it as failure.89// Note that None is considered success.90static MlirLogicalResult logicalResultFromObject(const nb::object &obj) {91 if (obj.is_none())92 return mlirLogicalResultSuccess();93 94 return nb::cast<bool>(obj) ? mlirLogicalResultFailure()95 : mlirLogicalResultSuccess();96}97 98/// Owning Wrapper around a PDLPatternModule.99class PyPDLPatternModule {100public:101 PyPDLPatternModule(MlirPDLPatternModule module) : module(module) {}102 PyPDLPatternModule(PyPDLPatternModule &&other) noexcept103 : module(other.module) {104 other.module.ptr = nullptr;105 }106 ~PyPDLPatternModule() {107 if (module.ptr != nullptr)108 mlirPDLPatternModuleDestroy(module);109 }110 MlirPDLPatternModule get() { return module; }111 112 void registerRewriteFunction(const std::string &name,113 const nb::callable &fn) {114 mlirPDLPatternModuleRegisterRewriteFunction(115 get(), mlirStringRefCreate(name.data(), name.size()),116 [](MlirPatternRewriter rewriter, MlirPDLResultList results,117 size_t nValues, MlirPDLValue *values,118 void *userData) -> MlirLogicalResult {119 nb::handle f = nb::handle(static_cast<PyObject *>(userData));120 return logicalResultFromObject(121 f(PyPatternRewriter(rewriter), results,122 objectsFromPDLValues(nValues, values)));123 },124 fn.ptr());125 }126 127 void registerConstraintFunction(const std::string &name,128 const nb::callable &fn) {129 mlirPDLPatternModuleRegisterConstraintFunction(130 get(), mlirStringRefCreate(name.data(), name.size()),131 [](MlirPatternRewriter rewriter, MlirPDLResultList results,132 size_t nValues, MlirPDLValue *values,133 void *userData) -> MlirLogicalResult {134 nb::handle f = nb::handle(static_cast<PyObject *>(userData));135 return logicalResultFromObject(136 f(PyPatternRewriter(rewriter), results,137 objectsFromPDLValues(nValues, values)));138 },139 fn.ptr());140 }141 142private:143 MlirPDLPatternModule module;144};145#endif // MLIR_ENABLE_PDL_IN_PATTERNMATCH146 147/// Owning Wrapper around a FrozenRewritePatternSet.148class PyFrozenRewritePatternSet {149public:150 PyFrozenRewritePatternSet(MlirFrozenRewritePatternSet set) : set(set) {}151 PyFrozenRewritePatternSet(PyFrozenRewritePatternSet &&other) noexcept152 : set(other.set) {153 other.set.ptr = nullptr;154 }155 ~PyFrozenRewritePatternSet() {156 if (set.ptr != nullptr)157 mlirFrozenRewritePatternSetDestroy(set);158 }159 MlirFrozenRewritePatternSet get() { return set; }160 161 nb::object getCapsule() {162 return nb::steal<nb::object>(163 mlirPythonFrozenRewritePatternSetToCapsule(get()));164 }165 166 static nb::object createFromCapsule(const nb::object &capsule) {167 MlirFrozenRewritePatternSet rawPm =168 mlirPythonCapsuleToFrozenRewritePatternSet(capsule.ptr());169 if (rawPm.ptr == nullptr)170 throw nb::python_error();171 return nb::cast(PyFrozenRewritePatternSet(rawPm), nb::rv_policy::move);172 }173 174private:175 MlirFrozenRewritePatternSet set;176};177 178class PyRewritePatternSet {179public:180 PyRewritePatternSet(MlirContext ctx)181 : set(mlirRewritePatternSetCreate(ctx)), ctx(ctx) {}182 ~PyRewritePatternSet() {183 if (set.ptr)184 mlirRewritePatternSetDestroy(set);185 }186 187 void add(MlirStringRef rootName, unsigned benefit,188 const nb::callable &matchAndRewrite) {189 MlirRewritePatternCallbacks callbacks;190 callbacks.construct = [](void *userData) {191 nb::handle(static_cast<PyObject *>(userData)).inc_ref();192 };193 callbacks.destruct = [](void *userData) {194 nb::handle(static_cast<PyObject *>(userData)).dec_ref();195 };196 callbacks.matchAndRewrite = [](MlirRewritePattern, MlirOperation op,197 MlirPatternRewriter rewriter,198 void *userData) -> MlirLogicalResult {199 nb::handle f(static_cast<PyObject *>(userData));200 201 PyMlirContextRef ctx =202 PyMlirContext::forContext(mlirOperationGetContext(op));203 nb::object opView = PyOperation::forOperation(ctx, op)->createOpView();204 205 nb::object res = f(opView, PyPatternRewriter(rewriter));206 return logicalResultFromObject(res);207 };208 MlirRewritePattern pattern = mlirOpRewritePatternCreate(209 rootName, benefit, ctx, callbacks, matchAndRewrite.ptr(),210 /* nGeneratedNames */ 0,211 /* generatedNames */ nullptr);212 mlirRewritePatternSetAdd(set, pattern);213 }214 215 PyFrozenRewritePatternSet freeze() {216 MlirRewritePatternSet s = set;217 set.ptr = nullptr;218 return mlirFreezeRewritePattern(s);219 }220 221private:222 MlirRewritePatternSet set;223 MlirContext ctx;224};225 226} // namespace227 228/// Create the `mlir.rewrite` here.229void mlir::python::populateRewriteSubmodule(nb::module_ &m) {230 //----------------------------------------------------------------------------231 // Mapping of the PatternRewriter232 //----------------------------------------------------------------------------233 nb::234 class_<PyPatternRewriter>(m, "PatternRewriter")235 .def_prop_ro("ip", &PyPatternRewriter::getInsertionPoint,236 "The current insertion point of the PatternRewriter.")237 .def(238 "replace_op",239 [](PyPatternRewriter &self, MlirOperation op,240 MlirOperation newOp) { self.replaceOp(op, newOp); },241 "Replace an operation with a new operation.", nb::arg("op"),242 nb::arg("new_op"),243 // clang-format off244 nb::sig("def replace_op(self, op: " MAKE_MLIR_PYTHON_QUALNAME("ir.Operation") ", new_op: " MAKE_MLIR_PYTHON_QUALNAME("ir.Operation") ") -> None")245 // clang-format on246 )247 .def(248 "replace_op",249 [](PyPatternRewriter &self, MlirOperation op,250 const std::vector<MlirValue> &values) {251 self.replaceOp(op, values);252 },253 "Replace an operation with a list of values.", nb::arg("op"),254 nb::arg("values"),255 // clang-format off256 nb::sig("def replace_op(self, op: " MAKE_MLIR_PYTHON_QUALNAME("ir.Operation") ", values: list[" MAKE_MLIR_PYTHON_QUALNAME("ir.Value") "]) -> None")257 // clang-format on258 )259 .def("erase_op", &PyPatternRewriter::eraseOp, "Erase an operation.",260 nb::arg("op"),261 // clang-format off262 nb::sig("def erase_op(self, op: " MAKE_MLIR_PYTHON_QUALNAME("ir.Operation") ") -> None")263 // clang-format on264 );265 266 //----------------------------------------------------------------------------267 // Mapping of the RewritePatternSet268 //----------------------------------------------------------------------------269 nb::class_<PyRewritePatternSet>(m, "RewritePatternSet")270 .def(271 "__init__",272 [](PyRewritePatternSet &self, DefaultingPyMlirContext context) {273 new (&self) PyRewritePatternSet(context.get()->get());274 },275 "context"_a = nb::none())276 .def(277 "add",278 [](PyRewritePatternSet &self, nb::handle root, const nb::callable &fn,279 unsigned benefit) {280 std::string opName =281 nb::cast<std::string>(root.attr("OPERATION_NAME"));282 self.add(mlirStringRefCreate(opName.data(), opName.size()), benefit,283 fn);284 },285 "root"_a, "fn"_a, "benefit"_a = 1,286 "Add a new rewrite pattern on the given root operation with the "287 "callable as the matching and rewriting function and the given "288 "benefit.")289 .def("freeze", &PyRewritePatternSet::freeze,290 "Freeze the pattern set into a frozen one.");291 292 //----------------------------------------------------------------------------293 // Mapping of the PDLResultList and PDLModule294 //----------------------------------------------------------------------------295#if MLIR_ENABLE_PDL_IN_PATTERNMATCH296 nb::class_<MlirPDLResultList>(m, "PDLResultList")297 .def(298 "append",299 [](MlirPDLResultList results, const PyValue &value) {300 mlirPDLResultListPushBackValue(results, value);301 },302 // clang-format off303 nb::sig("def append(self, value: " MAKE_MLIR_PYTHON_QUALNAME("ir.Value") ")")304 // clang-format on305 )306 .def(307 "append",308 [](MlirPDLResultList results, const PyOperation &op) {309 mlirPDLResultListPushBackOperation(results, op);310 },311 // clang-format off312 nb::sig("def append(self, op: " MAKE_MLIR_PYTHON_QUALNAME("ir.Operation") ")")313 // clang-format on314 )315 .def(316 "append",317 [](MlirPDLResultList results, const PyType &type) {318 mlirPDLResultListPushBackType(results, type);319 },320 // clang-format off321 nb::sig("def append(self, type: " MAKE_MLIR_PYTHON_QUALNAME("ir.Type") ")")322 // clang-format on323 )324 .def(325 "append",326 [](MlirPDLResultList results, const PyAttribute &attr) {327 mlirPDLResultListPushBackAttribute(results, attr);328 },329 // clang-format off330 nb::sig("def append(self, attr: " MAKE_MLIR_PYTHON_QUALNAME("ir.Attribute") ")")331 // clang-format on332 );333 nb::class_<PyPDLPatternModule>(m, "PDLModule")334 .def(335 "__init__",336 [](PyPDLPatternModule &self, MlirModule module) {337 new (&self)338 PyPDLPatternModule(mlirPDLPatternModuleFromModule(module));339 },340 // clang-format off341 nb::sig("def __init__(self, module: " MAKE_MLIR_PYTHON_QUALNAME("ir.Module") ") -> None"),342 // clang-format on343 "module"_a, "Create a PDL module from the given module.")344 .def(345 "__init__",346 [](PyPDLPatternModule &self, PyModule &module) {347 new (&self) PyPDLPatternModule(348 mlirPDLPatternModuleFromModule(module.get()));349 },350 // clang-format off351 nb::sig("def __init__(self, module: " MAKE_MLIR_PYTHON_QUALNAME("ir.Module") ") -> None"),352 // clang-format on353 "module"_a, "Create a PDL module from the given module.")354 .def(355 "freeze",356 [](PyPDLPatternModule &self) {357 return PyFrozenRewritePatternSet(mlirFreezeRewritePattern(358 mlirRewritePatternSetFromPDLPatternModule(self.get())));359 },360 nb::keep_alive<0, 1>())361 .def(362 "register_rewrite_function",363 [](PyPDLPatternModule &self, const std::string &name,364 const nb::callable &fn) {365 self.registerRewriteFunction(name, fn);366 },367 nb::keep_alive<1, 3>())368 .def(369 "register_constraint_function",370 [](PyPDLPatternModule &self, const std::string &name,371 const nb::callable &fn) {372 self.registerConstraintFunction(name, fn);373 },374 nb::keep_alive<1, 3>());375#endif // MLIR_ENABLE_PDL_IN_PATTERNMATCH376 nb::class_<PyFrozenRewritePatternSet>(m, "FrozenRewritePatternSet")377 .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR,378 &PyFrozenRewritePatternSet::getCapsule)379 .def(MLIR_PYTHON_CAPI_FACTORY_ATTR,380 &PyFrozenRewritePatternSet::createFromCapsule);381 m.def(382 "apply_patterns_and_fold_greedily",383 [](PyModule &module, PyFrozenRewritePatternSet &set) {384 auto status =385 mlirApplyPatternsAndFoldGreedily(module.get(), set.get(), {});386 if (mlirLogicalResultIsFailure(status))387 throw std::runtime_error("pattern application failed to converge");388 },389 "module"_a, "set"_a,390 // clang-format off391 nb::sig("def apply_patterns_and_fold_greedily(module: " MAKE_MLIR_PYTHON_QUALNAME("ir.Module") ", set: FrozenRewritePatternSet) -> None"),392 // clang-format on393 "Applys the given patterns to the given module greedily while folding "394 "results.")395 .def(396 "apply_patterns_and_fold_greedily",397 [](PyModule &module, MlirFrozenRewritePatternSet set) {398 auto status =399 mlirApplyPatternsAndFoldGreedily(module.get(), set, {});400 if (mlirLogicalResultIsFailure(status))401 throw std::runtime_error(402 "pattern application failed to converge");403 },404 "module"_a, "set"_a,405 // clang-format off406 nb::sig("def apply_patterns_and_fold_greedily(module: " MAKE_MLIR_PYTHON_QUALNAME("ir.Module") ", set: FrozenRewritePatternSet) -> None"),407 // clang-format on408 "Applys the given patterns to the given module greedily while "409 "folding "410 "results.")411 .def(412 "apply_patterns_and_fold_greedily",413 [](PyOperationBase &op, PyFrozenRewritePatternSet &set) {414 auto status = mlirApplyPatternsAndFoldGreedilyWithOp(415 op.getOperation(), set.get(), {});416 if (mlirLogicalResultIsFailure(status))417 throw std::runtime_error(418 "pattern application failed to converge");419 },420 "op"_a, "set"_a,421 // clang-format off422 nb::sig("def apply_patterns_and_fold_greedily(op: " MAKE_MLIR_PYTHON_QUALNAME("ir._OperationBase") ", set: FrozenRewritePatternSet) -> None"),423 // clang-format on424 "Applys the given patterns to the given op greedily while folding "425 "results.")426 .def(427 "apply_patterns_and_fold_greedily",428 [](PyOperationBase &op, MlirFrozenRewritePatternSet set) {429 auto status = mlirApplyPatternsAndFoldGreedilyWithOp(430 op.getOperation(), set, {});431 if (mlirLogicalResultIsFailure(status))432 throw std::runtime_error(433 "pattern application failed to converge");434 },435 "op"_a, "set"_a,436 // clang-format off437 nb::sig("def apply_patterns_and_fold_greedily(op: " MAKE_MLIR_PYTHON_QUALNAME("ir._OperationBase") ", set: FrozenRewritePatternSet) -> None"),438 // clang-format on439 "Applys the given patterns to the given op greedily while folding "440 "results.");441}442