[mlir][CAPI, python bindings] Expose Operation::setSuccessor

[makslevental]

This is useful for emitting (using the python bindings) cf.br to blocks that are declared lexically post block creation.

[llvmbot]

@llvm/pr-subscribers-mlir

Changes

Expose Operation::setSuccessor.

---

Full diff: https://github.com/llvm/llvm-project/pull/67922.diff

4 Files Affected:

• (modified) mlir/include/mlir-c/IR.h (+4)
• (modified) mlir/lib/Bindings/Python/IRCore.cpp (+13-2)
• (modified) mlir/lib/CAPI/IR/IR.cpp (+5)
• (added) mlir/test/python/dialects/cf.py (+42)

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
+@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

[jpienaar]

Looks good, thanks (I see this follows the C convention, I'd have flipped block and pos, but keeping consistent better).