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]: extend embedded implementation of premap() #1501

Merged
merged 51 commits into from
May 16, 2024

Conversation

egparedes
Copy link
Contributor

@egparedes egparedes commented Mar 19, 2024

Extend the implementation of the premap field operation (previously named remap, conceptually equivalent to a Contravariant Functor's contramap) to support more efficient implementations of different use cases depending on the contents of the connectivity field.

Added

  • gt4py.eve: new typing aliases and minor utilities

Changed

  • gt4py.next.common:

    • new typing aliases.
    • small refactoring of Domain to support creation of subdomains via slicing using the .slice_at attribute. The actual implementation comes from the now deleted gt4py.next.embedded.nd_array_field._relative_ranges_to_domain() function.
    • refactor ConnectivityKind to represent all known use cases
    • extend CartesianConnectivity to support translation and relocations
    • rename remap to premap
  • gt4py.next.embedded.nd_array_field:

    • full refactoring of premap() (old remap) and add usage documentation
    • some renamings (_hypercube() -> _hyperslice(), _compute_mask_ranges() -> _compute_mask_slices()

Removed

  • gt4py.next.embedded.nd_array_field: _relative_ranges_to_domain() function moved to an Domain attribute in gt4py.next.common

@egparedes egparedes changed the title Extend remap embedded feat(next): extend remap embedded Mar 22, 2024
@egparedes egparedes marked this pull request as ready for review April 18, 2024 10:10
@egparedes egparedes changed the title feat(next): extend remap embedded feat(next): extend embedded implementation of premap() Apr 18, 2024
@egparedes egparedes requested a review from havogt April 18, 2024 12:06
@egparedes
Copy link
Contributor Author

egparedes commented May 8, 2024

I tried to fix all mentioned issued but there are still a couple of open points to discuss. Could you do another review pass?

@egparedes egparedes requested a review from havogt May 8, 2024 09:47
Examples:
>>> I, J = Dimension("I"), Dimension("J")
>>> domain = Domain(NamedRange(I, UnitRange(0, 10)), NamedRange(J, UnitRange(5, 15)))
>>> domain.slice_at[2, 2:5]
Copy link
Contributor

Choose a reason for hiding this comment

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

It might be confusing that slicing with int behaves different than for fields/arrays as it doesn't reduce rank. Any arguments pro this behavior?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

To avoid inconsistencies, slice_at now only accepts slices, not ints.

src/gt4py/next/common.py Show resolved Hide resolved
src/gt4py/next/embedded/nd_array_field.py Outdated Show resolved Hide resolved
new_domain = self.domain.replace(dim_idx, *new_ranges)
# Select actual implementation of the transformation
if not (conn_fields[0].kind & common.ConnectivityKind.TRANSFORM_DATA):
return _domain_premap(self, *conn_fields)
Copy link
Contributor

Choose a reason for hiding this comment

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

I think we had a test that didn't use the short-circuit path for Cartesian connectivity, but went through the expensive index->index mapping.

src/gt4py/next/embedded/nd_array_field.py Show resolved Hide resolved
@egparedes egparedes changed the title feat(next): extend embedded implementation of premap() feat[next]: extend embedded implementation of premap() May 16, 2024
@egparedes egparedes merged commit 723b81f into GridTools:main May 16, 2024
43 checks passed
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