From 7d88773f3d1b96600863c22f386594da52ca7d6e Mon Sep 17 00:00:00 2001 From: Perry Gibson Date: Mon, 13 Oct 2025 14:25:48 +0100 Subject: [PATCH 1/6] feat: add dict-style to python attributes --- mlir/lib/Bindings/Python/IRCore.cpp | 47 ++++++++++++++++++++++++++++- mlir/test/python/ir/operation.py | 29 ++++++++++++++---- 2 files changed, 69 insertions(+), 7 deletions(-) diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index 7b1710656243a..a94674079bf1c 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -2730,6 +2730,16 @@ class PyOpAttributeMap { operation->get(), toMlirStringRef(name))); } + template + auto forEachAttr(F fn) { + intptr_t n = mlirOperationGetNumAttributes(operation->get()); + for (intptr_t i = 0; i < n; ++i) { + MlirNamedAttribute na = mlirOperationGetAttribute(operation->get(), i); + MlirStringRef name = mlirIdentifierStr(na.name); + fn(name, na.attribute); + } + } + static void bind(nb::module_ &m) { nb::class_(m, "OpAttributeMap") .def("__contains__", &PyOpAttributeMap::dunderContains) @@ -2737,7 +2747,42 @@ class PyOpAttributeMap { .def("__getitem__", &PyOpAttributeMap::dunderGetItemNamed) .def("__getitem__", &PyOpAttributeMap::dunderGetItemIndexed) .def("__setitem__", &PyOpAttributeMap::dunderSetItem) - .def("__delitem__", &PyOpAttributeMap::dunderDelItem); + .def("__delitem__", &PyOpAttributeMap::dunderDelItem) + .def("__iter__", + [](PyOpAttributeMap &self) { + nb::list keys; + self.forEachAttr([&](MlirStringRef name, MlirAttribute) { + keys.append(nb::str(name.data, name.length)); + }); + return nb::iter(keys); + }) + .def("keys", + [](PyOpAttributeMap &self) { + nb::list out; + self.forEachAttr([&](MlirStringRef name, MlirAttribute) { + out.append(nb::str(name.data, name.length)); + }); + return out; + }) + .def("values", + [](PyOpAttributeMap &self) { + nb::list out; + self.forEachAttr([&](MlirStringRef, MlirAttribute attr) { + out.append(PyAttribute(self.operation->getContext(), attr) + .maybeDownCast()); + }); + return out; + }) + .def("items", [](PyOpAttributeMap &self) { + nb::list out; + self.forEachAttr([&](MlirStringRef name, MlirAttribute attr) { + out.append( + nb::make_tuple(nb::str(name.data, name.length), + PyAttribute(self.operation->getContext(), attr) + .maybeDownCast())); + }); + return out; + }); } private: diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py index cb4cfc8c8a6ec..9d49cb1d25f9e 100644 --- a/mlir/test/python/ir/operation.py +++ b/mlir/test/python/ir/operation.py @@ -569,14 +569,31 @@ def testOperationAttributes(): # CHECK: Attribute value b'text' print(f"Attribute value {sattr.value_bytes}") + # Python dict-style iteration # We don't know in which order the attributes are stored. - # CHECK-DAG: NamedAttribute(dependent="text") - # CHECK-DAG: NamedAttribute(other.attribute=3.000000e+00 : f64) - # CHECK-DAG: NamedAttribute(some.attribute=1 : i8) - for attr in op.attributes: - print(str(attr)) + # CHECK-DAG: dependent + # CHECK-DAG: other.attribute + # CHECK-DAG: some.attribute + for name in op.attributes: + print(name) + + # Basic dict-like introspection + # CHECK: True + print("some.attribute" in op.attributes) + # CHECK: False + print("missing" in op.attributes) + # CHECK: Keys: ['dependent', 'other.attribute', 'some.attribute'] + print("Keys:", sorted(op.attributes.keys())) + # CHECK: Values count 3 + print("Values count", len(op.attributes.values())) + # CHECK: Items count 3 + print("Items count", len(op.attributes.items())) + + # Dict() conversion test + d = {k: v.value for k, v in dict(op.attributes).items()} + # CHECK: Dict mapping {'dependent': 'text', 'other.attribute': 3.0, 'some.attribute': 1} + print("Dict mapping", d) - # Check that exceptions are raised as expected. try: op.attributes["does_not_exist"] except KeyError: From 827870eb63814ec89a20d8a7efa1c168b15b4455 Mon Sep 17 00:00:00 2001 From: Perry Gibson Date: Mon, 13 Oct 2025 17:23:07 +0100 Subject: [PATCH 2/6] fix: re-add comment --- mlir/test/python/ir/operation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py index 9d49cb1d25f9e..1d4ede1ee336a 100644 --- a/mlir/test/python/ir/operation.py +++ b/mlir/test/python/ir/operation.py @@ -594,6 +594,7 @@ def testOperationAttributes(): # CHECK: Dict mapping {'dependent': 'text', 'other.attribute': 3.0, 'some.attribute': 1} print("Dict mapping", d) + # Check that exceptions are raised as expected. try: op.attributes["does_not_exist"] except KeyError: From 5c0753dc5200452adaff1285274e7364ccfbb222 Mon Sep 17 00:00:00 2001 From: Perry Gibson Date: Mon, 13 Oct 2025 17:32:00 +0100 Subject: [PATCH 3/6] refactor: replace templated forEachAttr with std::function version --- mlir/lib/Bindings/Python/IRCore.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index a94674079bf1c..082b87463160c 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -2730,8 +2730,8 @@ class PyOpAttributeMap { operation->get(), toMlirStringRef(name))); } - template - auto forEachAttr(F fn) { + void + forEachAttr(const std::function &fn) { intptr_t n = mlirOperationGetNumAttributes(operation->get()); for (intptr_t i = 0; i < n; ++i) { MlirNamedAttribute na = mlirOperationGetAttribute(operation->get(), i); From a7d5ca3a5cb254e6e61c89c4259814f91c43f33a Mon Sep 17 00:00:00 2001 From: Perry Gibson Date: Mon, 13 Oct 2025 17:52:09 +0100 Subject: [PATCH 4/6] refactor: replace std::function with llvm::function_ref --- mlir/lib/Bindings/Python/IRCore.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index 082b87463160c..7fc55cb5b870f 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -2730,8 +2730,7 @@ class PyOpAttributeMap { operation->get(), toMlirStringRef(name))); } - void - forEachAttr(const std::function &fn) { + void forEachAttr(llvm::function_ref fn) { intptr_t n = mlirOperationGetNumAttributes(operation->get()); for (intptr_t i = 0; i < n; ++i) { MlirNamedAttribute na = mlirOperationGetAttribute(operation->get(), i); From 446380a8562bbe1dc8c37151a6dcb7127685d68f Mon Sep 17 00:00:00 2001 From: Perry Gibson Date: Tue, 14 Oct 2025 09:05:11 +0100 Subject: [PATCH 5/6] chore: make forEachAttr static after review feedback --- mlir/lib/Bindings/Python/IRCore.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index 7fc55cb5b870f..0e7b1dd3a88d9 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -2730,7 +2730,8 @@ class PyOpAttributeMap { operation->get(), toMlirStringRef(name))); } - void forEachAttr(llvm::function_ref fn) { + static void + forEachAttr(llvm::function_ref fn) { intptr_t n = mlirOperationGetNumAttributes(operation->get()); for (intptr_t i = 0; i < n; ++i) { MlirNamedAttribute na = mlirOperationGetAttribute(operation->get(), i); From 73f12231e2a6a1c1883bcf0c3646baaba7cd5dda Mon Sep 17 00:00:00 2001 From: Perry Gibson Date: Tue, 14 Oct 2025 09:41:32 +0100 Subject: [PATCH 6/6] fix: apply consequences of static method --- mlir/lib/Bindings/Python/IRCore.cpp | 47 +++++++++++++++++------------ 1 file changed, 28 insertions(+), 19 deletions(-) diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index 0e7b1dd3a88d9..06d0256e6287b 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -2731,10 +2731,11 @@ class PyOpAttributeMap { } static void - forEachAttr(llvm::function_ref fn) { - intptr_t n = mlirOperationGetNumAttributes(operation->get()); + forEachAttr(MlirOperation op, + llvm::function_ref fn) { + intptr_t n = mlirOperationGetNumAttributes(op); for (intptr_t i = 0; i < n; ++i) { - MlirNamedAttribute na = mlirOperationGetAttribute(operation->get(), i); + MlirNamedAttribute na = mlirOperationGetAttribute(op, i); MlirStringRef name = mlirIdentifierStr(na.name); fn(name, na.attribute); } @@ -2751,36 +2752,44 @@ class PyOpAttributeMap { .def("__iter__", [](PyOpAttributeMap &self) { nb::list keys; - self.forEachAttr([&](MlirStringRef name, MlirAttribute) { - keys.append(nb::str(name.data, name.length)); - }); + PyOpAttributeMap::forEachAttr( + self.operation->get(), + [&](MlirStringRef name, MlirAttribute) { + keys.append(nb::str(name.data, name.length)); + }); return nb::iter(keys); }) .def("keys", [](PyOpAttributeMap &self) { nb::list out; - self.forEachAttr([&](MlirStringRef name, MlirAttribute) { - out.append(nb::str(name.data, name.length)); - }); + PyOpAttributeMap::forEachAttr( + self.operation->get(), + [&](MlirStringRef name, MlirAttribute) { + out.append(nb::str(name.data, name.length)); + }); return out; }) .def("values", [](PyOpAttributeMap &self) { nb::list out; - self.forEachAttr([&](MlirStringRef, MlirAttribute attr) { - out.append(PyAttribute(self.operation->getContext(), attr) - .maybeDownCast()); - }); + PyOpAttributeMap::forEachAttr( + self.operation->get(), + [&](MlirStringRef, MlirAttribute attr) { + out.append(PyAttribute(self.operation->getContext(), attr) + .maybeDownCast()); + }); return out; }) .def("items", [](PyOpAttributeMap &self) { nb::list out; - self.forEachAttr([&](MlirStringRef name, MlirAttribute attr) { - out.append( - nb::make_tuple(nb::str(name.data, name.length), - PyAttribute(self.operation->getContext(), attr) - .maybeDownCast())); - }); + PyOpAttributeMap::forEachAttr( + self.operation->get(), + [&](MlirStringRef name, MlirAttribute attr) { + out.append(nb::make_tuple( + nb::str(name.data, name.length), + PyAttribute(self.operation->getContext(), attr) + .maybeDownCast())); + }); return out; }); }