[mlir][python] Fix issues with block argument slices

[ftynse]

Summary

The type extraction helper function for block argument and op result
list objects was ignoring the slice entirely. So was the slice addition.
Both are caused by a misleading naming convention to implement slices
via CRTP. Make the convention more explicit and hide the helper
functions so users have harder time calling them directly.

Closes #56540.

[stellaraccident]

Good catch, and thank you!

[jpienaar]

mlir/lib/Bindings/Python/IRCore.cpp

1972

Could this be a range based loop?

[ftynse]

mlir/lib/Bindings/Python/IRCore.cpp

1972

There are no iterators in the Container class. We can do auto i : llvm::seq<int>(0, container.size())) but I am not convinced that improves readability.

[ingomueller-net]

Awesome, thanks! I confirm that this makes the code run that I had initially problems with!

mlir/lib/Bindings/Python/IRAffine.cpp

public:
 
   PyAffineMapExprList(const PyAffineMap &map, intptr_t startIndex = 0,
                       intptr_t length = -1, intptr_t step = 1)
       : Sliceable(startIndex,
                   length == -1 ? mlirAffineMapGetNumResults(map) : length,
                   step),
         affineMap(map) {}
 
-  intptr_t getNumElements() { return mlirAffineMapGetNumResults(affineMap); }
+private:
+  /// Give the parent CRTP class access to hook implementations below.
+  friend class Sliceable<PyAffineMapExprList, PyAffineExpr>;
+
+  intptr_t getRawNumElements() { return mlirAffineMapGetNumResults(affineMap); }
 
-  PyAffineExpr getElement(intptr_t pos) {
+  PyAffineExpr getRawElement(intptr_t pos) {
     return PyAffineExpr(affineMap.getContext(),
                         mlirAffineMapGetResult(affineMap, pos));
   }
 
   PyAffineMapExprList slice(intptr_t startIndex, intptr_t length,
                             intptr_t step) {
     return PyAffineMapExprList(affineMap, startIndex, length, step);
   }
 
-private:
   PyAffineMap affineMap;
 };
 } // namespace
 
 bool PyAffineMap::operator==(const PyAffineMap &other) {
   return mlirAffineMapEqual(affineMap, other.affineMap);
 }
 
public:
 
   PyIntegerSetConstraintList(const PyIntegerSet &set, intptr_t startIndex = 0,
                              intptr_t length = -1, intptr_t step = 1)
       : Sliceable(startIndex,
                   length == -1 ? mlirIntegerSetGetNumConstraints(set) : length,
                   step),
         set(set) {}
 
-  intptr_t getNumElements() { return mlirIntegerSetGetNumConstraints(set); }
+private:
+  /// Give the parent CRTP class access to hook implementations below.
+  friend class Sliceable<PyIntegerSetConstraintList, PyIntegerSetConstraint>;
+
+  intptr_t getRawNumElements() { return mlirIntegerSetGetNumConstraints(set); }
 
-  PyIntegerSetConstraint getElement(intptr_t pos) {
+  PyIntegerSetConstraint getRawElement(intptr_t pos) {
     return PyIntegerSetConstraint(set, pos);
   }
 
   PyIntegerSetConstraintList slice(intptr_t startIndex, intptr_t length,
                                    intptr_t step) {
     return PyIntegerSetConstraintList(set, startIndex, length, step);
   }
 
-private:
   PyIntegerSet set;
 };
 } // namespace
 
 bool PyIntegerSet::operator==(const PyIntegerSet &other) {
   return mlirIntegerSetEqual(integerSet, other.integerSet);
 }
 

mlir/lib/Bindings/Python/IRCore.cpp

public:
   }
 };
 
 /// Returns the list of types of the values held by container.
 template <typename Container>
 static std::vector<PyType> getValueTypes(Container &container,
                                          PyMlirContextRef &context) {
   std::vector<PyType> result;
-  result.reserve(container.getNumElements());
-  for (int i = 0, e = container.getNumElements(); i < e; ++i) {
+  result.reserve(container.size());
+  for (int i = 0, e = container.size(); i < e; ++i) {
     result.push_back(
         PyType(context, mlirValueGetType(container.getElement(i).get())));
   }
   return result;
 }
 
 /// A list of block arguments. Internally, these are stored as consecutive
 /// elements, random access is cheap. The argument list is associated with the
 /// operation that contains the block (detached blocks are not allowed in
 /// Python bindings) and extends its lifetime.
 class PyBlockArgumentList
     : public Sliceable<PyBlockArgumentList, PyBlockArgument> {
 public:
   static constexpr const char *pyClassName = "BlockArgumentList";
 
   PyBlockArgumentList(PyOperationRef operation, MlirBlock block,
                       intptr_t startIndex = 0, intptr_t length = -1,
                       intptr_t step = 1)
       : Sliceable(startIndex,
                   length == -1 ? mlirBlockGetNumArguments(block) : length,
                   step),
         operation(std::move(operation)), block(block) {}
 
+  static void bindDerived(ClassTy &c) {
+    c.def_property_readonly("types", [](PyBlockArgumentList &self) {
+      return getValueTypes(self, self.operation->getContext());
+    });
+  }
+
+private:
+  /// Give the parent CRTP class access to hook implementations below.
+  friend class Sliceable<PyBlockArgumentList, PyBlockArgument>;
+
   /// Returns the number of arguments in the list.
-  intptr_t getNumElements() {
+  intptr_t getRawNumElements() {
     operation->checkValid();
     return mlirBlockGetNumArguments(block);
   }
 
-  /// Returns `pos`-the element in the list. Asserts on out-of-bounds.
-  PyBlockArgument getElement(intptr_t pos) {
+  /// Returns `pos`-the element in the list.
+  PyBlockArgument getRawElement(intptr_t pos) {
     MlirValue argument = mlirBlockGetArgument(block, pos);
     return PyBlockArgument(operation, argument);
   }
 
   /// Returns a sublist of this list.
   PyBlockArgumentList slice(intptr_t startIndex, intptr_t length,
                             intptr_t step) {
     return PyBlockArgumentList(operation, block, startIndex, length, step);
   }
 
-  static void bindDerived(ClassTy &c) {
-    c.def_property_readonly("types", [](PyBlockArgumentList &self) {
-      return getValueTypes(self, self.operation->getContext());
-    });
-  }
-
-private:
   PyOperationRef operation;
   MlirBlock block;
 };
 
 /// A list of operation operands. Internally, these are stored as consecutive
 /// elements, random access is cheap. The result list is associated with the
 /// operation whose results these are, and extends the lifetime of this
 /// operation.
 class PyOpOperandList : public Sliceable<PyOpOperandList, PyValue> {
 public:
   static constexpr const char *pyClassName = "OpOperandList";
 
   PyOpOperandList(PyOperationRef operation, intptr_t startIndex = 0,
                   intptr_t length = -1, intptr_t step = 1)
       : Sliceable(startIndex,
                   length == -1 ? mlirOperationGetNumOperands(operation->get())
                                : length,
                   step),
         operation(operation) {}
 
-  intptr_t getNumElements() {
+  void dunderSetItem(intptr_t index, PyValue value) {
+    index = wrapIndex(index);
+    mlirOperationSetOperand(operation->get(), index, value.get());
+  }
+
+  static void bindDerived(ClassTy &c) {
+    c.def("__setitem__", &PyOpOperandList::dunderSetItem);
+  }
+
+private:
+  /// Give the parent CRTP class access to hook implementations below.
+  friend class Sliceable<PyOpOperandList, PyValue>;
+
+  intptr_t getRawNumElements() {
     operation->checkValid();
     return mlirOperationGetNumOperands(operation->get());
   }
 
-  PyValue getElement(intptr_t pos) {
+  PyValue getRawElement(intptr_t pos) {
     MlirValue operand = mlirOperationGetOperand(operation->get(), pos);
     MlirOperation owner;
     if (mlirValueIsAOpResult(operand))
       owner = mlirOpResultGetOwner(operand);
     else if (mlirValueIsABlockArgument(operand))
       owner = mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(operand));
     else
       assert(false && "Value must be an block arg or op result.");
     PyOperationRef pyOwner =
         PyOperation::forOperation(operation->getContext(), owner);
     return PyValue(pyOwner, operand);
   }
 
   PyOpOperandList slice(intptr_t startIndex, intptr_t length, intptr_t step) {
     return PyOpOperandList(operation, startIndex, length, step);
   }
 
-  void dunderSetItem(intptr_t index, PyValue value) {
-    index = wrapIndex(index);
-    mlirOperationSetOperand(operation->get(), index, value.get());
-  }
-
-  static void bindDerived(ClassTy &c) {
-    c.def("__setitem__", &PyOpOperandList::dunderSetItem);
-  }
-
-private:
   PyOperationRef operation;
 };
 
 /// A list of operation results. Internally, these are stored as consecutive
 /// elements, random access is cheap. The result list is associated with the
 /// operation whose results these are, and extends the lifetime of this
 /// operation.
 class PyOpResultList : public Sliceable<PyOpResultList, PyOpResult> {
 public:
   static constexpr const char *pyClassName = "OpResultList";
 
   PyOpResultList(PyOperationRef operation, intptr_t startIndex = 0,
                  intptr_t length = -1, intptr_t step = 1)
       : Sliceable(startIndex,
                   length == -1 ? mlirOperationGetNumResults(operation->get())
                                : length,
                   step),
         operation(operation) {}
 
-  intptr_t getNumElements() {
+  static void bindDerived(ClassTy &c) {
+    c.def_property_readonly("types", [](PyOpResultList &self) {
+      return getValueTypes(self, self.operation->getContext());
+    });
+  }
+
+private:
+  /// Give the parent CRTP class access to hook implementations below.
+  friend class Sliceable<PyOpResultList, PyOpResult>;
+
+  intptr_t getRawNumElements() {
     operation->checkValid();
     return mlirOperationGetNumResults(operation->get());
   }
 
-  PyOpResult getElement(intptr_t index) {
+  PyOpResult getRawElement(intptr_t index) {
     PyValue value(operation, mlirOperationGetResult(operation->get(), index));
     return PyOpResult(value);
   }
 
   PyOpResultList slice(intptr_t startIndex, intptr_t length, intptr_t step) {
     return PyOpResultList(operation, startIndex, length, step);
   }
 
-  static void bindDerived(ClassTy &c) {
-    c.def_property_readonly("types", [](PyOpResultList &self) {
-      return getValueTypes(self, self.operation->getContext());
-    });
-  }
-
-private:
   PyOperationRef operation;
 };
 
 /// A list of operation attributes. Can be indexed by name, producing
 /// attributes, or by index, producing named attributes.
 class PyOpAttributeMap {
 public:
   PyOpAttributeMap(PyOperationRef operation)

mlir/lib/Bindings/Python/PybindUtils.h

 /// This is intended for pseudo-containers that can refer to arbitrary slices of
 /// underlying storage indexed by a single integer. Indexing those with an
 /// integer produces an instance of ElementTy. Indexing those with a slice
 /// produces a new instance of Derived, which can be sliced further.
 ///
 /// A derived class must provide the following:
 ///   - a `static const char *pyClassName ` field containing the name of the
 ///     Python class to bind;
-///   - an instance method `intptr_t getNumElements()` that returns the number
+///   - an instance method `intptr_t getRawNumElements()` 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 `ElementTy getRawElement(intptr_t)` that returns a
+///     single element at the given linear index (NOT slice 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.
+/// The getRawNumElements() and getRawElement(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.
+  /// 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)
       return -1;
     return index;
   }
 
+  /// Computes the linear index given the current slice properties.
+  intptr_t linearizeIndex(intptr_t index) {
+    intptr_t linearIndex = index * step + startIndex;
+    assert(linearIndex >= 0 &&
+           linearIndex < static_cast<Derived *>(this)->getRawNumElements() &&
+           "linear index out of bounds, the slice is ill-formed");
+    return linearIndex;
+  }
+
   /// Returns the element at the given slice index. Supports negative indices
   /// 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 pybind11::cast(
-        static_cast<Derived *>(this)->getElement(linearIndex));
+        static_cast<Derived *>(this)->getRawElement(linearizeIndex(index)));
   }
 
   /// Returns a new instance of the pseudo-container restricted to the given
   /// slice. Returns a nullptr object on failure.
   pybind11::object getItemSlice(PyObject *slice) {
     ssize_t start, stop, extraStep, sliceLength;
     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 the `index`-th element in the slice, supports negative indices.
+  /// Throws if the index is out of bounds.
+  ElementTy getElement(intptr_t index) {
+    // Negative indices mean we count from the end.
+    index = wrapIndex(index);
+    if (index < 0) {
+      throw pybind11::index_error("index out of range");
+    }
+
+    return static_cast<Derived *>(this)->getRawElement(linearizeIndex(index));
+  }
+
+  /// Returns the size of slice.
+  intptr_t size() { return 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(static_cast<Derived *>(this)->getElement(i));
     }
     for (intptr_t i = 0; i < other.length; ++i) {
-      elements.push_back(static_cast<Derived *>(this)->getElement(i));
+      elements.push_back(static_cast<Derived *>(&other)->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())

mlir/test/python/ir/operation.py

print("Length: ",
           len(entry_block.arguments[:2] + entry_block.arguments[1:]))
 
     # CHECK: Type: i8
     # CHECK: Type: i16
     # CHECK: Type: i24
     for t in entry_block.arguments.types:
       print("Type: ", t)
 
+    # Check that slicing and type access compose.
+    # CHECK: Sliced type: i16
+    # CHECK: Sliced type: i24
+    for t in entry_block.arguments[1:].types:
+      print("Sliced type: ", t)
+
+    # Check that slice addition works as expected.
+    # CHECK: Argument 2, type i24
+    # CHECK: Argument 0, type i8
+    restructured = entry_block.arguments[-1:] + entry_block.arguments[:1]
+    for arg in restructured:
+      print(f"Argument {arg.arg_number}, type {arg.type}")
+
 
 # CHECK-LABEL: TEST: testOperationOperands
 @run
 def testOperationOperands():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func.func @f1(%arg0: i32) {