[mlir] Add Python binding for MLIR Dict Attribute

[kweisamx]

[mlir] Add Python binding for MLIR Dict Attribute

[ftynse]

Thanks!

mlir/lib/Bindings/Python/IRModules.cpp

1987

This should be getting a context of _this_ attribute instead of accepting any context as argument. PyConcreteAttribute has a field containing the current attribute IIRC.

1993–2018

Can this be defined as a member function rather than a long lambda?

1994

Do we really need to make this optional? Creating an empty dictionary attribute from an empty dictionary makes more sense to me.

2021–2022

Ideally, we would have the same interface querying for DictAttr as we do for the attrs field on ops, i.e. also have an indexed getitem.

mlir/test/Bindings/Python/ir_attributes.py

281

Please also add a check that the exception is thrown for non-existent attribute.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

1996

Instead of creating a vector of pair<string,MlirAttribute> and then use it to build a vector of MlirNamedAttribute, can you just build a vector of MlirNamedAttribute in the first place?
(also use SmallVector<MlirNamedAttribute>)

2003

Nit: remove the c_str(), it is only pessimizing here.

[kweisamx]

Thanks for advice. I refine it.

[kweisamx]

mlir/lib/Bindings/Python/IRModules.cpp

1996

In mlirNamedAttributeGet, it use MlirStringRef.
If we build vector with MlirNamedAttribute and pass name directly. It will be change after next iterator(Because we use reference)?
That I build MlirAttribute with two step.

Are there other ways better?

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

1996

Oh, it seems that MlirNamedAttribute are non-owning for the key... This is specific to the C API, the C++ API creates a context-owned string instead. I think we should change the C API.

I'm doing it here: https://reviews.llvm.org/D93133

[ftynse]

mlir/lib/Bindings/Python/IRModules.cpp

1996

Yes, we must have a todo somewhere to fix that, legacy from early iterations of C API.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

1996

Alright this is landed, if you rebase you can eliminate this extra vector now :)

[kweisamx]

Great!
I fixed it with eliminating extra vector .
Please help me review this commit.
Thanks.

[mehdi_amini]

LGTM!

Do you need me to land this for you? (if so please let me know the authorship information: name/email).

mlir/lib/Bindings/Python/IRModules.cpp

2016

Please run git clang-format (arc does it for you I believe)

[kweisamx]

Fix clang-format

[kweisamx]

Yes, i would contribute more for this project :)

name: kweisamx
email: [e-mail]

Thank!

mlir/lib/Bindings/Python/IRModules.cpp

py::int_ dunderGetItem(intptr_t pos) {
     throw SetPyError(PyExc_TypeError, "Unsupported integer type");
   }
 
   static void bindDerived(ClassTy &c) {
     c.def("__getitem__", &PyDenseIntElementsAttribute::dunderGetItem);
   }
 };
 
+class PyDictAttribute : public PyConcreteAttribute<PyDictAttribute> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirAttributeIsADictionary;
+  static constexpr const char *pyClassName = "DictAttr";
+  using PyConcreteAttribute::PyConcreteAttribute;
+
+  intptr_t dunderLen() { return mlirDictionaryAttrGetNumElements(*this); }
+
+  static void bindDerived(ClassTy &c) {
+    c.def("__len__", &PyDictAttribute::dunderLen);
+    c.def_static(
+        "get",
+        [](py::dict attributes, DefaultingPyMlirContext context) {
+          SmallVector<MlirNamedAttribute> mlirNamedAttributes;
+          mlirNamedAttributes.reserve(attributes.size());
+          for (auto &it : attributes) {
+            auto &mlir_attr = it.second.cast<PyAttribute &>();
+            auto name = it.first.cast<std::string>();
+            mlirNamedAttributes.push_back(mlirNamedAttributeGet(
+                mlirIdentifierGet(mlirAttributeGetContext(mlir_attr),
+                                  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(),
+        "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");
+      }
+      return PyAttribute(self.getContext(), attr);
+    });
+    c.def("__getitem__", [](PyDictAttribute &self, intptr_t index) {
+      if (index < 0 || index >= self.dunderLen()) {
+        throw SetPyError(PyExc_IndexError,
+                         "attempt to access out of bounds attribute");
+      }
+      MlirNamedAttribute namedAttr = mlirDictionaryAttrGetElement(self, index);
+      return PyNamedAttribute(
+          namedAttr.attribute,
+          std::string(mlirIdentifierStr(namedAttr.name).data));
+    });
+  }
+};
+
 /// Refinement of PyDenseElementsAttribute for attributes containing
 /// floating-point values. Supports element access.
 class PyDenseFPElementsAttribute
     : public PyConcreteAttribute<PyDenseFPElementsAttribute,
                                  PyDenseElementsAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsADenseFPElements;
   static constexpr const char *pyClassName = "DenseFPElementsAttr";
void mlir::python::populateIRSubmodule(py::module &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);
+  PyDictAttribute::bind(m);
   PyTypeAttribute::bind(m);
   PyUnitAttribute::bind(m);
 
   //----------------------------------------------------------------------------
   // Mapping of PyType.
   //----------------------------------------------------------------------------
   py::class_<PyType>(m, "Type")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyType::getCapsule)

mlir/test/Bindings/Python/ir_attributes.py

with Context():
 
     # CHECK: f32
     print(ShapedType(a.type).element_type)
 
 
 run(testDenseFPAttr)
 
 
+# CHECK-LABEL: TEST: testDictAttr
+def testDictAttr():
+  with Context():
+    dict_attr = {
+      'stringattr':  StringAttr.get('string'),
+      'integerattr' : IntegerAttr.get(
+        IntegerType.get_signless(32), 42)
+    }
+
+    a = DictAttr.get(dict_attr)
+
+    # CHECK attr: {integerattr = 42 : i32, stringattr = "string"}
+    print("attr:", a)
+
+    assert len(a) == 2
+
+    # CHECK: 42 : i32
+    print(a['integerattr'])
+
+    # CHECK: "string"
+    print(a['stringattr'])
+
+    # Check that exceptions are raised as expected.
+    try:
+      _ = a['does_not_exist']
+    except KeyError:
+      pass
+    else:
+      assert False, "Exception not produced"
+
+    try:
+      _ = a[42]
+    except IndexError:
+      pass
+    else:
+      assert False, "expected IndexError on accessing an out-of-bounds attribute"
+
+
+
+run(testDictAttr)
+
 # CHECK-LABEL: TEST: testTypeAttr
 def testTypeAttr():
   with Context():
     raw = Attribute.parse("vector<4xf32>")
     # CHECK: attr: vector<4xf32>
     print("attr:", raw)
     type_attr = TypeAttr(raw)
     # CHECK: f32