[mlir,python] Fix case when FuncOp.arg_attrs is not set

[Wheest]

FuncOps can have arg_attrs, an array of dictionary attributes associated with their arguments.

E.g.,

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.,

func.func @main(%arg0: tensor<8xf32>, %arg1: tensor<8x16xf32>)

This raises the error:

    return ArrayAttr(self.attributes[ARGUMENT_ATTRIBUTE_NAME])
                     ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'attempt to access a non-existent attribute'

This PR fixes this by returning the expected [{}, {}].

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Perry Gibson (Wheest)

Changes

FuncOps can have arg_attrs, an array of dictionary attributes associated with their arguments.

E.g.,

func.func @<!-- -->main(%arg0: tensor&lt;8xf32&gt; {test.attr_name = "value"}, %arg1: tensor&lt;8x16xf32&gt;)

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.,

func.func @<!-- -->main(%arg0: tensor&lt;8xf32&gt;, %arg1: tensor&lt;8x16xf32&gt;)

This raises the error:

    return ArrayAttr(self.attributes[ARGUMENT_ATTRIBUTE_NAME])
                     ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'attempt to access a non-existent attribute'

This PR fixes this by returning the expected [{}, {}].

---

Full diff: https://github.com/llvm/llvm-project/pull/117188.diff

2 Files Affected:

• (modified) mlir/python/mlir/dialects/func.py (+4)
• (modified) mlir/test/python/dialects/func.py (+29)

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..6b3932ce64f137 100644
--- a/mlir/test/python/dialects/func.py
+++ b/mlir/test/python/dialects/func.py
@@ -104,3 +104,32 @@ def testFunctionCalls():
 # CHECK:   %1 = call @qux() : () -> f32
 # CHECK:   return
 # CHECK: }
+
+
+# CHECK-LABEL: TEST: testFunctionArgAttrs
+@constructAndPrintInModule
+def testFunctionArgAttrs():
+    foo = func.FuncOp("foo", ([F32Type.get()], []))
+    foo.sym_visibility = StringAttr.get("private")
+    foo2 = func.FuncOp("foo2", ([F32Type.get(), F32Type.get()], []))
+    foo2.sym_visibility = StringAttr.get("private")
+
+    empty_attr = DictAttr.get({})
+    test_attr = DictAttr.get({"test.foo": StringAttr.get("bar")})
+    test_attr2 = DictAttr.get({"test.baz": StringAttr.get("qux")})
+
+    assert len(foo.arg_attrs) == 1
+    assert foo.arg_attrs[0] == empty_attr
+
+    foo.arg_attrs = [test_attr]
+    assert foo.arg_attrs[0]["test.foo"] == StringAttr.get("bar")
+
+    assert len(foo2.arg_attrs) == 2
+    assert foo2.arg_attrs == ArrayAttr.get([empty_attr, empty_attr])
+
+    foo2.arg_attrs = [empty_attr, test_attr2]
+    assert foo2.arg_attrs == ArrayAttr.get([empty_attr, test_attr2])
+
+
+# CHECK: func private @foo(f32 {test.foo = "bar"})
+# CHECK: func private @foo2(f32, f32  {test.baz = "qux"})

[makslevental]

I think it makes more sense to just return None here? I know it doesn't match the correct semantic but it seems odd to make empty dictionaries just to return them as empty?

[Wheest]

I can see your point, however in the code I was working on, I have something like:

for arg, arg_attr in zip(block.arguments, func_op.arg_attrs):
    if "my.value" in arg_attr:
        lst1.append(arg)
    else:
        lst2.append(arg)

In your proposal, I would need to add a None check such as:

if func_op.arg_attrs is None:
    lst2.extend(block.arguments)
else:
    for arg, arg_attr in zip(block.arguments, func_op.arg_attrs):
        if "my.value" in arg_attr:
            lst1.append(arg)
        else:
            lst2.append(arg)

This would improve performance, but this is Python and perhaps not the chief concern, rather simplicity.

Happy to make the change to return None if you prefer, but from my perspective the overhead of generating empty dicts and then walking over an empty list makes more sense.

[ftynse]

I don't necessarily mind returning a list of empty dictionary attributes when the attribute is missing (though it doesn't correspond to the underlying IR structure by mixing up two different cases in the IR), but having a property accessor mutate the operation by forcibly setting an attribute on it is wrong.

[Wheest]

I don't necessarily mind returning a list of empty dictionary attributes when the attribute is missing (though it doesn't correspond to the underlying IR structure by mixing up two different cases in the IR), but having a property accessor mutate the operation by forcibly setting an attribute on it is wrong.

I agree, the underlying IR shouldn't be changed by this. I've updated the return so that it gives the empty dict list, but not stored in the IR.

- self.attributes[ARGUMENT_ATTRIBUTE_NAME] = ArrayAttr.get(
-    [DictAttr.get({}) for _ in self.type.inputs]
- )
+ return ArrayAttr.get([DictAttr.get({}) for _ in self.type.inputs])

No further objections, but I'll let @maksleventhal comment

[Wheest]

I don't yet have committer permissions, would one of you be willing to do this please @makslevental @ftynse

[makslevental]

@Wheest if you want to keep contributing (and we want you to!), I highly encourage you to get commit access.