Skip to content
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

autodiff_deferred with derivatives of order greater than 2 #1173

Open
cgeoga opened this issue Nov 30, 2023 · 3 comments · May be fixed by JuliaGPU/GPUCompiler.jl#556
Open

autodiff_deferred with derivatives of order greater than 2 #1173

cgeoga opened this issue Nov 30, 2023 · 3 comments · May be fixed by JuliaGPU/GPUCompiler.jl#556

Comments

@cgeoga
Copy link

cgeoga commented Nov 30, 2023

I'm working on an application where I really want non-allocating higher-order derivatives. Here is my little code setup:

module EnzymeHigherOrder

  using Enzyme
  
  struct DefDerivative{F}
    fn::F
  end
  function (dd::DefDerivative{F})(x) where{F} 
    autodiff_deferred(Forward, dd.fn, DuplicatedNoNeed, 
                      Duplicated(x, one(x)))[1]
  end

  struct Derivative{F}
    fn::F
  end
  function (dv::Derivative{F})(x) where{F} 
    autodiff(Forward, dv.fn, DuplicatedNoNeed, Duplicated(x, one(x)))[1]
  end

  derivative(fn::F, ::Val{0}) where{F}   = fn
  derivative(fn::F, ::Val{1}) where{F}   = Derivative(fn)
  derivative(fn::F, ::Val{N}) where{F,N} = derivative(DefDerivative(fn), Val(N-1))

end

And I can use this to get second derivatives as

const ddcos = EnzymeHigherOrder.derivative(cos, Val(2))
@btime ddcos($(1.1)); # no allocations, great perf, woohoo!

But when I ask for a third derivative I get an error:

const dddcos = EnzymeHigherOrder.derivative(cos, Val(3))
dddcos(1.1)

gives

ERROR: Enzyme execution failed.
Enzyme: Not yet implemented, mixed activity for jl_new_struct constants=Bool[1, 0]   %5 = call noalias nonnull {} addrspace(10)* ({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)*, {} addrspace(10)*, ...) @julia.call({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)* noundef nonnull @jl_f_tuple, {} addrspace(10)* noundef null, {} addrspace(10)* nonnull %0, {} addrspace(10)* nonnull %3) #11

Stacktrace:
 [1] throwerr(cstr::Cstring)
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/rbuCz/src/compiler.jl:3044
 [2] macro expansion
   @ ~/.julia/packages/Enzyme/rbuCz/src/compiler.jl:9778 [inlined]
 [3] enzyme_call
   @ ~/.julia/packages/Enzyme/rbuCz/src/compiler.jl:9456 [inlined]
 [4] ForwardModeThunk
   @ ~/.julia/packages/Enzyme/rbuCz/src/compiler.jl:9422 [inlined]
 [5] autodiff
   @ ~/.julia/packages/Enzyme/rbuCz/src/Enzyme.jl:330 [inlined]
 [6] autodiff
   @ ~/.julia/packages/Enzyme/rbuCz/src/Enzyme.jl:222 [inlined]
 [7] (::Main.EnzymeHigherOrder.Derivative{Main.EnzymeHigherOrder.DefDerivative{Main.EnzymeHigherOrder.DefDerivative{typeof(cos)}}})(x::Float64)
   @ Main.EnzymeHigherOrder ~/Scratch/higherorderad/enzyme_ho.jl:18
 [8] top-level scope
   @ REPL[3]:1

Does this mean that autodiff_deferred can't differentiate other autodiff_deferred? So that if I want a third derivative in forward-mode I need to do something that looks like autodiff(autodiff(autodiff_deferred))? I would appreciate any thoughts on executing this as efficiently as possible! Thanks in advance.

@vchuravy
Copy link
Member

Does this mean that autodiff_deferred can't differentiate other autodiff_deferred?

No it definitely needs to be autodiff(autodiff_deferred(autodiff_deferred())) for third order.

@cgeoga
Copy link
Author

cgeoga commented Nov 30, 2023

Makes sense. Could you comment on the error then? Am I making some elementary mistake with the interface?

@jgreener64
Copy link
Contributor

I also run into a similar error. On Enzyme main (5e4e2ef) and Julia 1.10.2:

using Enzyme

f(x) = 4 * x^3
x = 3.0

function df_dx(x)
    autodiff_deferred(
        Forward,
        f,
        DuplicatedNoNeed,
        Duplicated(x, 1.0),
    )[1]
end

function d2f_dx2(x)
    autodiff_deferred(
        Forward,
        df_dx,
        DuplicatedNoNeed,
        Duplicated(x, 1.0),
    )[1]
end

# Forward over forward over forward fails
autodiff_deferred(
    Forward,
    d2f_dx2,
    DuplicatedNoNeed,
    Duplicated(x, 1.0),
)

# Reverse over forward over forward fails
autodiff_deferred(
    Reverse,
    d2f_dx2,
    Active,
    Active(x),
)

The errors for the two cases are similar, for example for the first:

ERROR: Enzyme execution failed.
Enzyme compilation failed.
Current scope: 
; Function Attrs: mustprogress willreturn
define double @preprocess_julia_df_dx_2320(double "enzymejl_parmtype"="140488835241888" "enzymejl_parmtype_ref"="0" %0) local_unnamed_addr #4 !dbg !36 {
top:
  %1 = call {}*** @julia.get_pgcstack() #5
  %ptls_field3 = getelementptr inbounds {}**, {}*** %1, i64 2
  %2 = bitcast {}*** %ptls_field3 to i64***
  %ptls_load45 = load i64**, i64*** %2, align 8, !tbaa !8
  %3 = getelementptr inbounds i64*, i64** %ptls_load45, i64 2
  %safepoint = load i64*, i64** %3, align 8, !tbaa !12, !invariant.load !7
  fence syncscope("singlethread") seq_cst
  call void @julia.safepoint(i64* %safepoint) #5, !dbg !37
  fence syncscope("singlethread") seq_cst
  %4 = call i64 @deferred_codegen(i64 noundef 140488246046696) #5, !dbg !38
  %5 = inttoptr i64 %4 to [1 x double] (double, double)*, !dbg !45
  %6 = call [1 x double] %5(double %0, double 1.000000e+00) #5, !dbg !45
  %.fca.0.extract = extractvalue [1 x double] %6, 0, !dbg !45
  ret double %.fca.0.extract, !dbg !49
}

No forward mode derivative found for deferred_codegen
 at context:   %4 = call i64 @deferred_codegen(i64 noundef 140488246046696) #5, !dbg !15

Stacktrace:
 [1] macro expansion
   @ ~/.julia/dev/Enzyme/src/compiler.jl:5665
 [2] deferred_codegen (repeats 3 times)
   @ ~/.julia/dev/Enzyme/src/compiler.jl:5623
 [3] autodiff_deferred
   @ ~/.julia/dev/Enzyme/src/Enzyme.jl:481
 [4] autodiff_deferred
   @ ~/.julia/dev/Enzyme/src/Enzyme.jl:495
 [5] df_dx
   @ ~/dms/molly_dev/enzyme_err38.jl:8


Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:1289
  [2] macro expansion
    @ ~/.julia/dev/Enzyme/src/compiler.jl:5665 [inlined]
  [3] deferred_codegen (repeats 3 times)
    @ ~/.julia/dev/Enzyme/src/compiler.jl:5623 [inlined]
  [4] autodiff_deferred
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:481 [inlined]
  [5] autodiff_deferred
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:495 [inlined]
  [6] df_dx
    @ ~/dms/molly_dev/enzyme_err38.jl:8 [inlined]
  [7] fwddiffejulia_df_dx_2320wrap
    @ ~/dms/molly_dev/enzyme_err38.jl:0 [inlined]
  [8] macro expansion
    @ ~/.julia/dev/Enzyme/src/compiler.jl:5440 [inlined]
  [9] enzyme_call
    @ ~/.julia/dev/Enzyme/src/compiler.jl:5118 [inlined]
 [10] ForwardModeThunk
    @ ~/.julia/dev/Enzyme/src/compiler.jl:5003 [inlined]
 [11] autodiff_deferred
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:483 [inlined]
 [12] autodiff_deferred
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:495 [inlined]
 [13] d2f_dx2
    @ ~/dms/molly_dev/enzyme_err38.jl:17 [inlined]
 [14] fwddiffejulia_d2f_dx2_2317wrap
    @ ~/dms/molly_dev/enzyme_err38.jl:0
 [15] macro expansion
    @ ~/.julia/dev/Enzyme/src/compiler.jl:5440 [inlined]
 [16] enzyme_call
    @ ~/.julia/dev/Enzyme/src/compiler.jl:5118 [inlined]
 [17] ForwardModeThunk
    @ ~/.julia/dev/Enzyme/src/compiler.jl:5003 [inlined]
 [18] autodiff_deferred
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:483 [inlined]
 [19] autodiff_deferred(mode::ForwardMode{FFIABI}, f::typeof(d2f_dx2), ::Type{DuplicatedNoNeed}, args::Duplicated{Float64})
    @ Enzyme ~/.julia/dev/Enzyme/src/Enzyme.jl:495
 [20] top-level scope
    @ REPL[11]:1

Does anyone have an example where third order AD works?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging a pull request may close this issue.

3 participants