[Mlir-commits] [mlir] [MLIR] [python] Fixed a bug in `PyRegionList` iterator (PR #137232)

[Sergei Lebedev]

https://github.com/superbobry updated https://github.com/llvm/llvm-project/pull/137232

>From fa454352b77a7108b54c82b7487b125347406272 Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <slebedev at google.com>
Date: Thu, 24 Apr 2025 19:36:52 +0100
Subject: [PATCH 1/2] [MLIR] [python] Fixed a bug in `PyRegionList` iterator

`PyRegionIterator` did not accoutn for start index/step, so this commit removes
it in favor of the sequence iterator provided by CPython.
---
 mlir/lib/Bindings/Python/IRCore.cpp | 37 -----------------------------
 mlir/test/python/ir/operation.py    | 16 ++++++++-----
 2 files changed, 10 insertions(+), 43 deletions(-)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index b5720b7ad8b21..77f6be7da31ea 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -332,33 +332,6 @@ nb::object PyBlock::getCapsule() {
 
 namespace {
 
-class PyRegionIterator {
-public:
-  PyRegionIterator(PyOperationRef operation)
-      : operation(std::move(operation)) {}
-
-  PyRegionIterator &dunderIter() { return *this; }
-
-  PyRegion dunderNext() {
-    operation->checkValid();
-    if (nextIndex >= mlirOperationGetNumRegions(operation->get())) {
-      throw nb::stop_iteration();
-    }
-    MlirRegion region = mlirOperationGetRegion(operation->get(), nextIndex++);
-    return PyRegion(operation, region);
-  }
-
-  static void bind(nb::module_ &m) {
-    nb::class_<PyRegionIterator>(m, "RegionIterator")
-        .def("__iter__", &PyRegionIterator::dunderIter)
-        .def("__next__", &PyRegionIterator::dunderNext);
-  }
-
-private:
-  PyOperationRef operation;
-  int nextIndex = 0;
-};
-
 /// Regions of an op are fixed length and indexed numerically so are represented
 /// with a sequence-like container.
 class PyRegionList : public Sliceable<PyRegionList, PyRegion> {
@@ -373,15 +346,6 @@ class PyRegionList : public Sliceable<PyRegionList, PyRegion> {
                   step),
         operation(std::move(operation)) {}
 
-  PyRegionIterator dunderIter() {
-    operation->checkValid();
-    return PyRegionIterator(operation);
-  }
-
-  static void bindDerived(ClassTy &c) {
-    c.def("__iter__", &PyRegionList::dunderIter);
-  }
-
 private:
   /// Give the parent CRTP class access to hook implementations below.
   friend class Sliceable<PyRegionList, PyRegion>;
@@ -4105,7 +4069,6 @@ void mlir::python::populateIRCore(nb::module_ &m) {
   PyOpOperandList::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 b08fe98397fbc..355b53ceaa158 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -82,12 +82,16 @@ def walk_operations(indent, op):
     # CHECK:           OP 1: func.return
     walk_operations("", op)
 
-    # CHECK:    Region iter: <mlir.{{.+}}.RegionIterator
-    # CHECK:     Block iter: <mlir.{{.+}}.BlockIterator
-    # CHECK: Operation iter: <mlir.{{.+}}.OperationIterator
-    print("   Region iter:", iter(op.regions))
-    print("    Block iter:", iter(op.regions[-1]))
-    print("Operation iter:", iter(op.regions[-1].blocks[-1]))
+    # CHECK:    Regions: [<mlir.{{.+}}.Region object at {{.+}}>]
+    # CHECK:     Blocks: [<mlir.{{.+}}.Block object at {{.+}}>]
+    # CHECK: Operations: [<mlir.{{.+}}.FuncOp object at {{.+}}>]
+    print("   Regions:", [*op.regions])
+    print("    Blocks:", [*op.regions[-1]])
+    print("Operations:", [*op.regions[-1].blocks[-1]])
+
+    # Make sure the iterator uses offsets and the resulting list is empty.
+    # CHECK: []
+    print([*op.regions[1:]])
 
     try:
         op.regions[-42]

>From bd54eb07a060452b4c7f81a97cb3a8e68db21bcf Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <185856+superbobry at users.noreply.github.com>
Date: Fri, 25 Apr 2025 08:09:37 +0100
Subject: [PATCH 2/2] Update mlir/test/python/ir/operation.py

Co-authored-by: Maksim Levental <maksim.levental at gmail.com>
---
 mlir/test/python/ir/operation.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index 355b53ceaa158..ccc8153832c8f 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -85,9 +85,9 @@ def walk_operations(indent, op):
     # CHECK:    Regions: [<mlir.{{.+}}.Region object at {{.+}}>]
     # CHECK:     Blocks: [<mlir.{{.+}}.Block object at {{.+}}>]
     # CHECK: Operations: [<mlir.{{.+}}.FuncOp object at {{.+}}>]
-    print("   Regions:", [*op.regions])
-    print("    Blocks:", [*op.regions[-1]])
-    print("Operations:", [*op.regions[-1].blocks[-1]])
+    print("   Regions:", list(op.regions))
+    print("    Blocks:", list(op.regions[-1]))
+    print("Operations:", list(op.regions[-1].blocks[-1]))
 
     # Make sure the iterator uses offsets and the resulting list is empty.
     # CHECK: []