[Mlir-commits] [mlir] [MLIR][Python] Register `OpAttributeMap` as `Mapping` for `match` compatibility (PR #174292)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jan 9 14:55:28 PST 2026
https://github.com/MaPePeR updated https://github.com/llvm/llvm-project/pull/174292
>From 86ed86690a0a32fd06373859882a1ad9a6beec2c Mon Sep 17 00:00:00 2001
From: MaPePeR <MaPePeR at users.noreply.github.com>
Date: Wed, 31 Dec 2025 19:03:34 +0000
Subject: [PATCH 1/2] [MLIR][Python] Register OpAttributeMap as Mapping for
match compatibility
---
mlir/include/mlir/Bindings/Python/IRCore.h | 3 ++
mlir/lib/Bindings/Python/IRCore.cpp | 14 +++++
mlir/python/mlir/_mlir_libs/__init__.py | 3 +-
mlir/test/python/ir/operation.py | 63 ++++++++++++++++++++++
4 files changed, 82 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 8a8d2a1b2270f..b2eb6b9edb677 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1782,6 +1782,9 @@ class MLIR_PYTHON_API_EXPORTED PyOpAttributeMap {
PyNamedAttribute dunderGetItemIndexed(intptr_t index);
+ nanobind::object getWithDefaultNamed(const std::string &key,
+ nanobind::object defaultValue);
+
void dunderSetItem(const std::string &name, const PyAttribute &attr);
void dunderDelItem(const std::string &name);
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 8396569757183..eb5ebf15ea728 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2388,6 +2388,16 @@ PyOpAttributeMap::dunderGetItemNamed(const std::string &name) {
return PyAttribute(operation->getContext(), attr).maybeDownCast();
}
+nb::object PyOpAttributeMap::getWithDefaultNamed(const std::string &key,
+ nb::object defaultValue) {
+ MlirAttribute attr =
+ mlirOperationGetAttributeByName(operation->get(), toMlirStringRef(key));
+ if (mlirAttributeIsNull(attr)) {
+ return defaultValue;
+ }
+ return PyAttribute(operation->getContext(), attr).maybeDownCast();
+}
+
PyNamedAttribute PyOpAttributeMap::dunderGetItemIndexed(intptr_t index) {
if (index < 0) {
index += dunderLen();
@@ -2450,6 +2460,10 @@ void PyOpAttributeMap::bind(nb::module_ &m) {
"Sets an attribute with the given name.")
.def("__delitem__", &PyOpAttributeMap::dunderDelItem, "name"_a,
"Deletes an attribute with the given name.")
+ .def("get", &PyOpAttributeMap::getWithDefaultNamed, nb::arg("key"),
+ nb::arg("default") = nb::none(),
+ "Gets an attribute by name or the default value, if it does not "
+ "exist.")
.def(
"__iter__",
[](PyOpAttributeMap &self) {
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index ce7e6bf93012a..c0e8775149d41 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -2,7 +2,7 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-from typing import Any, Sequence
+from typing import Any, Mapping, Sequence
import os
@@ -245,6 +245,7 @@ def __str__(self):
Sequence.register(ir.OpResultList)
Sequence.register(ir.OpSuccessors)
Sequence.register(ir.RegionSequence)
+ Mapping.register(ir.OpAttributeMap)
_site_initialize()
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index b9242a7cc2bd9..89f78ab1932a0 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -652,6 +652,34 @@ def testOperationAttributes():
# CHECK: Dict mapping {'dependent': 'text', 'other.attribute': 3.0, 'some.attribute': 1}
print("Dict mapping", d)
+ # Structural pattern matching test using Mapping
+
+ # CHECK: Matched Mapping Attribute 'some.attribute': 1
+ # CHECK: Matched Mapping Attribute 'other.attribute': 3.0
+ # CHECK: Matched Mapping Attribute 'dependent': text
+ match op:
+ case OpView(attributes={"does_not_exist": a0}):
+ assert False
+ case OpView(
+ attributes={
+ "some.attribute": IntegerAttr(value=some_attr_val) as some_attr,
+ "other.attribute": FloatAttr() as other_attr,
+ "dependent": StringAttr() as dep_attr,
+ **other_attributes,
+ }
+ ):
+ print(f"Matched Mapping Attribute 'some.attribute': {some_attr_val}")
+ print(f"Matched Mapping Attribute 'other.attribute': {other_attr.value}")
+ print(f"Matched Mapping Attribute 'dependent': {dep_attr.value}")
+ assert type(other_attributes) == dict
+ assert len(other_attributes) == 0
+ assert some_attr == op.attributes.get("some.attribute")
+ assert other_attr == op.attributes.get("other.attribute", None)
+ assert dep_attr == op.attributes.get("dependent", "Default value")
+ case _:
+ print("Did not match!")
+ assert False
+
# Check that exceptions are raised as expected.
try:
op.attributes["does_not_exist"]
@@ -667,6 +695,41 @@ def testOperationAttributes():
else:
assert False, "expected IndexError on accessing an out-of-bounds attribute"
+ # Check that exceptions are raised when `get` is used with non-str arg.
+ try:
+ op.attributes.get(0)
+ except TypeError:
+ pass
+ else:
+ assert False, "expected TypeError using int as key for get()"
+
+ try:
+ op.attributes.get(0, None)
+ except TypeError:
+ pass
+ else:
+ assert False, "expected TypeError using int as key for get()"
+
+ try:
+ op.attributes.get([], None)
+ except TypeError:
+ pass
+ else:
+ assert False, "expected TypeError using list as key for get()"
+
+ try:
+ match op:
+ case OpView(attributes={0: a}):
+ assert False
+ except TypeError:
+ pass
+ else:
+ assert False, "expected TypeError matching OpAttributeMap with int-key "
+
+ # get() does not throw for non existent attributes.
+ assert op.attributes.get("does_not_exist") is None
+ assert op.attributes.get("does_not_exist", "default_value") == "default_value"
+
# CHECK-LABEL: TEST: testOperationPrint
@run
>From f05200899b88303a46241b1e822abd7e9d860482 Mon Sep 17 00:00:00 2001
From: MaPePeR <MaPePeR at users.noreply.github.com>
Date: Fri, 9 Jan 2026 23:55:20 +0100
Subject: [PATCH 2/2] Update mlir/lib/Bindings/Python/IRCore.cpp
Co-authored-by: Maksim Levental <maksim.levental at gmail.com>
---
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 eb5ebf15ea728..151183855a7b8 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2392,9 +2392,8 @@ nb::object PyOpAttributeMap::getWithDefaultNamed(const std::string &key,
nb::object defaultValue) {
MlirAttribute attr =
mlirOperationGetAttributeByName(operation->get(), toMlirStringRef(key));
- if (mlirAttributeIsNull(attr)) {
+ if (mlirAttributeIsNull(attr))
return defaultValue;
- }
return PyAttribute(operation->getContext(), attr).maybeDownCast();
}
More information about the Mlir-commits
mailing list