Page MenuHomePhabricator

Add MLIR Python binding for Array Attribute
ClosedPublic

Authored by mehdi_amini on Dec 9 2020, 10:26 AM.

Diff Detail

Event Timeline

mehdi_amini created this revision.Dec 9 2020, 10:26 AM
mehdi_amini requested review of this revision.Dec 9 2020, 10:26 AM
ftynse added inline comments.Dec 9 2020, 2:29 PM
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 added inline comments.Dec 9 2020, 6:05 PM
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!

Add exception handling for invalid type in ArrayAttr ctor

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 accepted this revision.Dec 9 2020, 6:49 PM

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.

This revision is now accepted and ready to land.Dec 9 2020, 6:49 PM

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 added inline comments.Dec 10 2020, 1:56 AM
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?

Update removing the keep_alive

Remove keep_alive (for real!)

ftynse accepted this revision.Dec 10 2020, 11:02 AM
This revision was landed with ongoing or failed builds.Dec 10 2020, 12:51 PM
This revision was automatically updated to reflect the committed changes.