Make recursive_add/accumulate more recursive #1852
Draft
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.
This is to explore functionality for realizing JuliaMath/QuadGK.jl#120. The current draft cuts time and allocations in half for the MWE in that PR compared to the
make_zero
hack from the comments. Not sure if modifying the existingrecursive_*
functions like this is appropriate or whether it would be better to implement a separatedeep_recursive_accumulate
.This probably breaks some existing uses of
recursive_accumulate
, like the Holomorphic derivative code, becauserecursive_accumulate
now traverses most/all of the structure on its own and will double-accumulate when combined with the iteration over theseen
IdDicts. Curious to see the total impact on the test suite.This doesn't yet have any concept of
seen
and will thus double-accumulate if the structure has internal aliasing. That obviously needs to be fixed. Perhaps we can factor out and share the recursion code frommake_zero
.A bit of a tangent, but perhaps a final version of this PR should include migrating
ClosureVector
to Enzyme from the QuadGK ext as suggested in JuliaMath/QuadGK.jl#110 (comment). Looks like that's the most relevant application of fully recursive accumulation at the moment.Let me also throw out another suggestion: what if we implement a recursive generalization of broadcasting with an arbitrary number of arguments, i.e.,
recursive_broadcast!(f, a, b, c, ...)
as a recursive generalization ofa .= f.(b, c, ...)
, free of intermediate allocations whenever possible (and similarly an out-of-placerecursive_broadcast(f, a, b, c...)
generalizingf.(a, b, c...)
that only materializes/allocates once if possible). That would enable more optimized custom rules with Duplicated args, such as having the QuadGK rule call the in-place versionquadgk!(f!, result, segs...)
. Not sure if it would be hard to correctly handle aliasing without being overly defensive, or if that could mostly be taken care of by proper reuse of the existing broadcasting functionality.