[Mlir-commits] [mlir] [MLIR] [Python] Fixed a long standing bug in the `PyRegionListIterator` (PR #188475)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Sergei Lebedev (superbobry)

<details>
<summary>Changes</summary>

`PyRegionIterator` did not account for start index/step, so this commit removes it in favor of the sequence iterator provided by CPython.

The previous attempt in #<!-- -->137232 bitrotted, so I decided to open a new PR.

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


3 Files Affected:

- (modified) mlir/include/mlir/Bindings/Python/IRCore.h (-20) 
- (modified) mlir/lib/Bindings/Python/IRCore.cpp (-30) 
- (modified) mlir/test/python/ir/operation.py (+12-1) 


``````````diff
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index eefc51d519d62..e834fca61522f 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1378,22 +1378,6 @@ struct MLIR_PYTHON_API_EXPORTED PyAttrBuilderMap {
 // Collections.
 //------------------------------------------------------------------------------
 
-class MLIR_PYTHON_API_EXPORTED PyRegionIterator {
-public:
-  PyRegionIterator(PyOperationRef operation, int nextIndex)
-      : operation(std::move(operation)), nextIndex(nextIndex) {}
-
-  PyRegionIterator &dunderIter() { return *this; }
-
-  nanobind::typed<nanobind::object, PyRegion> dunderNext();
-
-  static void bind(nanobind::module_ &m);
-
-private:
-  PyOperationRef operation;
-  intptr_t nextIndex = 0;
-};
-
 /// Regions of an op are fixed length and indexed numerically so are represented
 /// with a sequence-like container.
 class MLIR_PYTHON_API_EXPORTED PyRegionList
@@ -1404,10 +1388,6 @@ class MLIR_PYTHON_API_EXPORTED PyRegionList
   PyRegionList(PyOperationRef operation, intptr_t startIndex = 0,
                intptr_t length = -1, intptr_t step = 1);
 
-  PyRegionIterator dunderIter();
-
-  static void bindDerived(ClassTy &c);
-
 private:
   /// Give the parent CRTP class access to hook implementations below.
   friend class Sliceable<PyRegionList, PyRegion>;
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index f3f1ee4ce343f..12f06a3cb2b5b 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -189,25 +189,6 @@ nb::object PyBlock::getCapsule() {
 // Collections.
 //------------------------------------------------------------------------------
 
-nb::typed<nb::object, PyRegion> PyRegionIterator::dunderNext() {
-  operation->checkValid();
-  if (nextIndex >= mlirOperationGetNumRegions(operation->get())) {
-    PyErr_SetNone(PyExc_StopIteration);
-    // python functions should return NULL after setting any exception
-    return nb::object();
-  }
-  MlirRegion region = mlirOperationGetRegion(operation->get(), nextIndex++);
-  return nb::cast(PyRegion(operation, region));
-}
-
-void PyRegionIterator::bind(nb::module_ &m) {
-  nb::class_<PyRegionIterator>(m, "RegionIterator")
-      .def("__iter__", &PyRegionIterator::dunderIter,
-           "Returns an iterator over the regions in the operation.")
-      .def("__next__", &PyRegionIterator::dunderNext,
-           "Returns the next region in the iteration.");
-}
-
 PyRegionList::PyRegionList(PyOperationRef operation, intptr_t startIndex,
                            intptr_t length, intptr_t step)
     : Sliceable(startIndex,
@@ -216,16 +197,6 @@ PyRegionList::PyRegionList(PyOperationRef operation, intptr_t startIndex,
                 step),
       operation(std::move(operation)) {}
 
-PyRegionIterator PyRegionList::dunderIter() {
-  operation->checkValid();
-  return PyRegionIterator(operation, startIndex);
-}
-
-void PyRegionList::bindDerived(ClassTy &c) {
-  c.def("__iter__", &PyRegionList::dunderIter,
-        "Returns an iterator over the regions in the sequence.");
-}
-
 intptr_t PyRegionList::getRawNumElements() {
   operation->checkValid();
   return mlirOperationGetNumRegions(operation->get());
@@ -5012,7 +4983,6 @@ void populateIRCore(nb::module_ &m) {
   PyOpOperands::bind(m);
   PyOpResultList::bind(m);
   PyOpSuccessors::bind(m);
-  PyRegionIterator::bind(m);
   PyRegionList::bind(m);
 
   // Debug bindings.
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index f561a1bc624d8..9f1ad01cbcc88 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -86,7 +86,7 @@ def walk_operations(indent, op):
     # CHECK:           OP 1: func.return
     walk_operations("", op)
 
-    # CHECK:    Region iter: <mlir.{{.+}}.RegionIterator
+    # CHECK:    Region iter: <iterator
     # CHECK:     Block iter: <mlir.{{.+}}.BlockIterator
     # CHECK: Operation iter: <mlir.{{.+}}.OperationIterator
     print("   Region iter:", iter(op.regions))
@@ -109,6 +109,17 @@ def walk_operations(indent, op):
         # CHECK: attempt to access out of bounds operation
         print(e)
 
+    # Verify that iterating a sliced region list yields the correct
+    # number of elements (i.e. respects length and step).
+    with Location.unknown(ctx):
+        op4 = Operation.create("custom.op", regions=4)
+        r = op4.regions
+        assert len(list(r[:])) == 4
+        assert len(list(r[1:])) == 3
+        assert len(list(r[::2])) == 2
+        assert len(list(r[1:3])) == 2
+        assert len(list(r[::-1])) == 4
+
 
 # Verify index based traversal of the op/region/block hierarchy.
 # CHECK-LABEL: TEST: testTraverseOpRegionBlockIndices

``````````

</details>


https://github.com/llvm/llvm-project/pull/188475