EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
455 stars 64 forks source link

Computing third order derivatives #1707

Closed Arpit-Babbar closed 2 months ago

Arpit-Babbar commented 3 months ago

I am able to go up to second order derivatives with Enzyme. However, when I use the same approach to go for the third derivative, I get an error. Here is a minimal working example to show that. The last function call (supposed to compute the third order derivative) gives the error.

using Enzyme
using StaticArrays
function g(x) # Some function to be differentiated
    x[1]^4 + x[1]^7 + x[2]^4 + x[1]^2 * x[2]^2
end

# Enzyme.jl derivatives
dg(x, dx) = autodiff_deferred(Forward, g, Duplicated, Duplicated(x, dx))
ddg(x, dx, ddx) = autodiff_deferred(Forward, dg, Duplicated, Duplicated(x, dx), Duplicated(dx, ddx))
dddg(x, dx, ddx, dddx) = autodiff_deferred(Forward, ddg, Duplicated, Duplicated(x, dx),
                                           Duplicated(dx, ddx), Duplicated(ddx, dddx))

# Get random inputs for x and some directions for the directional derivatives
x = SVector{2}(rand(2)); dx = SVector{2}(rand(2)); ddx = SVector{2}(rand(2)); dddx = SVector{2}(rand(2))

# Call the functions with inputs
dg(x, dx) # Works
ddg(x, dx, ddx) # Works
dddg(x, dx, ddx, dddx) # Error

The error message is below.

Error message ```julia ERROR: Enzyme execution failed. Enzyme compilation failed. No forward mode derivative found for deferred_codegen at context: %5 = call i64 @deferred_codegen(i64 noundef 139926815969256) #5, !dbg !19, !noalias !8 Stacktrace: [1] macro expansion @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:7113 [2] deferred_codegen (repeats 3 times) @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:7072 [3] autodiff_deferred @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:524 [4] autodiff_deferred @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:538 [5] dg @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:78 [6] dg @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0 Stacktrace: [1] throwerr(cstr::Cstring) @ Enzyme.Compiler ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:1797 [2] macro expansion @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:7113 [inlined] [3] deferred_codegen (repeats 3 times) @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:7072 [inlined] [4] autodiff_deferred @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:524 [inlined] [5] autodiff_deferred @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:538 [inlined] [6] dg @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:78 [inlined] [7] dg @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0 [inlined] [8] fwddiffejulia_dg_5482_inner_1wrap @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0 [inlined] [9] macro expansion @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6837 [inlined] [10] enzyme_call @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6437 [inlined] [11] ForwardModeThunk @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6317 [inlined] [12] autodiff_deferred @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:526 [inlined] [13] autodiff_deferred @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:538 [inlined] [14] ddg @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:79 [inlined] [15] ddg @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0 [inlined] [16] fwddiffejulia_ddg_5479_inner_1wrap @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:0 [17] macro expansion @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6837 [inlined] [18] enzyme_call @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6437 [inlined] [19] ForwardModeThunk @ ~/.julia/packages/Enzyme/fS5Q3/src/compiler.jl:6317 [inlined] [20] autodiff_deferred @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:526 [inlined] [21] autodiff_deferred @ ~/.julia/packages/Enzyme/fS5Q3/src/Enzyme.jl:538 [inlined] [22] dddg(x::SVector{2, Float64}, dx::SVector{2, Float64}, ddx::SVector{2, Float64}, dddx::SVector{2, Float64}) @ Main ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:80 [23] top-level scope @ ~/repositories/2024_ad_lw_dev/TenkaiAD/misc/third_derivative_test_arpit.jl:85 ```

I am using Enzyme v0.12.27.

versioninfo() Julia Version 1.10.3 Commit 0b4590a5507 (2024-04-30 10:59 UTC) Build Info: Official https://julialang.org/ release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 80 × Intel(R) Xeon(R) Gold 6148 CPU @ 2.40GHz WORD_SIZE: 64 LIBM: libopenlibm LLVM: libLLVM-15.0.7 (ORCJIT, skylake-avx512) Threads: 1 default, 0 interactive, 1 GC (on 80 virtual cores) Environment: JULIA_EDITOR = JULIA_NUM_THREADS = 1
vchuravy commented 3 months ago

This is likely the same as https://github.com/EnzymeAD/Enzyme.jl/issues/1173, the goal is that https://github.com/JuliaGPU/GPUCompiler.jl/pull/599 will fix this.

wsmoses commented 2 months ago

Closing in favor of https://github.com/JuliaGPU/GPUCompiler.jl/issues/629