[Mlir-commits] [mlir] [MLIR][Python] Improve Iterator performance. Don't `throw` in `dunderNext` methods. (PR #175377)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jan 10 11:07:42 PST 2026
https://github.com/MaPePeR created https://github.com/llvm/llvm-project/pull/175377
In https://github.com/llvm/llvm-project/pull/174139#issuecomment-3733259370 I wrote a scuffed benchmark that mostly iterates MLIR Container Types in Python. My changes from that PR made the performance worse, so I closed it.
However, when experimetning with that I also saw a large(?) performance gain by changing the `dunderNext` methods of the various Iterators to use `PyErr_SetNone(PyExc_StopIteration);` instead of `throw nb::stop_iteration();`.
<details><summary>Benchmark attempt script</summary>
```python
import timeit
from mlir.ir import Context, Location, Module, InsertionPoint, Block, Region, OpView
from mlir.dialects import func, builtin, scf, arith
def generate_module():
m = Module.create()
with InsertionPoint(m.body):
f = func.FuncOp("main", builtin.FunctionType.get([], []))
with InsertionPoint(f.body.blocks.append()):
generate_ops(10, 2)
func.ReturnOp([])
return m
def generate_ops(count: int, depth: int):
if depth == 0:
return
lower = arith.ConstantOp(builtin.IntegerType.get_signless(64), 0)
upper = arith.ConstantOp(builtin.IntegerType.get_signless(64), 100)
step = arith.ConstantOp(builtin.IntegerType.get_signless(64), 1)
for i in range(count):
forop = scf.ForOp(lower, upper, step)
with InsertionPoint(forop.region.blocks[0]):
generate_ops(count, depth - 1)
scf.YieldOp([])
def walk_module(m: Module):
walk_block(m.body)
def walk_region(region: Region):
for block in region.blocks:
walk_block(block)
def walk_block(block: Block):
for predecessors in block.predecessors:
pass
for successors in block.successors:
pass
for op in block.operations:
walk_op(op)
def walk_op(op: OpView):
for result in op.results:
pass
for successors in op.successors:
pass
for operands in op.operands:
pass
for region in op.regions:
walk_region(region)
with Context(), Location.unknown():
m = generate_module()
# From timeit.main:
t = timeit.Timer(lambda: walk_module(m))
number, _ = t.autorange()
repeats = 5
raw_timings = t.repeat(repeats, number)
timings = [dt / number for dt in raw_timings]
best = min(timings)
print(f"{number} loops, best of {repeats}: {best * 1000:.3g} msecs per loop")
```
</details>
The performance of the benchmark went from
```
50 loops, best of 5: 5.97 msecs per loop
```
to
```
50 loops, best of 5: 5.12 msecs per loop
```
which is a ~14% improvement. (Though you should validate that yourself, probably. My test setup is very scuffed)
The functions were previously set to return a C++ type like `PyRegion`. Because of the removal of the `throw` they now had to [return a `NULL` value to Python](https://github.com/python/cpython/blob/aa8578dc54df2af9daa3353566359e602e5905cf/Objects/call.c#L49-L61), so I changed the return type to `nanobind::typed<nanobind::object,PyRegion>` so I could return an `nb::object()` in case an error was set and otherwise `nb::cast` the `PyRegion` value to `nb::object` instead of returning it directly.
I'm not a huge fan, that this changes the external "Usage" of the functions, because now they won't bubble up exceptions, when they are called from C++ The return type and Python Error State have to be checked instead.
I couldn't find any location that called them in llvm itself, though. Maybe these functions should not be public, because they are only supposed to be called from Python anyway?
>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] [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) {
More information about the Mlir-commits
mailing list