-
-
Notifications
You must be signed in to change notification settings - Fork 5.5k
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
Customizable lazy fused broadcasting in pure Julia #26891
Merged
Merged
Changes from 3 commits
Commits
Show all changes
7 commits
Select commit
Hold shift + click to select a range
dbe1ae0
Customizable lazy fused broadcasting in pure Julia
mbauman c9a0a41
Rename Broadcast.make to Broadcast.broadcasted
mbauman 4cd10d3
Test and fix #26127
mbauman 8a2d88a
Remove hack that avoided materialize in fallback broadcast
mbauman 75798f9
Remove the last of the base broadcast methods...
mbauman edd4dcf
Correct NEWS.md
mbauman 18ad6a8
Change lowering of `A .= x` to match fused RHS cases
mbauman File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains 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
This file contains 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
Oops, something went wrong.
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.
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.
It looks like the lowering actually calls
Broadcast.materialize
orBroadcast.materialize!
?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.
Ah, yes, that's true. I still want to highlight
copy
andcopyto!
as those are the methods to specialize. Let's just remove that "parser will call" to make it technically correct.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.
Is there a reason why the parser doesn't lower to these
copy
calls directly?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.
Good question. There are several reasons here, and they're different for in-place (
broadcast!
/.=
) vs. out-of-place broadcasting.For in-place broadcasts, we use
materialize!
for two purposes: (1) the outermostbroadcasted
might not return aBroadcasted
object, but we need to support broadcasting into the destination (e.g.,copyto!(rand(3,3), 1:3)
doesn't broadcast) and we want to use the same code-paths here. It solves both problems at once by wrapping the non-Broadcasted
objectx
in aBroadcasted(identity, (x,))
. (2) TheBroadcasted
argument wants to know about its axes, but those are dependent upon the axes of the destination. It explicitly sets the axes of theBroadcasted
object to those of the destination, and then those are checked to ensure they are a compatible shape withinstantiate
.For out-of-place broadcasts,
materialize
also serves two purposes: (1) forBroadcasted
arguments, it computes and caches the axes of the outermostBroadcasted
, recursively checking the axes of its arguments byinstantiate
-ing them. Now, we may be able to get away with not storing theaxes
sinceaxes
should theoretically only be called once and it can be computed on demand, but I haven't tested it and it seems strange to throw aDimensionError
from a call toaxes
. (2) for all other arguments it defaults toidentity
but can serve as a hook for custom types to do similar sorts of specializations decoupled fromcopy
.I'm also aiming for some semblance of symmetry between the two cases.
Now there
are twois one caveat here:(1) A wart that I really dislike is(Fixed in 8a2d88a — this hack is no longer necessary). (2) The expressionbroadcast
itself is hard-coded to callcopy(instantiate(broadcasted(…)))
instead ofmaterialize
. For reasons beyond my comprehension,materialize
sometimes failed to inline even though it is explicitly marked as@inline
.A .= x
still lowers tobroadcast!(identity, A, x)
. If otherwise unspecialized, it will fall back to the default implementation, which constructs the appropriateBroadcasted
object and then goes through the above machinery as expected. We may want to change this purely to be able to simplify our description of this parser transformation.