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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

<details>
<summary>Changes</summary>

`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.

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


1 Files Affected:

- (modified) mlir/include/mlir/Bindings/Python/IRAttributes.h (+10-5) 


``````````diff
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.

``````````

</details>


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