feat[eve]: Preserve annex in custom visitors #1874
Conversation
tehrengruber
changed the title
src/gt4py/eve/visitors.py
Outdated
| if (value := getattr(old_annex, key, NOTHING)) is not NOTHING: | ||
| assert key not in new_annex_dict | ||
| new_annex_dict[key] = value | ||
| _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS) |
There was a problem hiding this comment.
What do you think about leaving the condition outside of the function? This is probably the hottest loop in the toolchain so I'd rather try to avoid a unnecessary function call if possible. Additionally, it would also make the function name fully correct 😄
| _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS) | |
| if preserved_annex_attrs and hasattr(node, "__node_annex__"): | |
| _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS) |
There was a problem hiding this comment.
I thought about this, but went for not duplicating the condition. Makes sense to me though, so I reverted that.
| @@ -436,8 +436,12 @@ def _infer_stmt( | |||
| **kwargs: Unpack[InferenceOptions], | |||
| ): | |||
| if isinstance(stmt, itir.SetAt): | |||
| # constant fold once otherwise constant folding after domain inference might create (syntactic) differences | |||
There was a problem hiding this comment.
Does this change really belong to this PR?
There was a problem hiding this comment.
I initially checked that the new annex attribute is equal to the old one. That would be a desirable property of the IR and this change here is a step in that direction. This also motivated the comment above
# Note: The annex value of the new node might not be equal
# (in the sense that an equality comparison returns false),
# but in the context of the pass, they are equivalent.
# Therefore, we don't assert equality here.
However there are annex values that are equivalent, but not equal e.g. the domain annex with a value of
(domain, domain) (tuple of domains) is equivalent to domain. I therefore gave up on the assert, but kept this change here. Does that make sense?
Preserve annex attributes in the
NodeTranslatorfor nodes returned from customvisit_methods, not only from nodes which are recreated ingeneric_visit. This is often desired and with this PR doesn't need to be done manually in child classes ofNodeTranslator. Thevisitmethod can still decide if the node it returns should have a different annex attribute as_preserve_annexdoes not overwrite in case the attribute already exists.