106 lines · cpp
1//===- TransformInterpreter.cpp -------------------------------------------===//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// Pybind classes for the transform dialect interpreter.10//11//===----------------------------------------------------------------------===//12 13#include "mlir-c/Dialect/Transform/Interpreter.h"14#include "mlir-c/IR.h"15#include "mlir-c/Support.h"16#include "mlir/Bindings/Python/Diagnostics.h"17#include "mlir/Bindings/Python/NanobindAdaptors.h"18#include "mlir/Bindings/Python/Nanobind.h"19 20namespace nb = nanobind;21 22namespace {23struct PyMlirTransformOptions {24 PyMlirTransformOptions() { options = mlirTransformOptionsCreate(); };25 PyMlirTransformOptions(PyMlirTransformOptions &&other) {26 options = other.options;27 other.options.ptr = nullptr;28 }29 PyMlirTransformOptions(const PyMlirTransformOptions &) = delete;30 31 ~PyMlirTransformOptions() { mlirTransformOptionsDestroy(options); }32 33 MlirTransformOptions options;34};35} // namespace36 37static void populateTransformInterpreterSubmodule(nb::module_ &m) {38 nb::class_<PyMlirTransformOptions>(m, "TransformOptions")39 .def(nb::init<>())40 .def_prop_rw(41 "expensive_checks",42 [](const PyMlirTransformOptions &self) {43 return mlirTransformOptionsGetExpensiveChecksEnabled(self.options);44 },45 [](PyMlirTransformOptions &self, bool value) {46 mlirTransformOptionsEnableExpensiveChecks(self.options, value);47 })48 .def_prop_rw(49 "enforce_single_top_level_transform_op",50 [](const PyMlirTransformOptions &self) {51 return mlirTransformOptionsGetEnforceSingleTopLevelTransformOp(52 self.options);53 },54 [](PyMlirTransformOptions &self, bool value) {55 mlirTransformOptionsEnforceSingleTopLevelTransformOp(self.options,56 value);57 });58 59 m.def(60 "apply_named_sequence",61 [](MlirOperation payloadRoot, MlirOperation transformRoot,62 MlirOperation transformModule, const PyMlirTransformOptions &options) {63 mlir::python::CollectDiagnosticsToStringScope scope(64 mlirOperationGetContext(transformRoot));65 66 // Calling back into Python to invalidate everything under the payload67 // root. This is awkward, but we don't have access to PyMlirContext68 // object here otherwise.69 nb::object obj = nb::cast(payloadRoot);70 71 MlirLogicalResult result = mlirTransformApplyNamedSequence(72 payloadRoot, transformRoot, transformModule, options.options);73 if (mlirLogicalResultIsSuccess(result))74 return;75 76 throw nb::value_error(77 ("Failed to apply named transform sequence.\nDiagnostic message " +78 scope.takeMessage())79 .c_str());80 },81 nb::arg("payload_root"), nb::arg("transform_root"),82 nb::arg("transform_module"),83 nb::arg("transform_options") = PyMlirTransformOptions());84 85 m.def(86 "copy_symbols_and_merge_into",87 [](MlirOperation target, MlirOperation other) {88 mlir::python::CollectDiagnosticsToStringScope scope(89 mlirOperationGetContext(target));90 91 MlirLogicalResult result = mlirMergeSymbolsIntoFromClone(target, other);92 if (mlirLogicalResultIsFailure(result)) {93 throw nb::value_error(94 ("Failed to merge symbols.\nDiagnostic message " +95 scope.takeMessage())96 .c_str());97 }98 },99 nb::arg("target"), nb::arg("other"));100}101 102NB_MODULE(_mlirTransformInterpreter, m) {103 m.doc() = "MLIR Transform dialect interpreter functionality.";104 populateTransformInterpreterSubmodule(m);105}106