Skip to content
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

Merged
merged 15 commits into from
Oct 11, 2024

Conversation

tehrengruber
Copy link
Contributor

@tehrengruber tehrengruber commented Sep 28, 2024

Adds a pass that transforms expressions like

as_fieldop(λ(__arg0, __arg1) → ·__arg0 + ·__arg1, c⟨ IDimₕ: [0, 1) ⟩)(
  as_fieldop(λ(__arg0, __arg1) → ·__arg0 × ·__arg1, c⟨ IDimₕ: [0, 1) ⟩)(inp1, inp2), inp3
)

into

as_fieldop(λ(inp1, inp2, inp3) → ·inp1 × ·inp2 + ·inp3, c⟨ IDimₕ: [0, 1) ⟩)(inp1, inp2, inp3)

@tehrengruber tehrengruber marked this pull request as draft September 28, 2024 09:56
@tehrengruber tehrengruber marked this pull request as ready for review October 9, 2024 22:57
@tehrengruber tehrengruber changed the title feat[next]: as_fieldop fusion pass feat[next]: GTIR as_fieldop fusion pass Oct 10, 2024
Copy link
Contributor

@egparedes egparedes left a 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):
Copy link
Contributor

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
Copy link
Contributor

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:

Suggested change
extracted_args[inner_arg.id] = inner_arg
extracted_args[inner_param.id] = inner_arg

Copy link
Contributor Author

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.

src/gt4py/next/iterator/transforms/fuse_as_fieldop.py Outdated Show resolved Hide resolved
src/gt4py/next/iterator/transforms/fuse_as_fieldop.py Outdated Show resolved Hide resolved
src/gt4py/next/iterator/transforms/fuse_as_fieldop.py Outdated Show resolved Hide resolved
src/gt4py/next/iterator/type_system/inference.py Outdated Show resolved Hide resolved
Copy link
Contributor

@egparedes egparedes left a comment

Choose a reason for hiding this comment

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

LGTM

@tehrengruber tehrengruber merged commit ed9d82d into GridTools:main Oct 11, 2024
31 checks passed
tehrengruber added a commit that referenced this pull request Oct 18, 2024
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants