[Mlir-commits] [mlir] [mlir][CAPI, python bindings] (PR #67922)
Maksim Levental
llvmlistbot at llvm.org
Sun Oct 1 09:02:59 PDT 2023
https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/67922
Expose `Operation::setSuccessor`.
>From 3b4bfcee2b9db723253d4ab55e1f1d1d4abf61bb 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 | 15 +++++++++--
mlir/lib/CAPI/IR/IR.cpp | 5 ++++
mlir/test/python/dialects/cf.py | 42 +++++++++++++++++++++++++++++
4 files changed, 64 insertions(+), 2 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..cafe592a0370007 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -618,6 +618,10 @@ MLIR_CAPI_EXPORTED bool
mlirOperationRemoveDiscardableAttributeByName(MlirOperation op,
MlirStringRef name);
+/// Set `pos`-th successor of the operation.
+MLIR_CAPI_EXPORTED void
+mlirOperationSetSuccessor(MlirOperation op, MlirBlock block, intptr_t pos);
+
/// Returns the number of attributes attached to the operation.
/// Deprecated, please use `mlirOperationGetNumInherentAttributes` or
/// `mlirOperationGetNumDiscardableAttributes`.
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index aad74f511e7ef11..4855300df44873f 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2924,7 +2924,17 @@ 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(
+ "num_successors",
+ [](PyOperation &self) { mlirOperationGetNumSuccessors(self); })
+ .def("get_successor",
+ [](PyOperation &self, int pos) {
+ mlirOperationGetSuccessor(self, pos);
+ })
+ .def("set_successor", [](PyOperation &self, PyBlock block, int pos) {
+ mlirOperationSetSuccessor(self, block.get(), pos);
+ });
auto opViewClass =
py::class_<PyOpView, PyOperationBase>(m, "OpView", py::module_local())
@@ -3448,7 +3458,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) {
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index 65b2b7466fb7c39..18de1cc00b47d5d 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, MlirBlock block,
+ intptr_t pos) {
+ 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..82e8678987a6a87
--- /dev/null
+++ b/mlir/test/python/dialects/cf.py
@@ -0,0 +1,42 @@
+# 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()
+ print(module)
+ print(module.operation.verify())
+ 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)
+ terminator = Operation.create("custom.terminator", ip=ip)
+
+ block1 = op1.regions[0].blocks.append()
+ ip = InsertionPoint(block1)
+ br1 = cf.BranchOp([], block1, ip=ip)
+
+ block2 = op1.regions[0].blocks.append()
+ ip = InsertionPoint(block2)
+ br2 = cf.BranchOp([], block1, ip=ip)
+
+ br1.operation.set_successor(block2, 0)
+
+
+# CHECK: ^bb1: // pred: ^bb2
+# CHECK: cf.br ^bb2
+# CHECK: ^bb2: // pred: ^bb1
+# CHECK: cf.br ^bb1
More information about the Mlir-commits
mailing list