[Mlir-commits] [mlir] [mlir, python] Fix case when `FuncOp.arg_attrs` is not set (PR #117188)
Perry Gibson
llvmlistbot at llvm.org
Thu Nov 21 08:47:26 PST 2024
https://github.com/Wheest created https://github.com/llvm/llvm-project/pull/117188
FuncOps can have `arg_attrs`, an array of dictionary attributes associated with their arguments.
E.g.,
```mlir
func.func @main(%arg0: tensor<8xf32> {test.attr_name = "value"}, %arg1: tensor<8x16xf32>)
```
These are exposed via the MLIR Python bindings with `my_funcop.arg_attrs`.
In this case, it would return `[{test.attr_name = "value"}, {}]`, i.e., `%arg1` has an empty `DictAttr`.
However, if I try and access this property from a FuncOp with an empty `arg_attrs`, e.g.,
```mlir
func.func @main(%arg0: tensor<8xf32>, %arg1: tensor<8x16xf32>)
```
This raises the error:
```python
return ArrayAttr(self.attributes[ARGUMENT_ATTRIBUTE_NAME])
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'attempt to access a non-existent attribute'
```
This PR fixes this by returning the expected `[{}, {}]`.
>From bad0d96c1f3be74b9fb7f9da9c7cd05c58b62375 Mon Sep 17 00:00:00 2001
From: pez <perry at fractile.ai>
Date: Thu, 21 Nov 2024 16:40:04 +0000
Subject: [PATCH] Add check to see if `FuncOp.arg_attrs` is set
---
mlir/python/mlir/dialects/func.py | 4 ++++
mlir/test/python/dialects/func.py | 13 +++++++++++++
2 files changed, 17 insertions(+)
diff --git a/mlir/python/mlir/dialects/func.py b/mlir/python/mlir/dialects/func.py
index 24fdcbcd85b29f..211027d88051a7 100644
--- a/mlir/python/mlir/dialects/func.py
+++ b/mlir/python/mlir/dialects/func.py
@@ -105,6 +105,10 @@ def add_entry_block(self, arg_locs: Optional[Sequence[Location]] = None):
@property
def arg_attrs(self):
+ if ARGUMENT_ATTRIBUTE_NAME not in self.attributes:
+ self.attributes[ARGUMENT_ATTRIBUTE_NAME] = ArrayAttr.get(
+ [DictAttr.get({}) for _ in self.type.inputs]
+ )
return ArrayAttr(self.attributes[ARGUMENT_ATTRIBUTE_NAME])
@arg_attrs.setter
diff --git a/mlir/test/python/dialects/func.py b/mlir/test/python/dialects/func.py
index a2014c64d2fa53..bcfaace853bc64 100644
--- a/mlir/test/python/dialects/func.py
+++ b/mlir/test/python/dialects/func.py
@@ -104,3 +104,16 @@ def testFunctionCalls():
# CHECK: %1 = call @qux() : () -> f32
# CHECK: return
# CHECK: }
+
+
+# CHECK-LABEL: TEST: testFunctionArgAttrs
+ at constructAndPrintInModule
+def testFunctionArgAttrs():
+ foo = func.FuncOp("foo", ([("arg0", F32Type.get())], []))
+
+ assert len(foo.arg_attrs) == 1
+ assert foo.arg_attrs[0] = ir.DictAttr.get({})
+
+ foo.arg_attrs = [DictAttr.get({"test.foo": StringAttr.get("bar")})]
+
+ assert foo.arg_attrs[0]["test.foo"] == StringAttr.get("bar")
More information about the Mlir-commits
mailing list