[mlir][python] Directly implement sequence protocol on Sliceable.

[stellaraccident]

Summary

• While annoying, this is the only way to get C++ exception handling out of the happy path for normal iteration.
• Implements sq_length and sq_item for the sequence protocol (used for iteration, including list() construction).
• Implements mp_subscript for general use (i.e. foo[1] and foo[1:1]).
• For constructing a list(op.results), this reduces the time from ~4-5us to ~1.5us on my machine (give or take measurement overhead) and eliminates C++ exceptions, which is a worthy goal in itself.
  • Compared to a baseline of similar construction of a three-integer list, which takes 450ns (might just be measuring function call overhead).
  • See issue discussed on the pybind side: https://github.com/pybind/pybind11/issues/2842

[ftynse]

mlir/lib/Bindings/Python/PybindUtils.h

237

Don't we want a non-empty error message here?

255

Ditto.

303

I suppose this cast can never throw (self is always an instance of Derived * if CRTP is used properly), but it may warrant a comment.

326–328

Can't we just do if(PySlice_Check(rawSubscript)) instead of casting and catching an exception?

[phawkins]

Thanks! I have little to add beyond @ftynse 's comments.

mlir/lib/Bindings/Python/PybindUtils.h

291–294

You might consider linking to https://github.com/pybind/pybind11/issues/2842 and you might also consider commenting on that issue linking to your solution, since I'm sure others may have the same need.

[jpienaar]

mlir/lib/Bindings/Python/PybindUtils.h

237

The kind of error set is an index error, OOC is there any other index error than OOB?

[ftynse]

mlir/lib/Bindings/Python/PybindUtils.h

237

It still has a message:

>>> a = [1,2,3]
>>> a[5]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: list index out of range

[stellaraccident]

Address comments.

mlir/lib/Bindings/Python/PybindUtils.h

237

Indeed. I was mis-remembering that this reports with an actual message.

303

I had added a comment with rationale above the block of functions.

326–328

Done - and just used the underlying C API vs pybind11::slice. It is trivial.

[stellaraccident]

Made some non-trivial changes. Another look?

[stellaraccident]

FYI - the benchmark script: https://gist.github.com/stellaraccident/ba86946ff70ff3084af62b476a05c7c9

[ftynse]

Still LGTM. Some test coverage that we actually raise the errors we expect wouldn't hurt.

[stellaraccident]

Add exception test.

mlir/lib/Bindings/Python/PybindUtils.h

 ///   - an instance method `intptr_t getNumElements()` that returns the number
 ///     of elements in the backing container (NOT that of the slice);
 ///   - an instance method `ElementTy getElement(intptr_t)` that returns a
 ///     single element at the given index.
 ///   - an instance method `Derived slice(intptr_t, intptr_t, intptr_t)` that
 ///     constructs a new instance of the derived pseudo-container with the
 ///     given slice parameters (to be forwarded to the Sliceable constructor).
 ///
+/// The getNumElements() and getElement(intptr_t) callbacks must not throw.
+///
 /// A derived class may additionally define:
 ///   - a `static void bindDerived(ClassTy &)` method to bind additional methods
 ///     the python class.
 template <typename Derived, typename ElementTy>
 class Sliceable {
 protected:
   using ClassTy = pybind11::class_<Derived>;
 
+  // Transforms `index` into a legal value to access the underlying sequence.
+  // Returns <0 on failure.
   intptr_t wrapIndex(intptr_t index) {
     if (index < 0)
       index = length + index;
-    if (index < 0 || index >= length) {
-      throw python::SetPyError(PyExc_IndexError,
-                               "attempt to access out of bounds");
-    }
+    if (index < 0 || index >= length)
+      return -1;
     return index;
   }
 
-public:
-  explicit Sliceable(intptr_t startIndex, intptr_t length, intptr_t step)
-      : startIndex(startIndex), length(length), step(step) {
-    assert(length >= 0 && "expected non-negative slice length");
-  }
-
-  /// Returns the length of the slice.
-  intptr_t dunderLen() const { return length; }
-
   /// Returns the element at the given slice index. Supports negative indices
-  /// by taking elements in inverse order. Throws if the index is out of bounds.
-  ElementTy dunderGetItem(intptr_t index) {
+  /// by taking elements in inverse order. Returns a nullptr object if out
+  /// of bounds.
+  pybind11::object getItem(intptr_t index) {
     // Negative indices mean we count from the end.
     index = wrapIndex(index);
+    if (index < 0) {
+      PyErr_SetString(PyExc_IndexError, "index out of range");
+      return {};
+    }
 
     // Compute the linear index given the current slice properties.
     int linearIndex = index * step + startIndex;
     assert(linearIndex >= 0 &&
            linearIndex < static_cast<Derived *>(this)->getNumElements() &&
            "linear index out of bounds, the slice is ill-formed");
-    return static_cast<Derived *>(this)->getElement(linearIndex);
+    return pybind11::cast(
+        static_cast<Derived *>(this)->getElement(linearIndex));
   }
 
   /// Returns a new instance of the pseudo-container restricted to the given
-  /// slice.
-  Derived dunderGetItemSlice(pybind11::slice slice) {
+  /// slice. Returns a nullptr object on failure.
+  pybind11::object getItemSlice(PyObject *slice) {
     ssize_t start, stop, extraStep, sliceLength;
-    if (!slice.compute(dunderLen(), &start, &stop, &extraStep, &sliceLength)) {
-      throw python::SetPyError(PyExc_IndexError,
-                               "attempt to access out of bounds");
-    }
-    return static_cast<Derived *>(this)->slice(startIndex + start * step,
-                                               sliceLength, step * extraStep);
+    if (PySlice_GetIndicesEx(slice, length, &start, &stop, &extraStep,
+                             &sliceLength) != 0) {
+      PyErr_SetString(PyExc_IndexError, "index out of range");
+      return {};
+    }
+    return pybind11::cast(static_cast<Derived *>(this)->slice(
+        startIndex + start * step, sliceLength, step * extraStep));
+
+public:
+  explicit Sliceable(intptr_t startIndex, intptr_t length, intptr_t step)
+      : startIndex(startIndex), length(length), step(step) {
+    assert(length >= 0 && "expected non-negative slice length");
   }
 
   /// Returns a new vector (mapped to Python list) containing elements from two
   /// slices. The new vector is necessary because slices may not be contiguous
   /// or even come from the same original sequence.
   std::vector<ElementTy> dunderAdd(Derived &other) {
     std::vector<ElementTy> elements;
     elements.reserve(length + other.length);
     for (intptr_t i = 0; i < length; ++i) {
-      elements.push_back(dunderGetItem(i));
+      elements.push_back(static_cast<Derived *>(this)->getElement(i));
     }
     for (intptr_t i = 0; i < other.length; ++i) {
-      elements.push_back(other.dunderGetItem(i));
+      elements.push_back(static_cast<Derived *>(this)->getElement(i));
     }
     return elements;
   }
 
   /// Binds the indexing and length methods in the Python class.
   static void bind(pybind11::module &m) {
     auto clazz = pybind11::class_<Derived>(m, Derived::pyClassName,
                                            pybind11::module_local())
-                     .def("__len__", &Sliceable::dunderLen)
-                     .def("__getitem__", &Sliceable::dunderGetItem)
-                     .def("__getitem__", &Sliceable::dunderGetItemSlice)
                      .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 pybind11, 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/pybind11/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 = pybind11::cast<Derived *>(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 = pybind11::cast<Derived *>(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 = pybind11::cast<Derived *>(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.
   static void bindDerived(ClassTy &) {}
 
 private:
   intptr_t startIndex;
   intptr_t length;
   intptr_t step;
 };
 
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_PYBINDUTILS_H

mlir/test/python/ir/operation.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 import io
 import itertools
 from mlir.ir import *
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
   return f
 
 
+def expect_index_error(callback):
+  try:
+    _ = callback()
+    raise RuntimeError("Expected IndexError")
+  except IndexError:
+    pass
+
+
 # Verify iterator based traversal of the op/region/block hierarchy.
 # CHECK-LABEL: TEST: testTraverseOpRegionBlockIterators
 @run
 def testTraverseOpRegionBlockIterators():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(
       r"""
for res in call.results:
     print(f"Result {res.result_number}, type {res.type}")
 
   # CHECK: Result type i32
   # CHECK: Result type f64
   # CHECK: Result type index
   for t in call.results.types:
     print(f"Result type {t}")
 
-
+  # Out of range
+  expect_index_error(lambda: call.results[3])
+  expect_index_error(lambda: call.results[-4])
 
 
 # CHECK-LABEL: TEST: testOperationResultListSlice
 @run
 def testOperationResultListSlice():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
with Context() as ctx:
 
     # CHECK: Result 3, type i4
     # CHECK: Result 1, type i2
     inverted_middle = producer.results[-2:0:-2]
     for res in inverted_middle:
       print(f"Result {res.result_number}, type {res.type}")
 
 
-
-
 # CHECK-LABEL: TEST: testOperationAttributes
 @run
 def testOperationAttributes():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(
       r"""
     "some.op"() { some.attribute = 1 : i8,