[Mlir-commits] [mlir] [MLIR][Python] Remove pybind workaround for not throwing IndexError (PR #174139)

[llvmlistbot at llvm.org]

https://github.com/MaPePeR created https://github.com/llvm/llvm-project/pull/174139

With 429b0cf1de14471c9d258467bfc9936c3a9d52f7 a workaround was introduced to improve the performance of the `Sequence` Protocol for `Sliceable`s, which main benefit was to not require throwing C++ exceptions (according to the code-coment):

https://github.com/llvm/llvm-project/blob/429b0cf1de14471c9d258467bfc9936c3a9d52f7/mlir/lib/Bindings/Python/PybindUtils.h#L291-L299
(Note, that this explanation was [later](https://github.com/llvm/llvm-project/blob/b56d1ec6cb8b5cb3ff46cba39a1049ecf3831afb/mlir/lib/Bindings/Python/NanobindUtils.h#L321-L329) automatically replaced to state that this was "4x faster than via nanobind", but I'm pretty certain that this is just a `s/pybind11/nanobind/` gone wrong, because it also broke the issue URL)

This workaround caused me some confusion, because I wrongly assumed that it was the reason the `Sliceable` Sequences were not compatible with the `match` statement (See #174091). So I looked into it and with `nanobind` it seems to be perfectly fine to call `PyErr_SetString` directly instead of throwing a C++ exception, so the `Sequence` protocol can be implemented the "normal" way without the performance implications of C++ exceptions.
It only required registering the already existing functions with nanobind and changing one function signature slightly.

Though I did not benchmark this change, so this is only a Draft for now.

There are also a lot of other places where `PyErr_SetString` could be used to report Exceptions back to Python (e.g. 72x`throw nb::` in `IRCore.cpp`), but its not as easy as simply replacing `throw` with `Py_Err_Set*`, because then it would result in unexpected behavior when function calls are nested.

@stellaraccident I'm mentioning you here, because you implemented the original workaround, I hope that is fine with you. Do you maybe still have the benchmark and could take a look at this? Thank you very much.

>From 86543e29894243c953c76aabe023e4b845e2284c Mon Sep 17 00:00:00 2001
From: MaPePeR <MaPePeR at users.noreply.github.com>
Date: Wed, 31 Dec 2025 20:25:22 +0000
Subject: [PATCH] [MLIR][Python] Remove pybind workaround for not throwing
 IndexError

The workaround was required to avoid the throwing of C++ exceptions with
pybind. This is not required anymore, because with nanobind the
`PyErr_SetString` function can be called directly.
---
 mlir/lib/Bindings/Python/NanobindUtils.h | 50 +++---------------------
 1 file changed, 5 insertions(+), 45 deletions(-)

diff --git a/mlir/lib/Bindings/Python/NanobindUtils.h b/mlir/lib/Bindings/Python/NanobindUtils.h
index aea195fecae82..b0f5069201735 100644
--- a/mlir/lib/Bindings/Python/NanobindUtils.h
+++ b/mlir/lib/Bindings/Python/NanobindUtils.h
@@ -296,9 +296,9 @@ class Sliceable {
 
   /// Returns a new instance of the pseudo-container restricted to the given
   /// slice. Returns a nullptr object on failure.
-  nanobind::object getItemSlice(PyObject *slice) {
+  nanobind::object getItemSlice(nanobind::slice slice) {
     ssize_t start, stop, extraStep, sliceLength;
-    if (PySlice_GetIndicesEx(slice, length, &start, &stop, &extraStep,
+    if (PySlice_GetIndicesEx(slice.ptr(), length, &start, &stop, &extraStep,
                              &sliceLength) != 0) {
       PyErr_SetString(PyExc_IndexError, "index out of range");
       return {};
@@ -355,51 +355,11 @@ class Sliceable {
                       nanobind::cast<std::string>(elemTyName) + "])";
     auto clazz = nanobind::class_<Derived>(m, Derived::pyClassName,
                                            nanobind::sig(sig.c_str()))
+                     .def("__len__", &Sliceable::size)
+                     .def("__getitem__", &Sliceable::getItem)
+                     .def("__getitem__", &Sliceable::getItemSlice)
                      .def("__add__", &Sliceable::dunderAdd);
     Derived::bindDerived(clazz);
-
-    // Manually implement the sequence protocol via the C API. We do this
-    // because it is approx 4x faster than via nanobind, largely because that
-    // formulation requires a C++ exception to be thrown to detect end of
-    // sequence.
-    // Since we are in a C-context, any C++ exception that happens here
-    // will terminate the program. There is nothing in this implementation
-    // that should throw in a non-terminal way, so we forgo further
-    // exception marshalling.
-    // See: https://github.com/pybind/nanobind/issues/2842
-    auto heap_type = reinterpret_cast<PyHeapTypeObject *>(clazz.ptr());
-    assert(heap_type->ht_type.tp_flags & Py_TPFLAGS_HEAPTYPE &&
-           "must be heap type");
-    heap_type->as_sequence.sq_length = +[](PyObject *rawSelf) -> Py_ssize_t {
-      auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
-      return self->length;
-    };
-    // sq_item is called as part of the sequence protocol for iteration,
-    // list construction, etc.
-    heap_type->as_sequence.sq_item =
-        +[](PyObject *rawSelf, Py_ssize_t index) -> PyObject * {
-      auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
-      return self->getItem(index).release().ptr();
-    };
-    // mp_subscript is used for both slices and integer lookups.
-    heap_type->as_mapping.mp_subscript =
-        +[](PyObject *rawSelf, PyObject *rawSubscript) -> PyObject * {
-      auto self = nanobind::cast<Derived *>(nanobind::handle(rawSelf));
-      Py_ssize_t index = PyNumber_AsSsize_t(rawSubscript, PyExc_IndexError);
-      if (!PyErr_Occurred()) {
-        // Integer indexing.
-        return self->getItem(index).release().ptr();
-      }
-      PyErr_Clear();
-
-      // Assume slice-based indexing.
-      if (PySlice_Check(rawSubscript)) {
-        return self->getItemSlice(rawSubscript).release().ptr();
-      }
-
-      PyErr_SetString(PyExc_ValueError, "expected integer or slice");
-      return nullptr;
-    };
   }
 
   /// Hook for derived classes willing to bind more methods.