Enzyme: Autodiff Functor which contains cache

Specific Domains
,
1

Hello,
I have a RHS of an ODE in the form of a Functor which includes some caches, which i am trying to make Enzym compatible. Heres an MWE of the problem

using Pkg
pkg"activate --temp"
pkg"add Enzyme, ForwardDiff, PreallocationTools, DifferentiationInterface"
using Enzyme: Enzyme
using ForwardDiff: ForwardDiff
using PreallocationTools
using DifferentiationInterface: DifferentiationInterface as DI

struct Functor{C}
    cache::C
end
function (f::Functor)(du, u, p, t)
    tmp = get_tmp(f.cache, u)
    tmp .= u .* p
    du .= tmp
    nothing
end
f = Functor(DiffCache(zeros(2)))

single_arg_wrapper = u -> begin
    du = zeros(eltype(u), 2)
    f(du, u, [1,2], nothing)
    return du
end

x0 = zeros(2)
ForwardDiff.jacobian(single_arg_wrapper, x0) # correct result
Enzyme.jacobian(Enzyme.Forward, single_arg_wrapper, x0)[1] # returns zeros(2,2)
Enzyme.jacobian(Enzyme.Reverse, single_arg_wrapper, x0)[1] # returns zeros(2,2)

Enzyme gives zero jacobians. However when I used DifferentiationInterface instead

DI.jacobian(f, zeros(2), DI.AutoEnzyme(), x0, DI.Constant([1,2]), DI.Constant(nothing))

I don’t need the single_arg_wrapper and get a more meaningful error:

ERROR: Function argument passed to autodiff cannot be proven readonly.
If the the function argument cannot contain derivative data, instead call autodiff(Mode, Const(f), ...)
See https://enzyme.mit.edu/index.fcgi/julia/stable/faq/#Activity-of-temporary-storage for more information.

Whats the “correct” way of dealing with this? I guess Enzyme being tricked into believing f is constant by the wrapper might be a bug. Is there a simple way to provide Enzyme with a Duplicated version of the functor?

2

Hi @hexaeder!
The current answer for DI: we introduced an additional setting to AutoEnzyme in order to make functors work properly. Can you try it with AutoEnzyme(function_annotation=Enzyme.Duplicated)?
Note that this functionality has not been tested extensively, but if you encounter bugs we’ll fix them together.
The future answer for DI: Add `Cache` as another type of context · Issue #551 · gdalle/DifferentiationInterface.jl · GitHub

1 Like
3

Thanks for the reply, function_annotation together with mode=Enzyme.set_runtime_activity(Enzyme.Forward) made forward work mode works now on my actual code.

Reverse mode errors now with the below error, I’ll try to recreate a MWE for that…

ERROR: MethodError: no method matching addrspace(::LLVM.ArrayType)

Closest candidates are:
  addrspace(::LLVM.PointerType)
   @ LLVM ~/.julia/packages/LLVM/joxPv/src/core/type.jl:312

Stacktrace:
  [1] get_array_len(B::LLVM.IRBuilder, array::LLVM.ExtractValueInst)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:1698
  [2] jl_array_del_end_rev(B::LLVM.IRBuilder, orig::LLVM.CallInst, gutils::Enzyme.Compiler.GradientUtils, tape::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Vjlrr/src/rules/llvmrules.jl:1530
  [3] jl_array_del_end_rev_cfunc(B::Ptr{…}, OrigCI::Ptr{…}, gutils::Ptr{…}, tape::Ptr{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Vjlrr/src/rules/llvmrules.jl:48
  [4] EnzymeCreatePrimalAndGradient(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnValue::Bool, dretUsed::Bool, mode::Enzyme.API.CDerivativeMode, runtimeActivity::Bool, width::Int64, additionalArg::Ptr{…}, forceAnonymousTape::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, augmented::Ptr{…}, atomicAdd::Bool)
    @ Enzyme.API ~/.julia/packages/Enzyme/Vjlrr/src/api.jl:253
  [5] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{…}, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:4706
  [6] codegen(output::Symbol, job::GPUCompiler.CompilerJob{…}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:7801
  [7] codegen
    @ ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:6638 [inlined]
  [8] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:8909
  [9] _thunk
    @ ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:8909 [inlined]
 [10] cached_compilation
    @ ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:8950 [inlined]
 [11] thunkbase(ctx::LLVM.Context, mi::Core.MethodInstance, ::Val{…}, ::Type{…}, ::Type{…}, tt::Type{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Type{…}, ::Val{…}, ::Val{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:9082
 [12] #s2067#19055
    @ ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:9219 [inlined]
 [13] 
    @ Enzyme.Compiler ./none:0
 [14] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any})
    @ Core ./boot.jl:602
 [15] autodiff
    @ ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:473 [inlined]
 [16] value_and_pullback(::Network{…}, ::Vector{…}, ::DifferentiationInterface.NoPullbackPrep, ::AutoEnzyme{…}, ::Vector{…}, ::NTuple{…}, ::Constant{…}, ::Constant{…})
    @ DifferentiationInterfaceEnzymeExt ~/.julia/packages/DifferentiationInterface/v1iM4/ext/DifferentiationInterfaceEnzymeExt/reverse_twoarg.jl:105
 [17] pullback
    @ ~/.julia/packages/DifferentiationInterface/v1iM4/src/first_order/pullback.jl:322 [inlined]
 [18] (::DifferentiationInterface.var"#42#43"{…})(a::Int64)
    @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/v1iM4/src/first_order/jacobian.jl:269
 [19] iterate
    @ ./generator.jl:47 [inlined]
 [20] _collect(c::Base.OneTo{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
    @ Base ./array.jl:854
 [21] collect_similar
    @ ./array.jl:763 [inlined]
 [22] map
    @ ./abstractarray.jl:3285 [inlined]
 [23] _jacobian_aux
    @ ~/.julia/packages/DifferentiationInterface/v1iM4/src/first_order/jacobian.jl:268 [inlined]
 [24] jacobian
    @ ~/.julia/packages/DifferentiationInterface/v1iM4/src/first_order/jacobian.jl:183 [inlined]
 [25] jacobian(::Network{…}, ::Vector{…}, ::AutoEnzyme{…}, ::Vector{…}, ::Constant{…}, ::Constant{…})
    @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/v1iM4/src/fallbacks/no_prep.jl:75
 [26] top-level scope
    @ REPL[55]:1
 [27] top-level scope
    @ none:1
Some type information was truncated. Use `show(err)` to see complete types.
1 Like