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