[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

1904

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.

1910–1935

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

1911

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

1938–1939

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

1913

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

1920

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

[kweisamx]

Thanks for advice. I refine it.

[kweisamx]

mlir/lib/Bindings/Python/IRModules.cpp

1913

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

1913

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

1913

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

1913

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

1933

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); }
+
+  PyAttribute dunderGetItemNamed(const std::string &name,
+                                 DefaultingPyMlirContext context) {
+    MlirAttribute attr =
+        mlirDictionaryAttrGetElementByName(*this, toMlirStringRef(name));
+    if (mlirAttributeIsNull(attr)) {
+      throw SetPyError(PyExc_KeyError,
+                       "attempt to access a non-existent attribute");
+    }
+    return PyAttribute(context->getRef(), attr);
+  }
+
+  static void bindDerived(ClassTy &c) {
+    c.def("__len__", &PyDictAttribute::dunderLen);
+    c.def_static(
+        "get",
+        [](llvm::Optional<py::dict> attributes,
+           DefaultingPyMlirContext context) {
+          std::vector<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.c_str()),
+                                        mlir_attr.get());
+          }
+          if (!mlirAttributes.empty()) {
+            std::vector<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::dunderGetItemNamed, py::arg("name"),
+          py::arg("context") = py::none());
+  }
+};
+
 /// 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) {
   // Builtin attribute bindings.
   PyFloatAttribute::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: integerattr: 42
+    print(a['integerattr'])
+
+    # CHECK: stringattr: string
+    print(a['stringattr'])
+
+
+
+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
     print(ShapedType(type_attr.value).element_type)
 
 
 run(testTypeAttr)