[Mlir-commits] [mlir] [MLIR][Python] Improve Iterator performance. Don't `throw` in `dunderNext` methods. (PR #175377)

[llvmlistbot at llvm.org]

https://github.com/MaPePeR updated https://github.com/llvm/llvm-project/pull/175377

>From 281d10ef0a0325204628d90501b9eb366aaee648 Mon Sep 17 00:00:00 2001
From: MaPePeR <MaPePeR at users.noreply.github.com>
Date: Sat, 10 Jan 2026 17:59:36 +0000
Subject: [PATCH 1/2] [MLIR][Python] Improve Iterator performance. Don't
 `throw` in `dunderNext` methods.

---
 mlir/include/mlir/Bindings/Python/IRCore.h |  6 ++---
 mlir/lib/Bindings/Python/IRAttributes.cpp  |  6 +++--
 mlir/lib/Bindings/Python/IRCore.cpp        | 29 +++++++++++++---------
 3 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 330318683c15e..bd0c715813bd9 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1374,7 +1374,7 @@ class MLIR_PYTHON_API_EXPORTED PyRegionIterator {
 
   PyRegionIterator &dunderIter() { return *this; }
 
-  PyRegion dunderNext();
+  nanobind::typed<nanobind::object, PyRegion> dunderNext();
 
   static void bind(nanobind::module_ &m);
 
@@ -1417,7 +1417,7 @@ class MLIR_PYTHON_API_EXPORTED PyBlockIterator {
 
   PyBlockIterator &dunderIter() { return *this; }
 
-  PyBlock dunderNext();
+  nanobind::typed<nanobind::object, PyBlock> dunderNext();
 
   static void bind(nanobind::module_ &m);
 
@@ -1508,7 +1508,7 @@ class MLIR_PYTHON_API_EXPORTED PyOpOperandIterator {
 
   PyOpOperandIterator &dunderIter() { return *this; }
 
-  PyOpOperand dunderNext();
+  nanobind::typed<nanobind::object, PyOpOperand> dunderNext();
 
   static void bind(nanobind::module_ &m);
 
diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index 3685ff0d602e2..0db0aa53f3ebb 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -226,8 +226,10 @@ PyArrayAttribute::PyArrayAttributeIterator::dunderNext() {
   // TODO: Throw is an inefficient way to stop iteration.
   if (PyArrayAttribute::PyArrayAttributeIterator::nextIndex >=
       mlirArrayAttrGetNumElements(
-          PyArrayAttribute::PyArrayAttributeIterator::attr.get()))
-    throw nb::stop_iteration();
+          PyArrayAttribute::PyArrayAttributeIterator::attr.get())) {
+    PyErr_SetNone(PyExc_StopIteration);
+    return nb::object();
+  }
   return PyAttribute(
              this->PyArrayAttribute::PyArrayAttributeIterator::attr
                  .getContext(),
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 19db41fae4fe2..d5d9929724788 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -194,13 +194,14 @@ nb::object PyBlock::getCapsule() {
 // Collections.
 //------------------------------------------------------------------------------
 
-PyRegion PyRegionIterator::dunderNext() {
+nb::typed<nb::object, PyRegion> PyRegionIterator::dunderNext() {
   operation->checkValid();
   if (nextIndex >= mlirOperationGetNumRegions(operation->get())) {
-    throw nb::stop_iteration();
+    PyErr_SetNone(PyExc_StopIteration);
+    return nb::object();
   }
   MlirRegion region = mlirOperationGetRegion(operation->get(), nextIndex++);
-  return PyRegion(operation, region);
+  return nb::cast(PyRegion(operation, region));
 }
 
 void PyRegionIterator::bind(nb::module_ &m) {
@@ -244,15 +245,16 @@ PyRegionList PyRegionList::slice(intptr_t startIndex, intptr_t length,
   return PyRegionList(operation, startIndex, length, step);
 }
 
-PyBlock PyBlockIterator::dunderNext() {
+nb::typed<nb::object, PyBlock> PyBlockIterator::dunderNext() {
   operation->checkValid();
   if (mlirBlockIsNull(next)) {
-    throw nb::stop_iteration();
+    PyErr_SetNone(PyExc_StopIteration);
+    return nb::object();
   }
 
   PyBlock returnBlock(operation, next);
   next = mlirBlockGetNextInRegion(next);
-  return returnBlock;
+  return nb::cast(returnBlock);
 }
 
 void PyBlockIterator::bind(nb::module_ &m) {
@@ -327,13 +329,14 @@ void PyBlockList::bind(nb::module_ &m) {
 nb::typed<nb::object, PyOpView> PyOperationIterator::dunderNext() {
   parentOperation->checkValid();
   if (mlirOperationIsNull(next)) {
-    throw nb::stop_iteration();
+    PyErr_SetNone(PyExc_StopIteration);
+    return nb::object();
   }
 
   PyOperationRef returnOperation =
       PyOperation::forOperation(parentOperation->getContext(), next);
   next = mlirOperationGetNextInBlock(next);
-  return returnOperation->createOpView();
+  return nb::cast(returnOperation->createOpView());
 }
 
 void PyOperationIterator::bind(nb::module_ &m) {
@@ -410,13 +413,15 @@ void PyOpOperand::bind(nb::module_ &m) {
                    "Returns the operand number in the owning operation.");
 }
 
-PyOpOperand PyOpOperandIterator::dunderNext() {
-  if (mlirOpOperandIsNull(opOperand))
-    throw nb::stop_iteration();
+nb::typed<nb::object, PyOpOperand> PyOpOperandIterator::dunderNext() {
+  if (mlirOpOperandIsNull(opOperand)) {
+    PyErr_SetNone(PyExc_StopIteration);
+    return nb::object();
+  }
 
   PyOpOperand returnOpOperand(opOperand);
   opOperand = mlirOpOperandGetNextUse(opOperand);
-  return returnOpOperand;
+  return nb::cast(returnOpOperand);
 }
 
 void PyOpOperandIterator::bind(nb::module_ &m) {

>From 70b5a2343bccc6fb89a7ae4480af40f3f3657c01 Mon Sep 17 00:00:00 2001
From: MaPePeR <MaPePeR at users.noreply.github.com>
Date: Tue, 13 Jan 2026 16:33:15 +0100
Subject: [PATCH 2/2] Update mlir/lib/Bindings/Python/IRAttributes.cpp

Co-authored-by: Maksim Levental <maksim.levental at gmail.com>
---
 mlir/lib/Bindings/Python/IRAttributes.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index 0db0aa53f3ebb..b2e9d9887e098 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -228,6 +228,7 @@ PyArrayAttribute::PyArrayAttributeIterator::dunderNext() {
       mlirArrayAttrGetNumElements(
           PyArrayAttribute::PyArrayAttributeIterator::attr.get())) {
     PyErr_SetNone(PyExc_StopIteration);
+    // python functions should return NULL after setting any exception
     return nb::object();
   }
   return PyAttribute(