[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

1981

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.

1987–2012

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

1988

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

2015–2016

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

1990

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>)

1997

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

[kweisamx]

Thanks for advice. I refine it.

[kweisamx]

mlir/lib/Bindings/Python/IRModules.cpp

1990

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

1990

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

1990

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

1990

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

2010

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

 //===- IRModules.cpp - IR Submodules of pybind module ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
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<std::pair<std::string, MlirAttribute>> mlirAttributes;
+          MlirAttribute attr;
+          mlirAttributes.reserve(attributes.size());
+          for (auto &it : attributes) {
+            auto &mlir_attr = it.second.cast<PyAttribute &>();
+            auto name = it.first.cast<std::string>();
+            mlirAttributes.emplace_back(std::move(name), mlir_attr.get());
+          }
+          if (!mlirAttributes.empty()) {
+            SmallVector<MlirNamedAttribute> mlirNamedAttributes;
+            mlirNamedAttributes.reserve(mlirAttributes.size());
+            for (auto &it : mlirAttributes) {
+              mlirNamedAttributes.push_back(
+                  mlirNamedAttributeGet(toMlirStringRef(it.first), it.second));
+            }
+            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(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