[Mlir-commits] [mlir] [mlir][CAPI, python bindings] Expose `Operation::setSuccessor` (PR #67922)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/67922

>From c70e3866a9610d8d009ec2aa8b0480779a5e25ac 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     | 50 +++++++++++++++++++++++++++++
 4 files changed, 72 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..1dbd29cb7d0c8bc 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) { return mlirOperationGetNumSuccessors(self); })
+      .def("get_successor",
+           [](PyOperation &self, int pos) -> PyBlock {
+             return {self.getRef(), 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..38a6bcd904b5e13
--- /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)
+    terminator = 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.operation.get_successor(0))
+    # CHECK: num_successors 1
+    print("num_successors", br1.operation.num_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.operation.get_successor(0))
+    # CHECK: num_successors 1
+    print("num_successors", br2.operation.num_successors)
+
+    br1.operation.set_successor(block2, 0)
+    # CHECK: ^bb2:  // pred: ^bb1
+    # CHECK:   cf.br ^bb1
+    print(br1.operation.get_successor(0))
+    # CHECK: ^bb1:  // pred: ^bb2
+    # CHECK:   cf.br ^bb2
+    print(br2.operation.get_successor(0))