-
Notifications
You must be signed in to change notification settings - Fork 49
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.
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
feat[next]: GTIR as_fieldop
fusion pass
#1670
feat[next]: GTIR as_fieldop
fusion pass
#1670
Conversation
…p' into gtir_fuse_as_fieldop
as_fieldop
fusion passas_fieldop
fusion pass
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM, only minor comments.
from gt4py.next.type_system import type_info, type_specifications as ts | ||
|
||
|
||
def _inline_as_fieldop_arg(arg: itir.Expr, uids: eve_utils.UIDGenerator): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Missing return type annotation in some functions of this module (why is mypy not complaining here?)
for inner_param, inner_arg in zip(stencil.params, inner_args, strict=True): | ||
if isinstance(inner_arg, itir.SymRef): | ||
stencil_params.append(inner_param) | ||
extracted_args[inner_arg.id] = inner_arg |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If the comment above is not wrong, I think this should be:
extracted_args[inner_arg.id] = inner_arg | |
extracted_args[inner_param.id] = inner_arg |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I changed the comment above. It is the stencil param of the stencil of the outer as_fieldop
that the current stencil is being inlined into.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
New temporary extraction pass. Transforms an `itir.Program` like ``` testee(inp, out) { out @ c⟨ IDimₕ: [0, 1) ⟩ ← as_fieldop(deref, c⟨ IDimₕ: [0, 1) ⟩)(as_fieldop(deref, c⟨ IDimₕ: [0, 1) ⟩)(inp)); } ``` into ``` testee(inp, out) { __tmp_1 = temporary(domain=c⟨ IDimₕ: [0, 1) ⟩, dtype=float64); __tmp_1 @ c⟨ IDimₕ: [0, 1) ⟩ ← as_fieldop(deref, c⟨ IDimₕ: [0, 1) ⟩)(inp); out @ c⟨ IDimₕ: [0, 1) ⟩ ← as_fieldop(deref, c⟨ IDimₕ: [0, 1) ⟩)(__tmp_1); } ``` Note that this pass intentionally unconditionally extracts. In case you don't want a temporary you should fuse the `as_fieldop` before. As such the fusion pass (see #1670) contains the heuristics on what to fuse.
Adds a pass that transforms expressions like
into