Add recursive map generalizing the make_zero mechanism

[danielwe]

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 existing recursive_* functions like this is appropriate or whether it would be better to implement a separate deep_recursive_accumulate.

This probably breaks some existing uses of recursive_accumulate, like the Holomorphic derivative code, because recursive_accumulate now traverses most/all of the structure on its own and will double-accumulate when combined with the iteration over the seen 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 from make_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 of a .= f.(b, c, ...), free of intermediate allocations whenever possible (and similarly an out-of-place recursive_broadcast(f, a, b, c...) generalizing f.(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 version quadgk!(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.

[danielwe]

Alright, I could take some feedback/discussion on this now.

• This implements a generic recursive_map for mapping a function over the differentiable values in arbitrary tuples of identical data structures. There's an in-place equivalent recursive_map! for mutable values, built on top of a bangbang-style recursive_map!! that works on arbitrary types and reuses all the mutable memory (similar to the old make_zero_immutable! but without code duplication).
• The code is diffed with the old make_zero(!) code on github, but this is a complete rewrite and the diff will probably not be helpful for reviewing. Let me know if you want me to rename files or something to get rid github's diff view.
  • The implementation is leaner and simpler than the old one, and even though recursive_map{!!} aren't public I wrote extensive docstrings to clarify the spec for myself and others, so I don't think the code should be too difficult to review from scratch.
• I added fast paths such that new structs are allocated using splatnew with a tuple instead of ccall with a vector in the common case where there are no undefined fields. This gives a substantial speedup in many cases, which is good since recursive_map will be called in hot loops in custom quadrature rules and the like.
• I have refactored make_zero and make_zero! to be minimal wrappers around recursive_map{!}, without changing their public API.
  • To stay safe while doing such a big refactoring, I wrote extensive tests with ~full coverage of both the old and new implementations of make_zero(!) (a small number of edge case branches aren't covered only because I can't find a way to reach them from any public entry point).
  • These tests uncovered quite a few bugs in the existing make_zero(!) implementations. See the following commit on a separate branch in my fork for the necessary fixes to get the old code to pass the new tests: danielwe@7a6ca9f. (Note, one of the tests still errors due to active_reg_inner with justActive = true incorrect with immutable types that can be incompletely initialized #1935.)
• I have not refactored recursive_add and recursive_accumulate, since these currently have different semantics where they don't recurse into mutable values. I'm happy to go ahead and do the refactoring if you're OK with changing their semantics.

TLDR: Should I rewrite recursive_add and recursive_accumulate to be based on recursive_map{!} and have full recursion semantics? Anything else?

@gdalle Promised to tag you when this was ready for review, but note that this PR only deals with the low-level, non-public guts of the implementation. I'll do the vector space wrapper in a separate PR as soon as this is merged (hopefully that won't be long, I really need that QuadGK rule for my research 📐)

[danielwe]

Wondering whether this is necessary or if the leaf types could just be hardcoded to Union{_RealOrComplexFloat,Array{<:_RealOrComplexFloat}}. I'll make a prototype of the vector space wrapper and the updated QuadGK rules to see if customizable leaf types comes in handy.

[danielwe]

Update for anyone who's following: I've implemented the VectorSpace wrapper, which prompted me to adjust the recursive_map implementation a bit, all for the better. It's looking good and will make writing custom higher-order rules as well as the DI wrappers a lot nicer for arbitrary types. However, it dawned on me that you probably want make_zero to be easily extensible by just adding methods, like what's already done in the StaticArrays extension. That will require a bit of redesign, nothing too hard, but I've got weekend plans so might not get to it until next week.

[wsmoses]

awesome, sorry I haven't had a chance to review let [just a bunch of schenanigans atm], I'll try to take a closer look next week and ping me if not

[danielwe]

No worries! I restored the draft label when I realized there was a bit more to do and will remove it again once I think this is ready for review. No need to look at it until then, the current state here on github doesn't reflect what I'm working with locally anyway.

[wsmoses]

Just to confirm, this is only for make_zero, and not for add/etc?

Because this case here already feels specific to the context

[danielwe]

It's going to look a bit different once I push the next update (hopefully tomorrow), but no, after some experimenting it seemed best to me to always skip guaranteed inactive subtrees and restrict recursive_map to applying f to the differentiable values only. I tried doing the opposite initially, leaving it as part of the isleaftype filter and handling the possible deepcopy within the mapped function f, but it made things a lot more complicated. I think the main issue was that the whole mechanism with seen and keeping track of object identity then becomes the purview of the mapped function f instead of recursive_map itself, increasing boilerplate and complicating the contract between recursive_map and its callers. I couldn't think of a use case within Enzyme where you're interested in mapping over the guaranteed inactive parts anyway, and not recursing through inactive subtrees saves you from having to deal with deconstruction/reconstruction of a few specialized types (deepcopy has a lot more methods than recursive_map). So I went with this solution instead.

Of course, adding a skip_guaranteed_const flag would be straightforward (or combining it with copy_if_inactive into a single inactive_mode parameter). Do you think this is warranted?

[danielwe]

At long last, I think this one's ready for you to take a look. Hit me with any questions and concerns, from major design issues to bikeshedding over names.

I put both the implementation and tests in their own modules because they define a lot of helpers and I didn't want to pollute other modules' namespaces.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

Attention: Patch coverage is 86.49518% with 42 lines in your changes missing coverage. Please review.

Project coverage is 75.21%. Comparing base (037dfed) to head (4e1e58f).
Report is 334 commits behind head on main.

Files with missing lines	Patch %	Lines
src/typeutils/recursive_maps.jl	88.21%	31 Missing ⚠️
src/typeutils/recursive_add.jl	74.35%	10 Missing ⚠️
src/internal_rules.jl	0.00%	1 Missing ⚠️

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1852      +/-   ##
==========================================
+ Coverage   67.50%   75.21%   +7.70%     
==========================================
  Files          31       56      +25     
  Lines       12668    16618    +3950     
==========================================
+ Hits         8552    12499    +3947     
- Misses       4116     4119       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[wsmoses]

I'm still knee deep in 1.11 land and don't have cycles to review this immediately. @vchuravy can you take a look?

[danielwe]

1.11 efforts deeply appreciated! Don't rush this. I'll keep using 1.10 and a local fork for my own needs, and occiasionally push small changes here as my tinkering surfaces new concerns/opportunities.

[wsmoses]

@danielwe is this the one now to review or is there a different related PR I should review first? (and also would you mind rebasing)

[danielwe]

This is the one wrt. recursive maps and all that, but I've been continually refining stuff locally, so I need to both rebase and push the latest changes, hang on! Will remove draft status when ready for review.

[danielwe]

Finally wrapped up and rebased this! I'll come back later and write a little blurb, but the code should be ready for review as-is.

[vchuravy]

Nice!

WARNING: Method definition inactive_type(Type{var"#s2630"} where var"#s2630"<:(RecursiveMapTests.CustomVector{T} where T)) in module RecursiveMapTests at /home/runner/work/Enzyme.jl/Enzyme.jl/test/recursive_maps.jl:867 overwritten at /home/runner/work/Enzyme.jl/Enzyme.jl/test/recursive_maps.jl:882.

[vchuravy]

This is a new API? It will need a version bump for EnzymeCore.

[danielwe]

We should discuss. The reason I put these helpers in EnzymeCore rather than keeping them internal was that the StaticArrays extension needed to add a method, so I figured there's a chance others might have to do the same for their custom types. However, subtyping DenseArray (and AbstractFloat if that ever becomes relevant) should almost always be sufficient. Either way, the point is only to make these extensible in package extensions. I don't think anyone should ever have to call them.

[danielwe]

However, the vector space wrapper functionality I've built on top of this (which will be a separate PR) will probably involve a new type in EnzymeCore, so if that gets accepted there will have to be a new EnzymeCore release anyway

[vchuravy]

Yeah, absolutely. I think this is the right place to add them, and EnzymeCore is basically meant for people to be able to extend things without having to bite the load time bullet that is Enzyme.

[danielwe]

Btw. names are infinitely bikesheddable, both in this case and elsewhere in the PR. My mindset working on this PR is to enable consistent treatment of arbitrary objects as vectors in a space spanned by the scalar (float) values reachable from the object, hence all the vector/scalar terminology, but I don't know if this works well or if it's confusing, especially as part of the public API.

[danielwe]

Blurb time!

Let's start with a quick overview of the API. I also wrote exhaustive docstrings in the code (just my way of clarifying my thinking, and I hope it's useful for code review too), but this should be a more conversational introduction highlighting the main points.

Out-of-place usage

(out1::T, out2::T, ...) = recursive_map([seen::IdDict,] f, Val(Nout), (in1::T, in2::T, ...), [Val(copy_if_inactive), [isinactivetype]])

This generalizes map to recurse through inputs of arbitrary type and map f over all differentiable leaves. The function f should have methods (leaf_out1::U, leaf_out2::U, ...) = f(leaf_in1::U, leaf_in2::U, ...) for every type U whose instances are such leaves. In various docstrings and helper functions I refer to these leaf types as "vector types", since they are supposed to be the lowest-level types that natively implement vector space operations like +, zero, et cetera.

Note how this supports mapped functions f that have an arbitrary number of outputs, assembling them into Nout instances of the input type T rather than a single output containing leaves of types NTuple{Nout,U}. This was necessary to replace accumulate_into from internal_rules.jl. The number of outputs Nout should be passed to recursive_map as Val(Nout). To avoid complexity in the implementation, f has to return a Tuple{U} rather than just U even when there is only a single output. Since this is an internal API I figured this was an acceptable tradeoff.

Partially-in-place usage

(new_out1::T, new_out2::T, ...) = recursive_map([seen::IdDict,] f!!, (out1::T, out2::T, ...), (in1::T, in2::T, ...), [Val(copy_if_inactive), [isinactivetype]])

This form has bangbang-like semantics where mutable storage within the (out1, out2, ...) is mutated and reused, but there is no requirement that every differentiable value is in mutable storage. New immutable containers will be instantiated as needed and you need to use the returned values (new_out1, new_out2, ...) downstream.

To use this form, the mapped function f!! should have the same kind of out-of-place method as f above for all vector types, as well as mutating in-place methods for mutable, non-scalar vector types: (leaf_out1::U, leaf_out2::U, ...) = f!!(leaf_out1::U, leaf_out2::U, ..., leaf_in1::U, leaf_in2::U, ...). (The function should return the mutated outputs because technically you're also allowed to register non-mutable but non-scalar/non-isbits vector types that contain some mutable storage internally and are only partially mutated in place by this method. In such cases, you need to return the new outputs. This is hardly relevant in practice, and there's an argument for simplifying this and requiring that a vector type is either mutable or scalar (or both).)

The design of the whole API clicked for me the moment I realized that if you get this version right, both the in-place and out-of-place versions are just special cases, so you only need a single core of recursive functions with a f ewdifferent entry points for the various use cases.

In-place usage

If you pass types where all differentiable values live in mutable storage, partially-in-place already implements in-place behavior for you. The in-place function recursive_map! is just a wrapper that checks that the inputs satisfy this condition and throws an error otherwise.

recursive_map!([seen::IdDict,] f!!, (out1::T, out2::T, ...), (in1::T, in2::T, ...), [Val(copy_if_inactive), [isinactivetype]])::Nothing

f!! should have the same methods as described above for partially-in-place usage. (The parenthetically mentioned partially-in-place method variant for f!! is not relevant here, as the corresponding vector types would fail the in-place validity check.)

Optional arguments

seen::IdDict and Val{copy_if_inactive} should be familiar from the existing make_zero implementation.

• seen tracks object identity to reproduce the graph topology of the input objects, i.e., multiple references to the same object and/or cyclical references. One difference is that the in-place form also uses an IdDict rather than an IdSet, since the inputs and outputs are not generally the same objects even in in-place usage---outputs may just be preallocated and uninitialized storage. (Inputs and outputs being identical is of course still allowed, otherwise make_zero! wouldn't work.)
  • I spent quite a bit of energy thinking about how to ensure/enforce consistency in graph topology with multiple inputs and outputs. I can elaborate more if needed, but the punchline is that the first input in1 is the reference version: other inputs must at least mirror all the aliasing seen in the first input; new outputs will reproduce the topology of the first input; and existing, reused outputs must have no more aliasing than the first input; all this so that every output leaf value is well-defined.
• Val{copy_if_inactive} decides whether non-differentiable subgraphs should be deepcopied or shared between inputs and outputs.
• isinactivetype is a callable that decides exactly what these "non-differentiable subgraphs" are. This was a tricky one for multiple reasons, so I'll give it a dedicated paragraph below.

isinactivetype and IsInactive

There were three concerns:

• Performance: This was never that important for make_zero, but a key goal of recursive_map is to be the basis for writing rules for higher-order functions like quadrature routines. When you need to map vector space operations over hundreds or even thousands of closure instances, performance is important.

  guaranteed_const_nongen ended up being a performance bottleneck due to frequent type instabilities. I therefore decided to make the compile-time version guaranteed_const the default. However, this means new methods added to EnzymeRules.inactive_type won't generally be picked up by subsequent calls to recursive_map. Hence, there should be a simple way to choose the _nongen version when needed.

• Consistency: When the in-place recursive_map! validates input/output, it should ignore non-differentiable subgraphs; what's important is that the differentiable values live in mutable storage. Hence this needs to take isinactivetype into account. When using the compile-time guaranteed_const, there exists a corresponding guaranteed_nonactive that ensures this consistency. I added guaranteed_nonactive_nongen to complete the 2x2 matrix. Now all that's needed is an API for choosing between guaranteed_const and guaranteed_const_nongen that also ensures that the corresponding version of guaranteed_nonactive* is used for validation.

• Extras: Some uses of recursive_map(!) need to terminate at (i.e., treat as inactive) additional parts of the object graph. Case in point: recursive_accumulate!, now renamed to accumulate_seen!, used for holomorphic derivatives---this should terminate whenever it encounters a type that would be cached in seen. Hence it must be possible to pass a custom treat_this_as_inactive_too callable to be used alongside guaranteed_const*, and for consistency, recursive_map! argument validation should also hook into this.

Rather than a proliferation of optional arguments and internal voodoo, this called for a dedicated abstraction. Hence the IsInactive struct, which works as follows:

• Instantiation: isinactivetype = IsInactive{runtime}([extra]) where runtime is a Bool and the optional extra is your treat_this_as_inactive_too function described above.
• The signature of isinactivetype as a callable is isinactivetype(T, [Val(nonactive)])::Bool
• If runtime == false you get compile-time evaluation (no *_nongen):
  • isinactivetype(T) and isinactivetype(T, Val(false)) call guaranteed_const(T) || extra(T)
  • isinactivetype(T, Val(true)) calls guaranteed_nonactive(T) || extra(T)
• If runtime == true you get runtime evaluation (*_nongen):
  • isinactivetype(T) and isinactivetype(T, Val(false)) call guaranteed_const_nongen(T) || extra(T)
  • isinactivetype(T, Val(true)) callsguaranteed_nonactive_nongen(T) || extra(T)

This functionality is the reason behind the method overwriting warning pointed out by @vchuravy above:

WARNING: Method definition inactive_type(Type{var"#s2630"} where var"#s2630"<:(RecursiveMapTests.CustomVector{T} where T)) in module RecursiveMapTests at /home/runner/work/Enzyme.jl/Enzyme.jl/test/recursive_maps.jl:867 overwritten at /home/runner/work/Enzyme.jl/Enzyme.jl/test/recursive_maps.jl:882.

This warning comes from the test suite, not the package code. The test suite overwrites inactive_type a couple of times to verify that the various modes have the expected behavior (runtime = true picks up the new methods, runtime = false does not).

isinactivetype in practice

The optional argument ::Val{runtime}=Val(false) was added to make_zero(!) to let the user choose between compile-time and runtime inactivity checking. Under the hood, this passes IsInactive{runtime}() to recursive_map(!).

A word on generated functions

Unrolling loops over tuples and struct fields is crucial for type stability and performance. I tried to only use ntuple and recursive tuple peeling idioms combined with aggressive @inlineing, but couldn't consistently eliminate as many allocations as I could with simple and arguably less convoluted generated functions. Moreover, I saw some segfaults when using ntuple to assemble arguments for splatnew. Hence, I decided to use @generated for a few core functions. However, I made them as benign as I possibly could: there is no manual Expr handling, only quoted code where select compile-time constants are interpolated into the Base.@ntuple and Base.Cartesian.@nexprs macros. Perhaps the most questionable choice is the single use of @goto, but how else do you emulate break in an unrolled loop? (However, that loop is the least important one to unroll, so we can remove the @goto if you want.) I avoided gratuitous use of @generated when recursive tuple peeling worked just as well, such as for the getitems and setitems! helpers.

Generality/GPU compatibility

When the recursion hits a DenseArray{T} where T is a bitstype, it uses broadcasting rather than scalar iteration to recurse into the array. This means that everything should hopefully work out of the box for GPU arrays even when the array eltype is non-scalar (e.g., a struct or tuple). This is tested for JLArray, but I haven't actually checked if it works for CuArray or MtlArray. (The question is just whether the broadcasted closure compiles on GPU; I don't see why it shouldn't as long as the mapped function f works on GPU and you use IsInactive{false} to avoid type instabilities, but this remains to be seen. I don't know how well GPUCompiler does with recursion.)

In the common case where the GPU array eltype is scalar (i.e., float), the array is considered a leaf and dispatched directly to the mapped function f without further broadcast/recursion, so this case should definitely work on GPU, provided f itself does.

[danielwe]

• guaranteed_const_nongen ended up being a performance bottleneck due to frequent type instabilities.

I did some more profiling, and this is not true anymore, possibly due to the @assume_effects annotations added in the meantime. Using *_nongen versions of the functions is now at least as fast as the alternative. I'll go ahead and simplify the implementation accordingly.

[danielwe]

Would a default arg value world=nothing be acceptable here and in guaranteed_const_nongen?

[danielwe]

@wsmoses @vchuravy This line in accumulate_into! is the only reason multiple outputs are supported in recursive_map(!). The reverse rule for deepcopy is the only place this is used.

We could change this to

    recursive_map!(accumulate_into_alt!!, into, (into, from))
    make_zero!(from)

to completely remove the need for multiple outputs. That would simplify the implementation of recursive_map! and make its signature/usage more intuitive.

However, this implementation of accumulate_into! recurses twice through from instead of once, allocating two IdDicts instead of one, so the deepcopy rule will perform somewhat worse.

Other uses of recursive_map may see slightly improved performance because the indirection of the ubiquitous tuples in the current implementation seems to add some extra runtime dispatch when objects have abstractly typed fields/elements. In type-stable cases, I don't think there will be any difference.

What's your opinion on this tradeoff?