[mlir][python] Usability improvements for Python bindings

[ftynse]

Summary

Provide a couple of quality-of-life usability improvements for Python bindings,
in particular:

• give access to the list of types for the list of op results or block arguments, similarly to ValueRange->TypeRange,
• allow for constructing empty dictionary arrays,
• support construction of array attributes by concatenating an existing attribute with a Python list of attributes.

All these are required for the upcoming customization of builtin and standard
ops.

[stellaraccident]

Awesome, thanks!

[ftynse]

Rebase.

mlir/lib/Bindings/Python/IRAttributes.cpp

 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
-using llvm::StringRef;
 using llvm::Twine;
 
 namespace {
 
 static MlirStringRef toMlirStringRef(const std::string &s) {
   return mlirStringRefCreate(s.data(), s.size());
 }
 
c.def_static(
         [](PyAffineMap &affineMap) {
           MlirAttribute attr = mlirAffineMapAttrGet(affineMap.get());
           return PyAffineMapAttribute(affineMap.getContext(), attr);
         },
         py::arg("affine_map"), "Gets an attribute wrapping an AffineMap.");
   }
 };
 
+template <typename T>
+static T pyTryCast(py::handle object) {
+  try {
+    return object.cast<T>();
+  } 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) {
+    std::string msg = std::string("Invalid attribute (None?) when attempting "
+                                  "to create an ArrayAttribute (") +
+                      err.what() + ")";
+    throw py::cast_error(msg);
+  }
+}
+
 class PyArrayAttribute : public PyConcreteAttribute<PyArrayAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAArray;
   static constexpr const char *pyClassName = "ArrayAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
 
   class PyArrayAttributeIterator {
   public:
static void bind(py::module &m) {
           .def("__next__", &PyArrayAttributeIterator::dunderNext);
     }
 
   private:
     PyAttribute attr;
     int nextIndex = 0;
   };
 
+  PyAttribute getItem(intptr_t i) {
+    return PyAttribute(getContext(), mlirArrayAttrGetElement(*this, i));
+  }
+
   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);
-            }
+            mlirAttributes.push_back(pyTryCast<PyAttribute>(attribute));
           }
           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));
+            return arr.getItem(i);
           })
         .def("__len__",
              [](const PyArrayAttribute &arr) {
                return mlirArrayAttrGetNumElements(arr);
              })
         .def("__iter__", [](const PyArrayAttribute &arr) {
           return PyArrayAttributeIterator(arr);
         });
+    c.def("__add__", [](PyArrayAttribute arr, py::list extras) {
+      std::vector<MlirAttribute> attributes;
+      intptr_t numOldElements = mlirArrayAttrGetNumElements(arr);
+      attributes.reserve(numOldElements + py::len(extras));
+      for (intptr_t i = 0; i < numOldElements; ++i)
+        attributes.push_back(arr.getItem(i));
+      for (py::handle attr : extras)
+        attributes.push_back(pyTryCast<PyAttribute>(attr));
+      MlirAttribute arrayAttr = mlirArrayAttrGet(
+          arr.getContext()->get(), attributes.size(), attributes.data());
+      return PyArrayAttribute(arr.getContext(), arrayAttr);
+    });
   }
 };
 
 /// Float Point Attribute subclass - FloatAttr.
 class PyFloatAttribute : public PyConcreteAttribute<PyFloatAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAFloat;
   static constexpr const char *pyClassName = "FloatAttr";
c.def_static(
                                   toMlirStringRef(name)),
                 mlir_attr));
           }
           MlirAttribute attr =
               mlirDictionaryAttrGet(context->get(), mlirNamedAttributes.size(),
                                     mlirNamedAttributes.data());
           return PyDictAttribute(context->getRef(), attr);
         },
-        py::arg("value"), py::arg("context") = py::none(),
+        py::arg("value") = py::dict(), py::arg("context") = py::none(),
         "Gets an uniqued dict attribute");
     c.def("__getitem__", [](PyDictAttribute &self, const std::string &name) {
       MlirAttribute attr =
           mlirDictionaryAttrGetElementByName(self, toMlirStringRef(name));
       if (mlirAttributeIsNull(attr)) {
         throw SetPyError(PyExc_KeyError,
                          "attempt to access a non-existent attribute");
       }

mlir/lib/Bindings/Python/IRCore.cpp

c.def_property_readonly("owner", [](PyOpResult &self) {
       return self.getParentOperation().getObject();
     });
     c.def_property_readonly("result_number", [](PyOpResult &self) {
       return mlirOpResultGetResultNumber(self.get());
     });
   }
 };
 
+/// 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.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";
public:
   }
 
   /// 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
PyOpResult getElement(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:

mlir/test/python/ir/attributes.py

with Context():
 
     try:
       _ = a[42]
     except IndexError:
       pass
     else:
       assert False, "expected IndexError on accessing an out-of-bounds attribute"
 
+    # CHECK "empty: {}"
+    print("empty: ", DictAttr.get())
+
 
 # CHECK-LABEL: TEST: testTypeAttr
 @run
 def testTypeAttr():
   with Context():
     raw = Attribute.parse("vector<4xf32>")
     # CHECK: attr: vector<4xf32>
     print("attr:", raw)
with Context():
     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
       print("Error: ", e)
+
+  with Context():
+    array = ArrayAttr.get([StringAttr.get("a"), StringAttr.get("b")])
+    array = array + [StringAttr.get("c")]
+    # CHECK: concat: ["a", "b", "c"]
+    print("concat: ", array)

mlir/test/python/ir/operation.py

with Context() as ctx:
     for arg in entry_block.arguments[1:]:
       print(f"Argument {arg.arg_number}, type {arg.type}")
 
     # Check that we can concatenate slices of argument lists.
     # CHECK: Length: 4
     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)
+
 
 run(testBlockArgumentList)
 
 
 # CHECK-LABEL: TEST: testOperationOperands
 def testOperationOperands():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True
def testOperationResultList():
   call = caller.regions[0].blocks[0].operations[0]
   assert len(call.results) == 3
   # CHECK: Result 0, type i32
   # CHECK: Result 1, type f64
   # CHECK: Result 2, type index
   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}")
+
 
 run(testOperationResultList)
 
 
 # CHECK-LABEL: TEST: testOperationResultListSlice
 def testOperationResultListSlice():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True