[llvm] [mlir][CAPI, python bindings] Expose `Operation::setSuccessor` (PR #67922)
Maksim Levental via llvm-commits
llvm-commits at lists.llvm.org
Sun Oct 1 17:03:52 PDT 2023
https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/67922
>From bb702b0c4c7f9cdcaf2ee1112141c8e9d6883b2a Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Sun, 1 Oct 2023 17:12:30 -0500
Subject: [PATCH 1/2] [flang-rt] fix missing llvm_gtest target
---
flang-rt/unittests/CMakeLists.txt | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/flang-rt/unittests/CMakeLists.txt b/flang-rt/unittests/CMakeLists.txt
index 2b813c6e16dfaa1..cfd6b4c973e50f7 100644
--- a/flang-rt/unittests/CMakeLists.txt
+++ b/flang-rt/unittests/CMakeLists.txt
@@ -33,8 +33,9 @@ function(add_flang_rt_unittest_offload_properties target)
endif()
endfunction()
-if(NOT TARGET llvm_gtest)
- message(FATAL_ERROR "Target llvm_gtest not found.")
+set(UNITTEST_DIR ${LLVM_THIRD_PARTY_DIR}/unittest)
+if (NOT TARGET llvm_gtest)
+ add_subdirectory(${UNITTEST_DIR} third-party/unittest)
endif()
function(add_flang_rt_unittest test_dirname)
>From b375d883830bcd4eae07e73ebcdf7e8426dbddbc 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 2/2] [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))
More information about the llvm-commits
mailing list