More parameters for `AutoEnzyme`

[gdalle]

I think it would make sense to add more parameters to AutoEnzyme, like:

• deferred = true / false to use autodiff_deferred instead of autodiff
• chunksize to imitate the AutoForwardDiff API and allow selecting a chunk size for Jacobians and Hessians

@wsmoses @vchuravy what do you think?

Related:

• Exception access violation with Enzyme ("Incorrect number of arguments passed to called function!") JuliaDiff/DifferentiationInterface.jl#819 (comment)

[vchuravy]

deferred should not be needed. Enzyme does this automatically now and it is just needed for directly using it in GPU kernels.

[wsmoses]

gentle bump on this @gdalle since it came up in the linked thread. would be great to have

[gdalle]

I would love some input on the options we'd like there. ForwardDiff has two:

• chunk size=nothing for automatic selection, upper-bounded by 12 for bandwidth reasons
• chunk size::Integer for fixed manual selection

Does Enzyme need a third option which always takes the batch size to be the length of the vector (aka computing the whole Jacobian in a single batch for example)?

[gdalle]

Related:

• Heuristic for picking the chunk/batch size? EnzymeAD/Enzyme.jl#1542
• Implement pick_batchsize EnzymeAD/Enzyme.jl#1545

[wsmoses]

yeah the "same size as vector" does have performance benefits, so it would be nice to represent

[gdalle]

I was thinking we could use chunksize=Inf for that one. Sure, Inf is technically a Float but typemax(Int) is much less readable and Inf gets the point across rather well. How does that sound?

[wsmoses]

I guess the question becomes what is the meaning if someone gives a float != Inf. Ideally we would restrict the type to only contain valid inputs.

Maybe a new type class?

[gdalle]

We can check type parameters at construction with an inner constructor:

struct AutoEnzyme{M,A,chunksize}
    mode::M
    function AutoEnzyme{M,A,chunksize}(mode::M) where {M,A,chunksize}
        if chunksize isa Integer
            @assert chunksize > 0
        elseif chunksize isa Float
            @assert chunksize == Inf
        else
            @assert isnothing(chunksize)
        end
    end
    return new{M,A,chunksize}(mode)
end

[wsmoses]

Is that cleaner than say AutoEnzyme(ForwardMode, FullBatch), where FullBatch is a new struct? (I'm not sure, genuinely asking this as a design question)

[gdalle]

I'd say that with ForwardDiff, ADTypes has already started a tradition of using built-in language values to represent specific batching behaviors. For instance, chunksize=nothing really means chunksize=AutoChunk(), in the same way that chunksize=Inf would mean chunksize=FullChunk(). Unfortunately, we can't change the ForwardDiff API without a breaking release, so I would suggest we continue down that path for the time being, but keep this change in mind for a future ADTypes v2.0?

[wsmoses]

I mean we could equally put FullChunk/etc in EnzymeCore proper [like mode]

[gdalle]

Down the road we'll need the same chunking options for Mooncake too (https://github.com/chalk-lab/Mooncake.jl/discussions/533), and we already have them for ForwardDiff. I think it would be nice for users if these three were handled in the same way?

[wsmoses]

perhaps, though per #123 users already need to import EnzymeCore.Forward anyways atm

[gdalle]

I'm not sure how to get around that one, because Enzyme's mode objects are rather sophisticated. For Mooncake I'll probably do AutoMooncake(; mode=ADTypes.ForwardMode()) and call it a day. Maybe we could add such a shortcut for Enzyme as well, but it wouldn't incorporate things like runtime activity, holomorphy, etc.