Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
makslevental merged 4 commits into from
Dec 2, 2024
Merged

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

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bad0d96
Add check to see if `FuncOp.arg_attrs` is set
Wheest Nov 21, 2024
049bace
Linting fix
Wheest Nov 21, 2024
2ae4815
Add additional test case
Wheest Nov 22, 2024
c6afea5
Return empty dict list, do not mutate IR
Wheest Nov 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions mlir/python/mlir/dialects/func.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Copy link
Contributor

@makslevental makslevental Nov 22, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor Author

@Wheest Wheest Nov 22, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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

@arg_attrs.setter
Expand Down
29 changes: 29 additions & 0 deletions mlir/test/python/dialects/func.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"})
Loading