[Mlir-commits] [mlir] 351d06a - [MLIR][Python] Improve Iterator performance. Don't `throw` in `dunderNext` methods. (#175377)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jan 13 08:02:43 PST 2026
Author: MaPePeR
Date: 2026-01-13T16:02:38Z
New Revision: 351d06a8191e3ae50a73149c3b67ae2decb49e26
URL: https://github.com/llvm/llvm-project/commit/351d06a8191e3ae50a73149c3b67ae2decb49e26
DIFF: https://github.com/llvm/llvm-project/commit/351d06a8191e3ae50a73149c3b67ae2decb49e26.diff
LOG: [MLIR][Python] Improve Iterator performance. Don't `throw` in `dunderNext` methods. (#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
```
in my setup, 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?
---------
Co-authored-by: Maksim Levental <maksim.levental at gmail.com>
Added:
Modified:
mlir/include/mlir/Bindings/Python/IRCore.h
mlir/lib/Bindings/Python/IRAttributes.cpp
mlir/lib/Bindings/Python/IRCore.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 050260e65d03b..599771f8a3283 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1378,7 +1378,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);
@@ -1421,7 +1421,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);
@@ -1512,7 +1512,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..b2e9d9887e098 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -226,8 +226,11 @@ 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);
+ // python functions should return NULL after setting any exception
+ 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 8b3542c717ca2..eb00363a54034 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -191,13 +191,15 @@ 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);
+ // python functions should return NULL after setting any exception
+ 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) {
@@ -241,15 +243,17 @@ 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);
+ // python functions should return NULL after setting any exception
+ return nb::object();
}
PyBlock returnBlock(operation, next);
next = mlirBlockGetNextInRegion(next);
- return returnBlock;
+ return nb::cast(returnBlock);
}
void PyBlockIterator::bind(nb::module_ &m) {
@@ -324,7 +328,9 @@ 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);
+ // python functions should return NULL after setting any exception
+ return nb::object();
}
PyOperationRef returnOperation =
@@ -407,13 +413,16 @@ 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);
+ // python functions should return NULL after setting any exception
+ 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