Add MLIR Python binding for Array Attribute

[ftynse]

mlir/lib/Bindings/Python/IRModules.cpp

1509

If you make it inherit Sliceable, it should be relatively easy to get Python array slicing behavior here.

1526–1527

Aren't attributes owned by the context anyway?

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

1509

I gave it a try, but that requires multiple-inheritance and it breaks down because PyConcreteAttribute wants the derived class to be copy-constructible, at which point I don't know how to preserve the slice.

1526–1527

Good point!

[mehdi_amini]

Add exception handling for invalid type in ArrayAttr ctor

[mehdi_amini]

Here is the diff of my attempt to add Sliceable support:

diff --git a/mlir/lib/Bindings/Python/IRModules.cpp b/mlir/lib/Bindings/Python/IRModules.cpp
index d30c1b333807..65c39112d2fc 100644
--- a/mlir/lib/Bindings/Python/IRModules.cpp
+++ b/mlir/lib/Bindings/Python/IRModules.cpp
@@ -1461,11 +1461,33 @@ public:
   static void bindDerived(ClassTy &m) {}
 };
 
-class PyArrayAttribute : public PyConcreteAttribute<PyArrayAttribute> {
+class PyArrayAttribute : public PyConcreteAttribute<PyArrayAttribute>, public Sliceable<PyArrayAttribute, PyAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAArray;
   static constexpr const char *pyClassName = "ArrayAttr";
-  using PyConcreteAttribute::PyConcreteAttribute;
+  using ClassTy = PyConcreteAttribute::ClassTy;
+
+  PyArrayAttribute(PyMlirContextRef contextRef,  MlirAttribute arr, intptr_t startIndex = 0,
+                  intptr_t length = -1, intptr_t step = 1)
+      : PyConcreteAttribute(contextRef, arr), Sliceable(startIndex,
+                  length == -1 ? mlirArrayAttrGetNumElements(arr)
+                               : length,
+                  step) {}
+
+  intptr_t getNumElements() {
+    return mlirArrayAttrGetNumElements(*this);
+  }
+
+  PyAttribute getElement(intptr_t pos) {
+    if (pos >= getNumElements())
+      throw py::index_error("ArrayAttribute index out of range");
+    return PyAttribute(getContext(), mlirArrayAttrGetElement(*this, pos));
+  }
+
+  PyArrayAttribute slice(intptr_t startIndex, intptr_t length, intptr_t step) {
+    return PyArrayAttribute(getContext(), *this, startIndex, length, step);
+  }
+
 
   class PyArrayAttributeIterator {
   public:
@@ -1492,6 +1514,12 @@ public:
     int nextIndex = 0;
   };
 
+  static void bind(py::module &m)  {
+     Sliceable::bind(m);
+     PyConcreteAttribute::bind(m);
+  }
+
+  static void bindDerived(Sliceable::ClassTy &) {}
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
diff --git a/mlir/test/Bindings/Python/ir_attributes.py b/mlir/test/Bindings/Python/ir_attributes.py
index 642c1f6a836c..99e2dfd60595 100644
--- a/mlir/test/Bindings/Python/ir_attributes.py
+++ b/mlir/test/Bindings/Python/ir_attributes.py
@@ -276,6 +276,8 @@ def testArrayAttr():
   with Context():
     raw = Attribute.parse("[42, true, vector<4xf32>]")
   # CHECK: attr: [42, true, vector<4xf32>]
+  print("raw attr:", raw[0:1])
+  # CHECK: attr: [42, true]
   print("raw attr:", raw)
   # CHECK: - 42
   # CHECK: - true

[stellaraccident]

Looks like Alex has taken a first stab at the review so I'll defer to him. Lgtm though.

As a preference, I would rather duplicate slicing behavior (or have utility functions) vs more complicated inheritance.

[mehdi_amini]

Actually it isn't clear to me that Attributes subclasses like ArrayAttribute can have data members.

For example when parsing: Attribute.parse("[42, true, vector<4xf32>]") we will allow to cast this to an ArrayAttr and the only "isa" method available is based on the underlying mlirAttribute member.
We could handle the slice in another wrapper class, but the slice would not be a subclass of PyConcreteAttribute and I'm not sure it'd play well with everything else.

[ftynse]

mlir/lib/Bindings/Python/IRModules.cpp

1509

I see. I am totally fine leaving this for the future and/or having extra code, potentially with helpers, like @stellaraccident suggests.

1526–1527

I still see the keep_alive, forgot to upload?

[mehdi_amini]

Update removing the keep_alive

[mehdi_amini]

Remove keep_alive (for real!)

mlir/lib/Bindings/Python/IRModules.cpp

static void bind(py::module &m) {
     cls.def(py::init<PyAttribute &>(), py::keep_alive<0, 1>());
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
+class PyArrayAttribute : public PyConcreteAttribute<PyArrayAttribute> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAArray;
+  static constexpr const char *pyClassName = "ArrayAttr";
+  using PyConcreteAttribute::PyConcreteAttribute;
+
+  class PyArrayAttributeIterator {
+  public:
+    PyArrayAttributeIterator(PyAttribute attr) : attr(attr) {}
+
+    PyArrayAttributeIterator &dunderIter() { return *this; }
+
+    PyAttribute dunderNext() {
+      if (nextIndex >= mlirArrayAttrGetNumElements(attr.get())) {
+        throw py::stop_iteration();
+      }
+      return PyAttribute(attr.getContext(),
+                         mlirArrayAttrGetElement(attr.get(), nextIndex++));
+    }
+
+    static void bind(py::module &m) {
+      py::class_<PyArrayAttributeIterator>(m, "ArrayAttributeIterator")
+          .def("__iter__", &PyArrayAttributeIterator::dunderIter)
+          .def("__next__", &PyArrayAttributeIterator::dunderNext);
+    }
+
+  private:
+    PyAttribute attr;
+    int nextIndex = 0;
+  };
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](py::list attributes, DefaultingPyMlirContext context) {
+          SmallVector<MlirAttribute> mlirAttributes;
+          mlirAttributes.reserve(py::len(attributes));
+          for (auto attribute : attributes) {
+            try {
+              mlirAttributes.push_back(attribute.cast<PyAttribute>());
+            } catch (py::cast_error &err) {
+              std::string msg = std::string("Invalid attribute when attempting "
+                                            "to create an ArrayAttribute (") +
+                                err.what() + ")";
+              throw py::cast_error(msg);
+            } catch (py::reference_cast_error &err) {
+              // This exception seems thrown when the value is "None".
+              std::string msg =
+                  std::string("Invalid attribute (None?) when attempting to "
+                              "create an ArrayAttribute (") +
+                  err.what() + ")";
+              throw py::cast_error(msg);
+            }
+          }
+          MlirAttribute attr = mlirArrayAttrGet(
+              context->get(), mlirAttributes.size(), mlirAttributes.data());
+          return PyArrayAttribute(context->getRef(), attr);
+        },
+        py::arg("attributes"), py::arg("context") = py::none(),
+        "Gets a uniqued Array attribute");
+    c.def("__getitem__",
+          [](PyArrayAttribute &arr, intptr_t i) {
+            if (i >= mlirArrayAttrGetNumElements(arr))
+              throw py::index_error("ArrayAttribute index out of range");
+            return PyAttribute(arr.getContext(),
+                               mlirArrayAttrGetElement(arr, i));
+          })
+        .def("__len__",
+             [](const PyArrayAttribute &arr) {
+               return mlirArrayAttrGetNumElements(arr);
+             })
+        .def(
+            "__iter__",
+            [](const PyArrayAttribute &arr) {
+              return PyArrayAttributeIterator(arr);
+            },
+            py::keep_alive<
+                0,
+                1>() /* Essential: keep object alive while iterator exists */);
+  }
+};
+
 /// Float Point Attribute subclass - FloatAttr.
 class PyFloatAttribute : public PyConcreteAttribute<PyFloatAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAFloat;
   static constexpr const char *pyClassName = "FloatAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
 
   static void bindDerived(ClassTy &c) {
py::class_<PyNamedAttribute>(m, "NamedAttribute")
                 mlirAttributeGetContext(self.namedAttr.attribute));
             return PyAttribute(std::move(contextRef), self.namedAttr.attribute);
           },
           py::keep_alive<0, 1>(),
           "The underlying generic attribute of the NamedAttribute binding");
 
   // Builtin attribute bindings.
   PyFloatAttribute::bind(m);
+  PyArrayAttribute::bind(m);
+  PyArrayAttribute::PyArrayAttributeIterator::bind(m);
   PyIntegerAttribute::bind(m);
   PyBoolAttribute::bind(m);
   PyStringAttribute::bind(m);
   PyDenseElementsAttribute::bind(m);
   PyDenseIntElementsAttribute::bind(m);
   PyDenseFPElementsAttribute::bind(m);
   PyTypeAttribute::bind(m);
   PyUnitAttribute::bind(m);

mlir/test/Bindings/Python/ir_attributes.py

with Context():
     # CHECK: attr: vector<4xf32>
     print("attr:", raw)
     type_attr = TypeAttr(raw)
     # CHECK: f32
     print(ShapedType(type_attr.value).element_type)
 
 
 run(testTypeAttr)
+
+
+# CHECK-LABEL: TEST: testArrayAttr
+def testArrayAttr():
+  with Context():
+    raw = Attribute.parse("[42, true, vector<4xf32>]")
+  # CHECK: attr: [42, true, vector<4xf32>]
+  print("raw attr:", raw)
+  # CHECK: - 42
+  # CHECK: - true
+  # CHECK: - vector<4xf32>
+  for attr in ArrayAttr(raw):
+    print("- ", attr)
+
+  with Context():
+    intAttr = Attribute.parse("42")
+    vecAttr = Attribute.parse("vector<4xf32>")
+    boolAttr = BoolAttr.get(True)
+    raw = ArrayAttr.get([vecAttr, boolAttr, intAttr])
+  # CHECK: attr: [vector<4xf32>, true, 42]
+  print("raw attr:", raw)
+  # CHECK: - vector<4xf32>
+  # CHECK: - true
+  # CHECK: - 42
+  arr = ArrayAttr(raw)
+  for attr in arr:
+    print("- ", attr)
+  # CHECK: attr[0]: vector<4xf32>
+  print("attr[0]:", arr[0])
+  # CHECK: attr[1]: true
+  print("attr[1]:", arr[1])
+  # CHECK: attr[2]: 42
+  print("attr[2]:", arr[2])
+  try:
+    print("attr[3]:", arr[3])
+  except IndexError as e:
+    # CHECK: Error: ArrayAttribute index out of range
+    print("Error: ", e)
+  with Context():
+    try:
+      ArrayAttr.get([None])
+    except RuntimeError as e:
+      # CHECK: Error: Invalid attribute (None?) when attempting to create an ArrayAttribute
+      print("Error: ", e)
+    try:
+      ArrayAttr.get([42])
+    except RuntimeError as e:
+      # CHECK: Error: Invalid attribute when attempting to create an ArrayAttribute (Unable to cast Python instance of type <class 'int'> to C++ type 'mlir::python::PyAttribute')
+      print("Error: ", e)
+run(testArrayAttr)
+