a1d7cda1d7
Provides the infrastructure for implementing and late-binding
OpInterfaces from Python.
* On the mlir-c API declaration side, each `XOpInterface` has a callback
struct, with a callback for each method and a userdata member (provided
as an arg to each method), and a
`mlirXOpInterfaceAttachFallbackModel(ctx, op_name, callbacks)` func.
* This CAPI is implemented by defining a subclass of
`XOpInterface::FallbackModel` that holds the callback struct and has
each method call the corresponding callback (with userdata as an arg).
Given a callback struct, a new `FallbackModel` is created and attached,
i.e. late bound, to the named op. (MLIR's interface infrastructure is
such that the thus registered `FallbackModel` will be returned in case
the op gets cast to the `XOpInterface`.)
* On the Python side, we expose a stand-in `XOpInterface` base class
which has one (class)method: `XOpInterface.attach(cls, op_name, ctx)`.
Python users subclass this class (`class MyInterfaceImpl(XOpInterface):
...`) and implement the interface's methods (with the right names and
signatures). The user calls `attach` on the subclass
(`MyInterfaceImpl.attach("my_dialect.my_op", ctx)`) which prepares the
callbacks struct _with userdata set to the subclass_ (as we use it to
lookup methods). These callbacks (and userdata) are then registered as
an `XOpInterface::FallbackModel` by
`mlirXOpInterfaceAttachFallbackModel(...)`. From then on the Python
methods will be used to respond to calls to the interface methods
(originating in C++).
This PR enables implementing the TransformOpInterface and the
MemoryEffectsOpInterface, both of which are required for making an op
into a transform op.
Everything besides the above linked code is there to facilitate exposing
the interfaces: the right types for the arguments of the methods are
exposed as are functions/methods for manipulating these arguments (e.g.
specifying side effects on `OpOperand`s and `OpResult`s and being able
to access and set the transform handles associated with args and
results).
84 lines
3.0 KiB
C++
84 lines
3.0 KiB
C++
//===- Rewrite.h - Rewrite Submodules of pybind module --------------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_BINDINGS_PYTHON_REWRITE_H
|
|
#define MLIR_BINDINGS_PYTHON_REWRITE_H
|
|
|
|
#include "mlir-c/Rewrite.h"
|
|
#include "mlir/Bindings/Python/IRCore.h"
|
|
|
|
#include <nanobind/nanobind.h>
|
|
|
|
namespace mlir {
|
|
namespace python {
|
|
namespace MLIR_BINDINGS_PYTHON_DOMAIN {
|
|
|
|
/// CRTP Base class for rewriter wrappers.
|
|
template <typename DerivedTy>
|
|
class MLIR_PYTHON_API_EXPORTED PyRewriterBase {
|
|
public:
|
|
PyRewriterBase(MlirRewriterBase rewriter)
|
|
: base(rewriter),
|
|
ctx(PyMlirContext::forContext(mlirRewriterBaseGetContext(base))) {}
|
|
|
|
PyInsertionPoint getInsertionPoint() const {
|
|
MlirBlock block = mlirRewriterBaseGetInsertionBlock(base);
|
|
MlirOperation op = mlirRewriterBaseGetOperationAfterInsertion(base);
|
|
|
|
if (mlirOperationIsNull(op)) {
|
|
MlirOperation owner = mlirBlockGetParentOperation(block);
|
|
auto parent = PyOperation::forOperation(ctx, owner);
|
|
return PyInsertionPoint(PyBlock(parent, block));
|
|
}
|
|
|
|
return PyInsertionPoint(PyOperation::forOperation(ctx, op));
|
|
}
|
|
|
|
static void bind(nanobind::module_ &m) {
|
|
nanobind::class_<DerivedTy>(m, DerivedTy::pyClassName)
|
|
.def_prop_ro("ip", &PyRewriterBase::getInsertionPoint,
|
|
"The current insertion point of the PatternRewriter.")
|
|
.def(
|
|
"replace_op",
|
|
[](DerivedTy &self, PyOperationBase &op, PyOperationBase &newOp) {
|
|
mlirRewriterBaseReplaceOpWithOperation(
|
|
self.base, op.getOperation(), newOp.getOperation());
|
|
},
|
|
"Replace an operation with a new operation.", nanobind::arg("op"),
|
|
nanobind::arg("new_op"))
|
|
.def(
|
|
"replace_op",
|
|
[](DerivedTy &self, PyOperationBase &op,
|
|
const std::vector<PyValue> &values) {
|
|
std::vector<MlirValue> values_(values.size());
|
|
std::copy(values.begin(), values.end(), values_.begin());
|
|
mlirRewriterBaseReplaceOpWithValues(
|
|
self.base, op.getOperation(), values_.size(), values_.data());
|
|
},
|
|
"Replace an operation with a list of values.", nanobind::arg("op"),
|
|
nanobind::arg("values"))
|
|
.def(
|
|
"erase_op",
|
|
[](DerivedTy &self, PyOperationBase &op) {
|
|
mlirRewriterBaseEraseOp(self.base, op.getOperation());
|
|
},
|
|
"Erase an operation.", nanobind::arg("op"));
|
|
}
|
|
|
|
private:
|
|
MlirRewriterBase base;
|
|
PyMlirContextRef ctx;
|
|
};
|
|
|
|
void MLIR_PYTHON_API_EXPORTED populateRewriteSubmodule(nanobind::module_ &m);
|
|
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
|
|
} // namespace python
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_BINDINGS_PYTHON_REWRITE_H
|