search

EnzymeAD / Enzyme.jl

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

Higher order AD error #1487

Open jgreener64 opened 1 month ago

jgreener64 commented 1 month ago

I am on Enzyme main (b5addb62c0c9d0d74825be5fc71deebfd26d793c), StaticArrays 1.9.4 and Julia 1.10.3. The following errors:

using Enzyme, StaticArrays

struct TriclinicBoundary{T}
    basis_vectors::T
end

function potential_energy(coords, boundary)
    E = 0.0
    for i in 1:length(coords), j in 1:length(coords)
        E += sum(abs, coords[i] - coords[j])
    end
    return E
end

function scale_frac_coord(fc, boundary)
    bv = boundary.basis_vectors
    return SVector(
        fc[1] * bv[1][1] + fc[2] * bv[2][1] + fc[3] * bv[3][1],
        fc[1] * bv[1][2] + fc[2] * bv[2][2] + fc[3] * bv[3][2],
        fc[1] * bv[1][3] + fc[2] * bv[2][3] + fc[3] * bv[3][3],
    )
end

function potential_energy_frac_coords(frac_coords, boundary)
    coords = zero(frac_coords)
    for i in 1:length(coords)
        coords[i] = scale_frac_coord(frac_coords[i], boundary)
    end
    return potential_energy(coords, boundary)
end

function potential_energy_eps(coords, boundary, epsilon)
    frac_coords = coords ./ 2.0
    bv, ep = boundary.basis_vectors, epsilon
    boundary_eps = TriclinicBoundary([
        SVector(
            bv[1][1] * (ep[1, 1] + 1.0) + bv[1][2] * ep[2, 1] + bv[1][3] * ep[3, 1],
            bv[1][1] * ep[1, 2] + bv[1][2] * (ep[2, 2] + 1.0) + bv[1][3] * ep[3, 2],
            bv[1][1] * ep[1, 3] + bv[1][2] * ep[2, 3] + bv[1][3] * (ep[3, 3] + 1.0),
        ),
        SVector(
            bv[2][1] * (ep[1, 1] + 1.0) + bv[2][2] * ep[2, 1] + bv[2][3] * ep[3, 1],
            bv[2][1] * ep[1, 2] + bv[2][2] * (ep[2, 2] + 1.0) + bv[2][3] * ep[3, 2],
            bv[2][1] * ep[1, 3] + bv[2][2] * ep[2, 3] + bv[2][3] * (ep[3, 3] + 1.0),
        ),
        SVector(
            bv[3][1] * (ep[1, 1] + 1.0) + bv[3][2] * ep[2, 1] + bv[3][3] * ep[3, 1],
            bv[3][1] * ep[1, 2] + bv[3][2] * (ep[2, 2] + 1.0) + bv[3][3] * ep[3, 2],
            bv[3][1] * ep[1, 3] + bv[3][2] * ep[2, 3] + bv[3][3] * (ep[3, 3] + 1.0),
        ),
    ])
    return potential_energy_frac_coords(frac_coords, boundary_eps)
end

function dU_dϵ(coords, boundary, epsilon, i, j)
    d_epsilon = zeros(3, 3)
    d_epsilon[i, j] = 1.0
    return autodiff_deferred(
        Enzyme.Forward,
        potential_energy_eps,
        Duplicated,
        Const(coords),
        Const(boundary),
        Duplicated(epsilon, d_epsilon),
    )[2]
end

function dU_dϵ(coords, boundary, epsilon)
    dU_dϵ_out = zeros(3, 3)
    for i in 1:3, j in 1:3
        dU_dϵ_out[i, j] = dU_dϵ(coords, boundary, epsilon, i, j)
    end
    return dU_dϵ_out
end

function loss(coords, boundary, epsilon)
    return sum(dU_dϵ(coords, boundary, epsilon))
end

coords = rand(SVector{3, Float64}, 4) .* 2.0
d_coords = zero(coords)
boundary = TriclinicBoundary([
    SVector(2.0  , 0.0  , 0.0),
    SVector(1e-12, 2.0  , 0.0),
    SVector(1e-12, 1e-12, 2.0),
])
epsilon = zeros(3, 3)

loss(coords, boundary, epsilon) # Works

autodiff(
    Enzyme.Reverse,
    loss,
    Active,
    Duplicated(coords, d_coords),
    Const(boundary),
    Const(epsilon),
)
┌ Error: Enzyme aligned size and Julia size disagree
│   AlignedSize = 512
│   esizeof(TT) = 488
│   fieldtypes(TT) = (Any, Any, @NamedTuple{1, 2, 3::Tuple{Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Tuple{Float64, Float64}, 0}, Core.LLVMPtr{Float64, 0}, Core.LLVMPtr{Tuple{Float64, Float64}, 0}, Core.LLVMPtr{Float64, 0}, Vararg{Core.LLVMPtr{Tuple{Float64, Float64}, 0}, 4}}, 4::UInt64, 5::Core.LLVMPtr{Float64, 0}, 6::Core.LLVMPtr{Float64, 0}, 7::Core.LLVMPtr{Float64, 0}, 8::Core.LLVMPtr{Float64, 0}, 9::Core.LLVMPtr{Float64, 0}, 10::Core.LLVMPtr{Float64, 0}, 11::Core.LLVMPtr{Float64, 0}, 12::Core.LLVMPtr{Float64, 0}, 13::Core.LLVMPtr{Float64, 0}, 14::Core.LLVMPtr{Float64, 0}, 15::Core.LLVMPtr{Float64, 0}, 16::Core.LLVMPtr{Float64, 0}, 17::Core.LLVMPtr{Float64, 0}, 18::Core.LLVMPtr{Float64, 0}, 19::Core.LLVMPtr{Float64, 0}, 20::Core.LLVMPtr{Float64, 0}, 21::Core.LLVMPtr{Tuple{Float64, Float64}, 0}, 22::Core.LLVMPtr{Tuple{Float64, Float64}, 0}, 23::Core.LLVMPtr{Tuple{Float64, Float64}, 0}, 24::Core.LLVMPtr{Tuple{Float64, Float64}, 0}}, Any, Any, Any, Any, UInt64, Core.LLVMPtr{NTuple{24, Float64}, 0}, Float64, UInt64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Tuple{Float64, Float64}, Float64, Float64, NTuple{4, Float64}, Tuple{Float64, Float64})
└ @ Enzyme.Compiler ~/.julia/dev/GPUCompiler/src/utils.jl:59
julia: /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/Support/Casting.h:578: decltype(auto) llvm::cast(From*) [with To = llvm::Instruction; From = llvm::Value]: Assertion `isa<To>(Val) && "cast<Ty>() argument of incompatible type!"' failed.

[2054785] signal (6.-6): Aborted
in expression starting at /home/jgreener/dms/molly_dev/enzyme_err39.jl:95
gsignal at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
abort at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x7fc37ddc0728)
__assert_fail at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
cast<llvm::Instruction, llvm::Value> at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/Support/Casting.h:578
cast<llvm::Instruction, llvm::Value> at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/Value.h:533 [inlined]
CreateAllocation at /workspace/srcdir/Enzyme/enzyme/Enzyme/Utils.cpp:309
createCacheForScope at /workspace/srcdir/Enzyme/enzyme/Enzyme/CacheUtility.cpp:815
ensureLookupCached at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:2408 [inlined]
ensureLookupCached at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:2394
cacheForReverse at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:2985
recursivelyHandleSubfunction at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:5521
visitCallInst at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:6512
visit at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/InstVisitor.h:111 [inlined]
CreateAugmentedPrimal at /workspace/srcdir/Enzyme/enzyme/Enzyme/EnzymeLogic.cpp:2544
recursivelyHandleSubfunction at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:5394
visitCallInst at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:6512
visit at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/InstVisitor.h:111 [inlined]
CreatePrimalAndGradient at /workspace/srcdir/Enzyme/enzyme/Enzyme/EnzymeLogic.cpp:4420
EnzymeCreatePrimalAndGradient at /workspace/srcdir/Enzyme/enzyme/Enzyme/CApi.cpp:622
EnzymeCreatePrimalAndGradient at /home/jgreener/.julia/dev/Enzyme/src/api.jl:154
unknown function (ip: 0x7fc24c614bfb)
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
enzyme! at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:3292
unknown function (ip: 0x7fc24c603118)
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
#codegen#1173 at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:5250
codegen at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:4592 [inlined]
_thunk at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:5968
_thunk at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:5968 [inlined]
cached_compilation at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:6006 [inlined]
#1220 at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:6073
#JuliaContext#147 at /home/jgreener/.julia/dev/GPUCompiler/src/driver.jl:52
unknown function (ip: 0x7fc24c5d5e96)
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
JuliaContext at /home/jgreener/.julia/dev/GPUCompiler/src/driver.jl:42
#s1957#1219 at /home/jgreener/.julia/dev/Enzyme/src/compiler.jl:6024 [inlined]
#s1957#1219 at ./none:0
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
GeneratedFunctionStub at ./boot.jl:602
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
jl_call_staged at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/method.c:540
ijl_code_for_staged at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/method.c:593
get_staged at ./compiler/utilities.jl:123
retrieve_code_info at ./compiler/utilities.jl:135 [inlined]
InferenceState at ./compiler/inferencestate.jl:430
typeinf_edge at ./compiler/typeinfer.jl:920
abstract_call_method at ./compiler/abstractinterpretation.jl:629
abstract_call_gf_by_type at ./compiler/abstractinterpretation.jl:95
abstract_call_known at ./compiler/abstractinterpretation.jl:2087
abstract_call at ./compiler/abstractinterpretation.jl:2169
abstract_call at ./compiler/abstractinterpretation.jl:2162
abstract_call at ./compiler/abstractinterpretation.jl:2354
abstract_eval_call at ./compiler/abstractinterpretation.jl:2370
abstract_eval_statement_expr at ./compiler/abstractinterpretation.jl:2380
abstract_eval_statement at ./compiler/abstractinterpretation.jl:2624
abstract_eval_basic_statement at ./compiler/abstractinterpretation.jl:2889
typeinf_local at ./compiler/abstractinterpretation.jl:3098
typeinf_nocycle at ./compiler/abstractinterpretation.jl:3186
_typeinf at ./compiler/typeinfer.jl:247
typeinf at ./compiler/typeinfer.jl:216
typeinf_edge at ./compiler/typeinfer.jl:930
abstract_call_method at ./compiler/abstractinterpretation.jl:629
abstract_call_gf_by_type at ./compiler/abstractinterpretation.jl:95
abstract_call_known at ./compiler/abstractinterpretation.jl:2087
abstract_call at ./compiler/abstractinterpretation.jl:2169
abstract_apply at ./compiler/abstractinterpretation.jl:1612
abstract_call_known at ./compiler/abstractinterpretation.jl:2004
abstract_call at ./compiler/abstractinterpretation.jl:2169
abstract_call at ./compiler/abstractinterpretation.jl:2162
abstract_call at ./compiler/abstractinterpretation.jl:2354
abstract_eval_call at ./compiler/abstractinterpretation.jl:2370
abstract_eval_statement_expr at ./compiler/abstractinterpretation.jl:2380
abstract_eval_statement at ./compiler/abstractinterpretation.jl:2624
abstract_eval_basic_statement at ./compiler/abstractinterpretation.jl:2913
typeinf_local at ./compiler/abstractinterpretation.jl:3098
typeinf_nocycle at ./compiler/abstractinterpretation.jl:3186
_typeinf at ./compiler/typeinfer.jl:247
typeinf at ./compiler/typeinfer.jl:216
typeinf_ext at ./compiler/typeinfer.jl:1051
typeinf_ext_toplevel at ./compiler/typeinfer.jl:1082
typeinf_ext_toplevel at ./compiler/typeinfer.jl:1078
jfptr_typeinf_ext_toplevel_35682.1 at /home/jgreener/soft/julia/julia-1.10.3/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
jl_apply at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
jl_type_infer at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:394
jl_generate_fptr_impl at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/jitlayers.cpp:504
jl_compile_method_internal at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2481 [inlined]
jl_compile_method_internal at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2368
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2887 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
jl_apply at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
do_call at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/interpreter.c:126
eval_value at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/interpreter.c:223
eval_stmt_value at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/interpreter.c:174 [inlined]
eval_body at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/interpreter.c:617
jl_interpret_toplevel_thunk at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/interpreter.c:775
jl_toplevel_eval_flex at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/toplevel.c:934
jl_toplevel_eval_flex at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/toplevel.c:877
ijl_toplevel_eval_in at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/toplevel.c:985
eval at ./boot.jl:385 [inlined]
include_string at ./loading.jl:2076
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
_include at ./loading.jl:2136
include at ./Base.jl:495
jfptr_include_46394.1 at /home/jgreener/soft/julia/julia-1.10.3/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
exec_options at ./client.jl:318
_start at ./client.jl:552
jfptr__start_82726.1 at /home/jgreener/soft/julia/julia-1.10.3/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:2895 [inlined]
ijl_apply_generic at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/gf.c:3077
jl_apply at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
true_main at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/jlapi.c:582
jl_repl_entrypoint at /cache/build/builder-amdci4-2/julialang/julia-release-1-dot-10/src/jlapi.c:731
main at julia (unknown line)
__libc_start_main at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x4010b8)
Allocations: 41014975 (Pool: 40949909; Big: 65066); GC: 51
Aborted (core dumped)

Possibly related to https://github.com/EnzymeAD/Enzyme.jl/issues/999.

wsmoses commented 1 month ago

Okay this is indeed a minimized version of https://github.com/EnzymeAD/Enzyme.jl/issues/999

@vchuravy any insights here?

wsmoses commented 1 month ago

is this just a layout padding thing

jgreener64 commented 1 week ago

@vchuravy could you look at this when you get a chance?

vchuravy commented 1 week ago

I am defending and then moving in July (and then starting a new position in August)

So I am not likely to find free time soon.

jgreener64 commented 1 week ago

Good luck with the defence!