transformations: Persist func arg names as arg_attr #3395
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #3395 +/- ##
==========================================
+ Coverage 90.13% 90.16% +0.02%
==========================================
Files 452 453 +1
Lines 57157 57220 +63
Branches 5498 5502 +4
==========================================
+ Hits 51519 51590 +71
+ Misses 4180 4175 -5
+ Partials 1458 1455 -3 ☔ View full report in Codecov by Sentry. |
| def match_and_rewrite(self, op: func.FuncOp, rewriter: PatternRewriter, /): | ||
| arg_attrs = list(op.arg_attrs or (len(op.args) * [DictionaryAttr({})])) | ||
| new_arg_attrs: list[DictionaryAttr] = [] | ||
| for arg, arg_attr in zip(op.args, arg_attrs): | ||
| if arg.name_hint: | ||
| d_attr = DictionaryAttr( | ||
| {"llvm.name": StringAttr(arg.name_hint), **arg_attr.data} | ||
| ) | ||
| else: | ||
| d_attr = arg_attr | ||
| new_arg_attrs.append(d_attr) | ||
|
|
||
| op.arg_attrs = ArrayAttr(new_arg_attrs) | ||
| rewriter.has_done_action = True |
There was a problem hiding this comment.
- It feels like the first thing to do would be to check whether any of the args aren't anonymous, and to bail out early if that's the case. Could you add a test case for this, also?
- could you also check whether the old attrs are the same as the new attrs, and avoid has_done_action if they are the same?
- nit: it would be nice to refactor this to a tuple comprehension instead of appending to a list, just for style reasons
There was a problem hiding this comment.
I had it as a tuple-comp list-comp, and after refactoring it am glad that is encouraged 😃
| class ArgNamesToArgAttrsPass(RewritePattern): | ||
| @op_type_rewrite_pattern | ||
| def match_and_rewrite(self, op: func.FuncOp, rewriter: PatternRewriter, /): | ||
| arg_attrs = list(op.arg_attrs or (len(op.args) * [DictionaryAttr({})])) |
There was a problem hiding this comment.
| arg_attrs = list(op.arg_attrs or (len(op.args) * [DictionaryAttr({})])) | |
| arg_attrs = op.arg_attrs.data if op.arg_attrs is not None else or (DictionaryAttr({}),) * len(op.args) |
There was a problem hiding this comment.
done as suggested (minus the 'else or')
Co-authored-by: Sasha Lopoukhine <superlopuh@gmail.com>
| def match_and_rewrite(self, op: func.FuncOp, rewriter: PatternRewriter, /): | ||
| arg_attrs = list(op.arg_attrs or (len(op.args) * [DictionaryAttr({})])) |
There was a problem hiding this comment.
Would it not be worth doing something like this?
| def match_and_rewrite(self, op: func.FuncOp, rewriter: PatternRewriter, /): | |
| arg_attrs = list(op.arg_attrs or (len(op.args) * [DictionaryAttr({})])) | |
| def match_and_rewrite(self, op: func.FuncOp, rewriter: PatternRewriter, /): | |
| if not any(arg.name_hint for arg in op.args): | |
| return | |
| arg_attrs = list(op.arg_attrs or (len(op.args) * [DictionaryAttr({})])) |
This way there's no intermediate thing created if there's nothing to do
There was a problem hiding this comment.
At the cost of an extra pass over the args? I actually didn't think that's what we wanted.
| @@ -257,6 +257,7 @@ def callback(x: str): | |||
| with ImplicitBuilder(expected_module.body): | |||
| function = func.FuncOp("hello", ((index,), (index,))) | |||
| with ImplicitBuilder(function.body) as (n,): | |||
| n.name_hint = "n" | |||
There was a problem hiding this comment.
This will otherwise show up as a difference in get_available_passes below, as the module specified as string has the name hint.
There was a problem hiding this comment.
Looks like this fixes a kind of bug in the test, which is that the name hints are not the same in the text and this IR
Persisting function arg names to arg attributes as suggested [here](https://discourse.llvm.org/t/mlir-capi-how-to-set-function-argument-names-arg0/3892/2). Note, upstream mlir will say `op arguments may only have dialect attributes`, requiring arg attribute names to be prefixed with a dialect. --------- Co-authored-by: n-io <n-io@users.noreply.github.com> Co-authored-by: Sasha Lopoukhine <superlopuh@gmail.com>
Persisting function arg names to arg attributes as suggested here.
Note, upstream mlir will say
op arguments may only have dialect attributes, requiring arg attribute names to be prefixed with a dialect.