[Mlir-commits] [mlir] [MLIR][Python] Don't throw in PyDenseArrayIterator dunderNext (PR #218193)

[Maksim Levental]

https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/218193

`PyDenseArrayIterator::dunderNext` signaled iterator exhaustion by throwing `nanobind::stop_iteration()`. Raising a C++ exception to signal `StopIteration` incurs stack-unwinding cost on every loop over a dense array attribute.

#175377 replaced this pattern with `PyErr_SetNone(PyExc_StopIteration)` (return a null object after setting the Python error indicator) for the other iterators in the bindings, measuring a ~14% improvement on a container-walk benchmark. `PyDenseArrayIterator` was missed in that change and still throws.

This PR applies the same conversion:
- Signal exhaustion via `PyErr_SetNone(PyExc_StopIteration)` and return `nanobind::object()`.
- Return type changes from `EltTy` to `nanobind::typed<nanobind::object, EltTy>` so the method can return a null object after setting the error, matching the existing `nanobind::typed<nanobind::object, PyAttribute>` iterator in the same header.

Standalone perf/consistency fix; independent of any nanobind version bump.

>From 232b40cec9c5208e4f06e4815d84269806038837 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Sat, 22 Aug 2026 21:47:07 -0700
Subject: [PATCH] [MLIR][Python] Don't throw in PyDenseArrayIterator dunderNext

Signal iterator exhaustion via PyErr_SetNone(PyExc_StopIteration) instead
of throwing nanobind::stop_iteration(), matching the pattern used by the
other iterators in the Python bindings (see #175377). Raising a C++
exception to signal StopIteration incurs stack-unwinding cost on every
loop, which PyDenseArrayIterator was still paying. The return type changes
to nanobind::typed<nanobind::object, EltTy> so the method can return a null
object after setting the Python error indicator.
---
 mlir/include/mlir/Bindings/Python/IRAttributes.h | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Bindings/Python/IRAttributes.h b/mlir/include/mlir/Bindings/Python/IRAttributes.h
index fcc003ab365fa..fbebea40f1ac8 100644
--- a/mlir/include/mlir/Bindings/Python/IRAttributes.h
+++ b/mlir/include/mlir/Bindings/Python/IRAttributes.h
@@ -131,11 +131,16 @@ class MLIR_PYTHON_API_EXPORTED PyDenseArrayAttribute
     PyDenseArrayIterator dunderIter() { return *this; }
 
     /// Return the next element.
-    EltTy dunderNext() {
-      // Throw if the index has reached the end.
-      if (nextIndex >= mlirDenseArrayGetNumElements(attr.get()))
-        throw nanobind::stop_iteration();
-      return DerivedT::getElement(attr.get(), nextIndex++);
+    nanobind::typed<nanobind::object, EltTy> dunderNext() {
+      // Set StopIteration if the index has reached the end. Signaling
+      // exhaustion via the Python error indicator rather than a C++ exception
+      // avoids the cost of stack unwinding on every iteration.
+      if (nextIndex >= mlirDenseArrayGetNumElements(attr.get())) {
+        PyErr_SetNone(PyExc_StopIteration);
+        // python functions should return NULL after setting any exception
+        return nanobind::object();
+      }
+      return nanobind::cast(DerivedT::getElement(attr.get(), nextIndex++));
     }
 
     /// Bind the iterator class.