[Mlir-commits] [mlir] [mlir][CAPI, python bindings] Expose `Operation::setSuccessor` (PR #67922)
Maksim Levental
llvmlistbot at llvm.org
Mon Oct 2 09:35:55 PDT 2023
https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/67922
>From 2807bd399cfe910814ce26b18b652a98ac8eff4d Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Sun, 1 Oct 2023 11:01:37 -0500
Subject: [PATCH] [mlir][CAPI, python bindings]
Expose `Operation::setSuccessor`.
---
mlir/include/mlir-c/IR.h | 4 ++
mlir/lib/Bindings/Python/IRCore.cpp | 82 +++++++++++++++++++++++++----
mlir/lib/CAPI/IR/IR.cpp | 5 ++
mlir/test/python/dialects/cf.py | 50 ++++++++++++++++++
4 files changed, 130 insertions(+), 11 deletions(-)
create mode 100644 mlir/test/python/dialects/cf.py
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index a6408317db69e61..e361f33a0d83641 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -576,6 +576,10 @@ MLIR_CAPI_EXPORTED intptr_t mlirOperationGetNumSuccessors(MlirOperation op);
MLIR_CAPI_EXPORTED MlirBlock mlirOperationGetSuccessor(MlirOperation op,
intptr_t pos);
+/// Set `pos`-th successor of the operation.
+MLIR_CAPI_EXPORTED void
+mlirOperationSetSuccessor(MlirOperation op, intptr_t pos, MlirBlock block);
+
/// Returns true if this operation defines an inherent attribute with this name.
/// Note: the attribute can be optional, so
/// `mlirOperationGetInherentAttributeByName` can still return a null attribute.
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index aad74f511e7ef11..84f980d79beebf6 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2207,9 +2207,9 @@ class PyBlockArgumentList
};
/// A list of operation operands. Internally, these are stored as consecutive
-/// elements, random access is cheap. The result list is associated with the
-/// operation whose results these are, and extends the lifetime of this
-/// operation.
+/// elements, random access is cheap. The (returned) operand list is associated
+/// with the operation whose operands these are, and thus extends the lifetime
+/// of this operation.
class PyOpOperandList : public Sliceable<PyOpOperandList, PyValue> {
public:
static constexpr const char *pyClassName = "OpOperandList";
@@ -2262,9 +2262,9 @@ class PyOpOperandList : public Sliceable<PyOpOperandList, PyValue> {
};
/// A list of operation results. Internally, these are stored as consecutive
-/// elements, random access is cheap. The result list is associated with the
-/// operation whose results these are, and extends the lifetime of this
-/// operation.
+/// elements, random access is cheap. The (returned) result list is associated
+/// with the operation whose results these are, and thus extends the lifetime of
+/// this operation.
class PyOpResultList : public Sliceable<PyOpResultList, PyOpResult> {
public:
static constexpr const char *pyClassName = "OpResultList";
@@ -2307,6 +2307,52 @@ class PyOpResultList : public Sliceable<PyOpResultList, PyOpResult> {
PyOperationRef operation;
};
+/// A list of operation successors. Internally, these are stored as consecutive
+/// elements, random access is cheap. The (returned) successor list is
+/// associated with the operation whose successors these are, and thus extends
+/// the lifetime of this operation.
+class PyOpSuccessors : public Sliceable<PyOpSuccessors, PyBlock> {
+public:
+ static constexpr const char *pyClassName = "OpSuccessors";
+
+ PyOpSuccessors(PyOperationRef operation, intptr_t startIndex = 0,
+ intptr_t length = -1, intptr_t step = 1)
+ : Sliceable(startIndex,
+ length == -1 ? mlirOperationGetNumSuccessors(operation->get())
+ : length,
+ step),
+ operation(operation) {}
+
+ void dunderSetItem(intptr_t index, PyBlock block) {
+ index = wrapIndex(index);
+ mlirOperationSetSuccessor(operation->get(), index, block.get());
+ }
+
+ static void bindDerived(ClassTy &c) {
+ c.def("__setitem__", &PyOpSuccessors::dunderSetItem);
+ }
+
+private:
+ /// Give the parent CRTP class access to hook implementations below.
+ friend class Sliceable<PyOpSuccessors, PyBlock>;
+
+ intptr_t getRawNumElements() {
+ operation->checkValid();
+ return mlirOperationGetNumSuccessors(operation->get());
+ }
+
+ PyBlock getRawElement(intptr_t pos) {
+ MlirBlock block = mlirOperationGetSuccessor(operation->get(), pos);
+ return PyBlock(operation, block);
+ }
+
+ PyOpSuccessors slice(intptr_t startIndex, intptr_t length, intptr_t step) {
+ return PyOpSuccessors(operation, startIndex, length, step);
+ }
+
+ PyOperationRef operation;
+};
+
/// A list of operation attributes. Can be indexed by name, producing
/// attributes, or by index, producing named attributes.
class PyOpAttributeMap {
@@ -2924,16 +2970,28 @@ void mlir::python::populateIRCore(py::module &m) {
&PyOperation::getCapsule)
.def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyOperation::createFromCapsule)
.def_property_readonly("operation", [](py::object self) { return self; })
- .def_property_readonly("opview", &PyOperation::createOpView);
+ .def_property_readonly("opview", &PyOperation::createOpView)
+ .def_property_readonly(
+ "successors",
+ [](PyOperationBase &self) {
+ return PyOpSuccessors(self.getOperation().getRef());
+ },
+ "Returns the list of Operation successors.");
auto opViewClass =
py::class_<PyOpView, PyOperationBase>(m, "OpView", py::module_local())
.def(py::init<py::object>(), py::arg("operation"))
.def_property_readonly("operation", &PyOpView::getOperationObject)
.def_property_readonly("opview", [](py::object self) { return self; })
- .def("__str__", [](PyOpView &self) {
- return py::str(self.getOperationObject());
- });
+ .def(
+ "__str__",
+ [](PyOpView &self) { return py::str(self.getOperationObject()); })
+ .def_property_readonly(
+ "successors",
+ [](PyOperationBase &self) {
+ return PyOpSuccessors(self.getOperation().getRef());
+ },
+ "Returns the list of Operation successors.");
opViewClass.attr("_ODS_REGIONS") = py::make_tuple(0, true);
opViewClass.attr("_ODS_OPERAND_SEGMENTS") = py::none();
opViewClass.attr("_ODS_RESULT_SEGMENTS") = py::none();
@@ -3448,7 +3506,8 @@ void mlir::python::populateIRCore(py::module &m) {
mlirOpPrintingFlagsUseLocalScope(flags);
valueState = mlirAsmStateCreateForValue(self.get(), flags);
}
- mlirValuePrintAsOperand(self.get(), valueState, printAccum.getCallback(),
+ mlirValuePrintAsOperand(self.get(), valueState,
+ printAccum.getCallback(),
printAccum.getUserData());
// Release state if allocated locally.
if (!state) {
@@ -3523,6 +3582,7 @@ void mlir::python::populateIRCore(py::module &m) {
PyOpOperandIterator::bind(m);
PyOpOperandList::bind(m);
PyOpResultList::bind(m);
+ PyOpSuccessors::bind(m);
PyRegionIterator::bind(m);
PyRegionList::bind(m);
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index 65b2b7466fb7c39..c1abbbe364611af 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -637,6 +637,11 @@ bool mlirOperationRemoveDiscardableAttributeByName(MlirOperation op,
return !!unwrap(op)->removeDiscardableAttr(unwrap(name));
}
+void mlirOperationSetSuccessor(MlirOperation op, intptr_t pos,
+ MlirBlock block) {
+ unwrap(op)->setSuccessor(unwrap(block), static_cast<unsigned>(pos));
+}
+
intptr_t mlirOperationGetNumAttributes(MlirOperation op) {
return static_cast<intptr_t>(unwrap(op)->getAttrs().size());
}
diff --git a/mlir/test/python/dialects/cf.py b/mlir/test/python/dialects/cf.py
new file mode 100644
index 000000000000000..469e74d44a7dae7
--- /dev/null
+++ b/mlir/test/python/dialects/cf.py
@@ -0,0 +1,50 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import cf
+
+
+def constructAndPrintInModule(f):
+ print("\nTEST:", f.__name__)
+ with Context() as ctx, Location.unknown():
+ ctx.allow_unregistered_dialects = True
+ module = Module.create()
+ with InsertionPoint(module.body):
+ f()
+ return f
+
+
+# CHECK-LABEL: TEST: testBranchAndSetSuccessor
+ at constructAndPrintInModule
+def testBranchAndSetSuccessor():
+ op1 = Operation.create("custom.op1", regions=1)
+
+ block0 = op1.regions[0].blocks.append()
+ ip = InsertionPoint(block0)
+ Operation.create("custom.terminator", ip=ip)
+
+ block1 = op1.regions[0].blocks.append()
+ ip = InsertionPoint(block1)
+ br1 = cf.BranchOp([], block1, ip=ip)
+ # CHECK: ^bb1: // pred: ^bb1
+ # CHECK: cf.br ^bb1
+ print(br1.successors[0])
+ # CHECK: num_successors 1
+ print("num_successors", len(br1.successors))
+
+ block2 = op1.regions[0].blocks.append()
+ ip = InsertionPoint(block2)
+ br2 = cf.BranchOp([], block1, ip=ip)
+ # CHECK: ^bb1: // 2 preds: ^bb1, ^bb2
+ # CHECK: cf.br ^bb1
+ print(br2.successors[0])
+ # CHECK: num_successors 1
+ print("num_successors", len(br2.successors))
+
+ br1.successors[0] = block2
+ # CHECK: ^bb2: // pred: ^bb1
+ # CHECK: cf.br ^bb1
+ print(br1.successors[0])
+ # CHECK: ^bb1: // pred: ^bb2
+ # CHECK: cf.br ^bb2
+ print(br2.operation.successors[0])
More information about the Mlir-commits
mailing list