Add ability to broadcast views to higher-dimensional arrays #383
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
tcclevenger
reviewed
Sep 8, 2025
src/kokkos/ekat_view_broadcast.hpp
Outdated
| KOKKOS_FORCEINLINE_FUNCTION | ||
| reference_type operator()(Is... indices) const { | ||
| int i=0; | ||
| ((indices *= coeff[i++]),...); |
Contributor
There was a problem hiding this comment.
Could we do some sort of assertion on i that it didn't exceed the rank of m_view_impl? Or check the size of ...indices at the beginning.
Contributor
Author
There was a problem hiding this comment.
Yeah, good idea. Maybe I should just add 8 overloads for ranks 1-8...
Contributor
Author
There was a problem hiding this comment.
I switched to 8 separate methods
- Use int[8] rather than another view to store outfacing layout - Implement 8 separate method for operator(), like Kokkos::View does - Improve docs for the constructor
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Motivation
This is motivated by some work in EAMxx. Namely, I have an array (nx, ny, nz), which I need to mask with another array, which could have shape (nx) or (nx,nz). To avoid lots of conditionals, I would like to do something similar to numpy's array broadcasting.
This PR allows to wrap a view into a higher dim array, in such a way that it is "replicated" along the new dimensions. E.g., one can take a 1d view
xrepresenting a column, and replicated it in a 2d matrix arrayA, in such a way thatA(i,j) = x(j)for alli. This, of course, does not increase storage requirement, and does not copy any data. It simply "hides" the slicing of "extruded" dimensions.E3SM Stakeholder Feedback
May be needed in EAMxx.
Testing
Added a unit test.