[Forward-mode] Implement multi-typed store

[sethaxen]

The following function that takes and returns a struct works fine with Float64 eltype within the struct and works fine if we pass the function bar itself to autodiff, but it errors if we use eltype Float32 and pass an anonymous function that calls bar to autodiff. If the return type is instead Const, then it segfaults.

using Enzyme

struct Foo{X,A}
    x::X
    a::A
end
bar(x::Foo) = Foo(x.a .* x.x, 2 * x.a)

x = Foo([1f0, 2f0, 3f0], 4f0)
∂x = Foo([5f0, 6f0, 7f0], 8f0)

autodiff(Forward, bar, Duplicated, Duplicated(x, ∂x))
# returns: (var"1" = Foo{Vector{Float32}, Float32}(Float32[4.0, 8.0, 12.0], 8.0f0), var"2" = Foo{Vector{Float32}, Float32}(Float32[28.0, 40.0, 52.0], 16.0f0))
autodiff(Forward, x -> bar(x), Duplicated, Duplicated(x, ∂x))
# errors, see attached

dump.txt

[wsmoses]

Minor, but this is a Julia error, not a segfault or assertion error

…

On Thu, May 18, 2023 at 1:42 PM Seth Axen ***@***.***> wrote: The following function that takes and returns a struct works fine with Float64 eltype within the struct and works fine if we pass the function bar itself to autodiff, but it errors if we use eltype Float32 and pass an anonymous function that calls bar to autodiff. If the return type is instead Const, then it segfaults. using Enzyme struct Foo{X,A} x::X a::Aendbar(x::Foo) = Foo(x.a .* x.x, 2 * x.a) x = Foo([1f0, 2f0, 3f0], 4f0) ∂x = Foo([5f0, 6f0, 7f0], 8f0) autodiff(Forward, bar, Duplicated, Duplicated(x, ∂x))# returns: (var"1" = Foo{Vector{Float32}, Float32}(Float32[4.0, 8.0, 12.0], 8.0f0), var"2" = Foo{Vector{Float32}, Float32}(Float32[28.0, 40.0, 52.0], 16.0f0))autodiff(Forward, x -> bar(x), Duplicated, Duplicated(x, ∂x))# errors, see attached dump.txt <https://github.com/EnzymeAD/Enzyme.jl/files/11510087/dump.txt> — Reply to this email directly, view it on GitHub <#857>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXCHR2YYYZE5QIJJUNDXGZNPPANCNFSM6AAAAAAYGYEKIU> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[sethaxen]

Minor, but this is a Julia error, not a segfault or assertion error

Ah, sorry, yes, if we change the return type to Const, Julia errors but does not segfault:

julia: /workspace/srcdir/Enzyme/enzyme/Enzyme/CApi.cpp:1237: void EnzymeFixupJuliaCallingConvention(LLVMValueRef): Assertion `Types.size()' failed.

[913111] signal (6.-6): Aborted
in expression starting at REPL[9]:1
pthread_kill at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
raise 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: 0x7fb44362871a)
__assert_fail at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
EnzymeFixupJuliaCallingConvention at /workspace/srcdir/Enzyme/enzyme/Enzyme/CApi.cpp:1237
EnzymeFixupJuliaCallingConvention at /home/sethaxen/.julia/packages/Enzyme/tHbH4/src/api.jl:541 [inlined]
post_optimze! at /home/sethaxen/.julia/packages/Enzyme/tHbH4/src/compiler/optimize.jl:1044
post_optimze! at /home/sethaxen/.julia/packages/Enzyme/tHbH4/src/compiler/optimize.jl:1041
unknown function (ip: 0x7fb40bf622b1)
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2758 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
_thunk at /home/sethaxen/.julia/packages/Enzyme/tHbH4/src/compiler.jl:9571
_thunk at /home/sethaxen/.julia/packages/Enzyme/tHbH4/src/compiler.jl:9549 [inlined]
cached_compilation at /home/sethaxen/.julia/packages/Enzyme/tHbH4/src/compiler.jl:9587 [inlined]
#s301#222 at /home/sethaxen/.julia/packages/Enzyme/tHbH4/src/compiler.jl:9645 [inlined]
#s301#222 at ./none:0
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2758 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
GeneratedFunctionStub at ./boot.jl:602
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2758 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
jl_apply at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/julia.h:1879 [inlined]
jl_call_staged at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/method.c:530
ijl_code_for_staged at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/method.c:581
get_staged at ./compiler/utilities.jl:115
retrieve_code_info at ./compiler/utilities.jl:127 [inlined]
InferenceState at ./compiler/inferencestate.jl:354
typeinf_edge at ./compiler/typeinfer.jl:924
abstract_call_method at ./compiler/abstractinterpretation.jl:611
abstract_call_gf_by_type at ./compiler/abstractinterpretation.jl:152
abstract_call_known at ./compiler/abstractinterpretation.jl:1949
abstract_call at ./compiler/abstractinterpretation.jl:2020
abstract_call at ./compiler/abstractinterpretation.jl:1999
abstract_eval_statement_expr at ./compiler/abstractinterpretation.jl:2183
abstract_eval_statement at ./compiler/abstractinterpretation.jl:2396
abstract_eval_basic_statement at ./compiler/abstractinterpretation.jl:2684
typeinf_local at ./compiler/abstractinterpretation.jl:2869
typeinf_nocycle at ./compiler/abstractinterpretation.jl:2957
_typeinf at ./compiler/typeinfer.jl:246
typeinf at ./compiler/typeinfer.jl:216
typeinf_edge at ./compiler/typeinfer.jl:933
abstract_call_method at ./compiler/abstractinterpretation.jl:611
abstract_call_gf_by_type at ./compiler/abstractinterpretation.jl:152
abstract_call_known at ./compiler/abstractinterpretation.jl:1949
abstract_call at ./compiler/abstractinterpretation.jl:2020
abstract_call at ./compiler/abstractinterpretation.jl:1999
abstract_eval_statement_expr at ./compiler/abstractinterpretation.jl:2183
abstract_eval_statement at ./compiler/abstractinterpretation.jl:2396
abstract_eval_basic_statement at ./compiler/abstractinterpretation.jl:2660
typeinf_local at ./compiler/abstractinterpretation.jl:2869
typeinf_nocycle at ./compiler/abstractinterpretation.jl:2957
_typeinf at ./compiler/typeinfer.jl:246
typeinf at ./compiler/typeinfer.jl:216
typeinf_ext at ./compiler/typeinfer.jl:1058
typeinf_ext_toplevel at ./compiler/typeinfer.jl:1091
typeinf_ext_toplevel at ./compiler/typeinfer.jl:1087
jfptr_typeinf_ext_toplevel_14294.clone_1 at /home/sethaxen/.julia/juliaup/julia-1.9.0+0.x64.linux.gnu/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2758 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
jl_apply at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/julia.h:1879 [inlined]
jl_type_infer at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:320
jl_generate_fptr_impl at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/jitlayers.cpp:444
jl_compile_method_internal at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2348 [inlined]
jl_compile_method_internal at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2237
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2750 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
autodiff at /home/sethaxen/.julia/packages/Enzyme/tHbH4/src/Enzyme.jl:214
unknown function (ip: 0x7fb40bf6b857)
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2758 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
jl_apply at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/julia.h:1879 [inlined]
do_call at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/interpreter.c:126
eval_value at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/interpreter.c:226
eval_stmt_value at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/interpreter.c:177 [inlined]
eval_body at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/interpreter.c:624
jl_interpret_toplevel_thunk at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/interpreter.c:762
jl_toplevel_eval_flex at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/toplevel.c:912
jl_toplevel_eval_flex at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/toplevel.c:856
ijl_toplevel_eval_in at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/toplevel.c:971
eval at ./boot.jl:370 [inlined]
eval_user_input at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:153
repl_backend_loop at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:249
#start_repl_backend#46 at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:234
start_repl_backend at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:231
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2758 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
#run_repl#59 at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:379
run_repl at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:365
jfptr_run_repl_60276.clone_1 at /home/sethaxen/.julia/juliaup/julia-1.9.0+0.x64.linux.gnu/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2758 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
#1017 at ./client.jl:421
jfptr_YY.1017_32248.clone_1 at /home/sethaxen/.julia/juliaup/julia-1.9.0+0.x64.linux.gnu/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2758 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
jl_apply at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/julia.h:1879 [inlined]
jl_f__call_latest at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/builtins.c:774
#invokelatest#2 at ./essentials.jl:816 [inlined]
invokelatest at ./essentials.jl:813 [inlined]
run_main_repl at ./client.jl:405
exec_options at ./client.jl:322
_start at ./client.jl:522
jfptr__start_37386.clone_1 at /home/sethaxen/.julia/juliaup/julia-1.9.0+0.x64.linux.gnu/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2758 [inlined]
ijl_apply_generic at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/gf.c:2940
jl_apply at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/julia.h:1879 [inlined]
true_main at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/jlapi.c:573
jl_repl_entrypoint at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/src/jlapi.c:717
main at /cache/build/default-amdci4-0/julialang/julia-release-1-dot-9/cli/loader_exe.c:59
unknown function (ip: 0x7fb443629d8f)
__libc_start_main at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x401098)
Allocations: 22922468 (Pool: 22893127; Big: 29341); GC: 32
Aborted (core dumped)

[wsmoses]

The return of const issue should now be resolved.

[wsmoses]

This will be fixed by EnzymeAD/Enzyme#1564 but requires a jll bump before being available here