Enzyme.jl icon indicating copy to clipboard operation
Enzyme.jl copied to clipboard

Vector support for generic call

Open stelmo opened this issue 3 years ago • 10 comments
trafficstars

Hi, I am trying to find the Jacobian of a vector function, but no matter what I try, jacobian crashes Julia. A minimal working example is shown below:

using Enzyme
using SparseArrays, LinearAlgebra

N = 10
_S = sprand(N, N, 0.1) 
_Q = sprand(N, N, 0.1)

F(z) = [ # some complicated sparse amenable function
    _S * z 
    _Q * z.^2
]
z = rand(N)

jac = Enzyme.jacobian(Forward, F, z) # crash

# Dense version also fails

N = 10
_S = rand(N, N) 
_Q = rand(N, N)

F(z) = [
    _S * z 
    _Q * z.^2
]
z = rand(N)

jac = Enzyme.jacobian(Forward, F, z) # crash

I am using this [7da242da] Enzyme v0.10.0-dev https://github.com/EnzymeAD/Enzyme.jl.git#main version of Enzyme, which should be the latest. Is there a workaround or should I just be patient? :)

stelmo avatar May 12 '22 10:05 stelmo

Can you post the full log as a Gist?

vchuravy avatar May 12 '22 11:05 vchuravy

sure thing, here it is:

julia> jac = Enzyme.jacobian(Forward, F, z) # crash
 o:   %7 = call cc37 nonnull {} addrspace(10)* bitcast ({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)* @jl_apply_generic to {} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)*, {} addrspace(10)*)*)({} addrspace(10)* addrspacecast ({}* inttoptr (i64 327951328 to {}*) to {} addrspace(10)*), {} addrspace(10)* %6, {} addrspace(10)* nonnull %0) #8, !dbg !17
 ot: {} addrspace(10)*
 ir:   %15 = load {} addrspace(10)*, {} addrspace(10)** %14, align 8, !dbg !17
 irt: {} addrspace(10)*
 p:   %"'dual_phi" = phi [10 x {} addrspace(10)*] , !dbg !26
 PT: [10 x {} addrspace(10)*]
 newCall:   %18 = call cc37 nonnull {} addrspace(10)* bitcast ({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)* @jl_apply_generic to {} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)*, {} addrspace(10)*)*)({} addrspace(10)* addrspacecast ({}* inttoptr (i64 327951328 to {}*) to {} addrspace(10)*), {} addrspace(10)* %3, {} addrspace(10)* nonnull %0) #8, !dbg !17
 newCallT: {} addrspace(10)*
Assertion failed: invertedReturn->getType() == gutils->getShadowType(orig->getType()), file /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h, line 8262

signal (22): SIGABRT
in expression starting at REPL[13]:1
crt_sig_handler at /cygdrive/c/buildbot/worker/package_win64/build/src\signals-win.c:92
raise at C:\Windows\System32\msvcrt.dll (unknown line)
abort at C:\Windows\System32\msvcrt.dll (unknown line)
assert at C:\Windows\System32\msvcrt.dll (unknown line)
visitCallInst at /workspace/srcdir/Enzyme/enzyme/Enzyme\AdjointGenerator.h:8261
delegateCallInst at /opt/x86_64-w64-mingw32/x86_64-w64-mingw32/sys-root/usr/local/include/llvm/IR\InstVisitor.h:299 [inlined]
visitCall at /opt/x86_64-w64-mingw32/x86_64-w64-mingw32/sys-root/usr/local/include/llvm/IR\Instruction.def:209 [inlined]
visit at /opt/x86_64-w64-mingw32/x86_64-w64-mingw32/sys-root/usr/local/include/llvm/IR\Instruction.def:209
visit at /opt/x86_64-w64-mingw32/x86_64-w64-mingw32/sys-root/usr/local/include/llvm/IR\InstVisitor.h:112 [inlined]
CreateForwardDiff at /workspace/srcdir/Enzyme/enzyme/Enzyme\EnzymeLogic.cpp:4227
EnzymeCreateForwardDiff at /workspace/srcdir/Enzyme/enzyme/Enzyme\CApi.cpp:414
EnzymeCreateForwardDiff at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\api.jl:125
enzyme! at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\compiler.jl:3172
unknown function (ip: 000000000a905c70)
#codegen#73 at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\compiler.jl:3930
codegen##kw at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\compiler.jl:3661 [inlined]
_thunk at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\compiler.jl:4337
unknown function (ip: 00000000631fec50)
cached_compilation at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\compiler.jl:4375
unknown function (ip: 00000000631e24f7)
#s502#108 at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\compiler.jl:4435 [inlined]
#s502#108 at .\none:0
GeneratedFunctionStub at .\boot.jl:580
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
jl_call_staged at /cygdrive/c/buildbot/worker/package_win64/build/src\method.c:431
jl_code_for_staged at /cygdrive/c/buildbot/worker/package_win64/build/src\method.c:482
get_staged at .\compiler\utilities.jl:111
retrieve_code_info at .\compiler\utilities.jl:123 [inlined]
InferenceState at .\compiler\inferencestate.jl:234
typeinf_edge at .\compiler\typeinfer.jl:814 [inlined]
abstract_call_method at .\compiler\abstractinterpretation.jl:504
abstract_call_gf_by_type at .\compiler\abstractinterpretation.jl:105
abstract_call_known at .\compiler\abstractinterpretation.jl:1342
abstract_call at .\compiler\abstractinterpretation.jl:1397
abstract_call at .\compiler\abstractinterpretation.jl:1382
abstract_eval_statement at .\compiler\abstractinterpretation.jl:1534
typeinf_local at .\compiler\abstractinterpretation.jl:1918
typeinf_nocycle at .\compiler\abstractinterpretation.jl:2014
_typeinf at .\compiler\typeinfer.jl:226
typeinf at .\compiler\typeinfer.jl:209
typeinf_edge at .\compiler\typeinfer.jl:823 [inlined]
abstract_call_method at .\compiler\abstractinterpretation.jl:504
abstract_call_gf_by_type at .\compiler\abstractinterpretation.jl:105
abstract_call_known at .\compiler\abstractinterpretation.jl:1342
abstract_call at .\compiler\abstractinterpretation.jl:1397
abstract_call at .\compiler\abstractinterpretation.jl:1382
abstract_eval_statement at .\compiler\abstractinterpretation.jl:1534
typeinf_local at .\compiler\abstractinterpretation.jl:1918
typeinf_nocycle at .\compiler\abstractinterpretation.jl:2014
_typeinf at .\compiler\typeinfer.jl:226
typeinf at .\compiler\typeinfer.jl:209
typeinf_edge at .\compiler\typeinfer.jl:823 [inlined]
abstract_call_method at .\compiler\abstractinterpretation.jl:504
abstract_call_gf_by_type at .\compiler\abstractinterpretation.jl:105
abstract_call_known at .\compiler\abstractinterpretation.jl:1342
abstract_call at .\compiler\abstractinterpretation.jl:1397
abstract_call at .\compiler\abstractinterpretation.jl:1382
abstract_eval_statement at .\compiler\abstractinterpretation.jl:1534
typeinf_local at .\compiler\abstractinterpretation.jl:1900
typeinf_nocycle at .\compiler\abstractinterpretation.jl:2014
_typeinf at .\compiler\typeinfer.jl:226
typeinf at .\compiler\typeinfer.jl:209
typeinf_ext at .\compiler\typeinfer.jl:909
typeinf_ext_toplevel at .\compiler\typeinfer.jl:942
typeinf_ext_toplevel at .\compiler\typeinfer.jl:938
jfptr_typeinf_ext_toplevel_16182.clone_1 at C:\Users\stelmo\AppData\Local\Programs\Julia-1.7.2\lib\julia\sys.dll (unknown line)
_jl_invoke at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:2247 [inlined]
jl_apply_generic at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:2429 [inlined]
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
jl_type_infer at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:295
jl_generate_fptr at /cygdrive/c/buildbot/worker/package_win64/build/src\jitlayers.cpp:338
jl_compile_method_internal at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:1980
jl_compile_method_internal at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:1934 [inlined]
_jl_invoke at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:2239 [inlined]
jl_apply_generic at /cygdrive/c/buildbot/worker/package_win64/build/src\gf.c:2429
#gradient#32 at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\Enzyme.jl:714 [inlined]
gradient at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\Enzyme.jl:714
#jacobian#41 at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\Enzyme.jl:787 [inlined]
jacobian at C:\Users\stelmo\.julia\packages\Enzyme\AxAzi\src\Enzyme.jl:787
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
do_call at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:126
eval_value at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:215
eval_stmt_value at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:166 [inlined]
eval_body at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:583
jl_interpret_toplevel_thunk at /cygdrive/c/buildbot/worker/package_win64/build/src\interpreter.c:731
jl_toplevel_eval_flex at /cygdrive/c/buildbot/worker/package_win64/build/src\toplevel.c:885
jl_toplevel_eval_flex at /cygdrive/c/buildbot/worker/package_win64/build/src\toplevel.c:830
jl_toplevel_eval_flex at /cygdrive/c/buildbot/worker/package_win64/build/src\toplevel.c:830
jl_toplevel_eval at /cygdrive/c/buildbot/worker/package_win64/build/src\toplevel.c:894 [inlined]
jl_toplevel_eval_in at /cygdrive/c/buildbot/worker/package_win64/build/src\toplevel.c:944
eval at .\boot.jl:373 [inlined]
eval_user_input at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:150
repl_backend_loop at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:246
start_repl_backend at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:231
#run_repl#47 at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:364
run_repl at C:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.7\REPL\src\REPL.jl:351
#930 at .\client.jl:394
jfptr_YY.930_36349.clone_1 at C:\Users\stelmo\AppData\Local\Programs\Julia-1.7.2\lib\julia\sys.dll (unknown line)
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
jl_f__call_latest at /cygdrive/c/buildbot/worker/package_win64/build/src\builtins.c:757
#invokelatest#2 at .\essentials.jl:716 [inlined]
invokelatest at .\essentials.jl:714 [inlined]
run_main_repl at .\client.jl:379
exec_options at .\client.jl:309
_start at .\client.jl:495
jfptr__start_21275.clone_1 at C:\Users\stelmo\AppData\Local\Programs\Julia-1.7.2\lib\julia\sys.dll (unknown line)
jl_apply at /cygdrive/c/buildbot/worker/package_win64/build/src\julia.h:1788 [inlined]
true_main at /cygdrive/c/buildbot/worker/package_win64/build/src\jlapi.c:559
jl_repl_entrypoint at /cygdrive/c/buildbot/worker/package_win64/build/src\jlapi.c:701
mainCRTStartup at /cygdrive/c/buildbot/worker/package_win64/build/cli\loader_exe.c:42
BaseThreadInitThunk at C:\Windows\System32\KERNEL32.DLL (unknown line)
RtlUserThreadStart at C:\Windows\SYSTEM32\ntdll.dll (unknown line)
Allocations: 46332246 (Pool: 46310729; Big: 21517); GC: 52

stelmo avatar May 12 '22 17:05 stelmo

The issue here is that we haven’t added vector mode to generic calls yet. Can you specify a chunk size of 1

On Thu, May 12, 2022 at 6:13 PM Valentin Churavy @.***> wrote:

Assigned #304 https://github.com/EnzymeAD/Enzyme.jl/issues/304 to @wsmoses https://github.com/wsmoses.

— Reply to this email directly, view it on GitHub https://github.com/EnzymeAD/Enzyme.jl/issues/304#event-6600121339, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXAMMRD3EZB34LRYXW3VJU33FANCNFSM5VX2O7XQ . You are receiving this because you were assigned.Message ID: @.***>

wsmoses avatar May 12 '22 17:05 wsmoses

julia> jac = Enzyme.jacobian(Forward, F, z, Val(1))
([0.0, 0.0, 0.9309410649689392, 0.19161192949503747, 1.395478078742109, 0.0, 0.0, 0.14002672400024546, 0.14227937225880918, 0.18290817583778968, 0.3712755068969665, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.08431178356516157, 0.40517985161797276, 0.4412912824781923], [0.0, 0.0, 1.8730243855277187, 0.19161192949503747, 0.5772122993893988, 0.0, 0.0, 0.14002672400024546, 0.14227937225880918, 0.18290817583778968, 0.3712755068969665, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.08431178356516157, 0.40517985161797276, 0.4412912824781923], [0.0, 0.0, 0.9309410649689392, 0.19161192949503747, 0.5772122993893988, 0.0, 0.0, 0.14002672400024546, 0.14227937225880918, 0.18290817583778968, 0.3712755068969665, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.08431178356516157, 0.40517985161797276, 0.4412912824781923], [0.0, 0.0, 0.9309410649689392, 0.19161192949503747, 0.5772122993893988, 0.0, 0.0, 0.14002672400024546, 0.14227937225880918, 0.18290817583778968, 0.3712755068969665, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.08431178356516157, 0.40517985161797276, 0.5869037781534985], [0.0, 0.0, 0.9309410649689392, 0.19161192949503747, 0.5772122993893988, 0.0, 0.0, 0.14002672400024546, 0.14227937225880918, 0.18290817583778968, 1.1788039954062128, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.08431178356516157, 0.40517985161797276, 0.4412912824781923], [0.0, 0.0, 0.9309410649689392, 0.19161192949503747, 1.317231352194884, 0.0, 0.0, 0.14002672400024546, 0.14227937225880918, 0.35919416790494735, 0.3712755068969665, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.08431178356516157, 1.447559234266814, 1.56022751123051], [0.0, 0.0, 0.9309410649689392, 0.19161192949503747, 0.5772122993893988, 0.0, 0.0, 0.14002672400024546, 0.14227937225880918, 0.18290817583778968, 0.3712755068969665, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.08431178356516157, 0.40517985161797276, 0.4412912824781923], [0.0, 0.0, 0.9309410649689392, 0.19161192949503747, 0.5772122993893988, 0.0, 0.0, 0.14002672400024546, 0.14227937225880918, 0.18290817583778968, 0.3712755068969665, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.08431178356516157, 0.40517985161797276, 0.4412912824781923], [0.0, 0.0, 0.9309410649689392, 0.7487083137329062, 0.5772122993893988, 0.0, 0.0, 0.14002672400024546, 0.14227937225880918, 0.316245611989353, 0.4278441883795144, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5745713399151094, 0.40517985161797276, 0.4412912824781923], [0.0, 0.0, 0.9309410649689392, 0.19161192949503747, 0.5772122993893988, 0.0, 0.0, 1.1176845169106329, 1.1356650138380477, 0.18290817583778968, 0.3712755068969665, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.08431178356516157, 0.40517985161797276, 0.4412912824781923])

vchuravy avatar May 12 '22 17:05 vchuravy

Note that what You might want to make your code type-stable.

julia> const S = sprand(N, N, 0.1)
10×10 SparseMatrixCSC{Float64, Int64} with 10 stored entries:
  ⋅   0.864191   ⋅         ⋅         ⋅        0.725971   ⋅         ⋅    ⋅         ⋅ 
  ⋅    ⋅         ⋅        0.454994   ⋅         ⋅         ⋅         ⋅   0.404086   ⋅ 
  ⋅    ⋅         ⋅         ⋅         ⋅        0.937271  0.321421   ⋅    ⋅         ⋅ 
  ⋅    ⋅         ⋅         ⋅        0.325655   ⋅         ⋅         ⋅    ⋅         ⋅ 
  ⋅   0.354443   ⋅         ⋅         ⋅         ⋅         ⋅         ⋅    ⋅         ⋅ 
  ⋅    ⋅        0.398672   ⋅         ⋅         ⋅         ⋅         ⋅    ⋅         ⋅ 
  ⋅    ⋅         ⋅         ⋅         ⋅         ⋅        0.14057    ⋅    ⋅         ⋅ 
  ⋅    ⋅         ⋅         ⋅         ⋅         ⋅         ⋅         ⋅    ⋅         ⋅ 
  ⋅    ⋅         ⋅         ⋅         ⋅         ⋅         ⋅         ⋅    ⋅         ⋅ 
  ⋅    ⋅         ⋅         ⋅         ⋅         ⋅         ⋅         ⋅    ⋅         ⋅ 

julia> const Q = _Q
10×10 SparseMatrixCSC{Float64, Int64} with 6 stored entries:
  ⋅    ⋅    ⋅    ⋅        0.450911   ⋅         ⋅    ⋅   0.0822345   ⋅ 
  ⋅    ⋅    ⋅    ⋅         ⋅         ⋅         ⋅    ⋅    ⋅          ⋅ 
  ⋅    ⋅    ⋅    ⋅         ⋅         ⋅         ⋅    ⋅    ⋅          ⋅ 
  ⋅    ⋅    ⋅    ⋅         ⋅         ⋅         ⋅    ⋅    ⋅          ⋅ 
  ⋅    ⋅    ⋅    ⋅         ⋅         ⋅         ⋅    ⋅    ⋅          ⋅ 
  ⋅    ⋅    ⋅    ⋅         ⋅         ⋅         ⋅    ⋅    ⋅          ⋅ 
  ⋅    ⋅    ⋅    ⋅         ⋅         ⋅         ⋅    ⋅    ⋅          ⋅ 
  ⋅    ⋅    ⋅    ⋅         ⋅         ⋅         ⋅    ⋅   0.712695    ⋅ 
  ⋅    ⋅    ⋅    ⋅         ⋅        0.670415   ⋅    ⋅    ⋅          ⋅ 
  ⋅    ⋅    ⋅   0.834334   ⋅        0.719653   ⋅    ⋅    ⋅          ⋅ 

julia> F(z) = [ # some complicated sparse amenable function
           S * z 
           Q * z.^2
       ]
F (generic function with 2 methods)

julia> jac = Enzyme.jacobian(Forward, F, z)
(([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8641913606126247, 0.0, 0.0, 0.0, 0.3544431260542895, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.39867161256143524, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.4549936171986305, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.14561249567530618], [0.0, 0.0, 0.0, 0.3256549706212486, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.8075284885092465, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7259714923753596, 0.0, 0.9372709878409331, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0423793826488412, 1.1189362287523177], [0.0, 0.0, 0.32142094179075853, 0.0, 0.0, 0.0, 0.14057035792574357, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.4040861036618695, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.05656868148254788, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.4902595563499479, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),)

vchuravy avatar May 12 '22 17:05 vchuravy

In particular the type-stable variant:

julia> @btime  Enzyme.jacobian(Forward, F, z)
  2.236 μs (59 allocations: 8.86 KiB)

Versus your variant:

julia> @btime  Enzyme.jacobian(Forward, F, z, Val(1))
  115.640 μs (855 allocations: 56.03 KiB)

vchuravy avatar May 12 '22 17:05 vchuravy

@vchuravy For the time for curiousity can you also try the type stable, but vector size 1 as well

wsmoses avatar May 12 '22 17:05 wsmoses

julia> @btime  Enzyme.jacobian(Forward, F, z, Val(1))
  3.336 μs (105 allocations: 14.78 KiB)

julia> @btime  Enzyme.jacobian(Forward, F, z)
  2.212 μs (59 allocations: 8.86 KiB)

vchuravy avatar May 12 '22 17:05 vchuravy

This worked! Thanks :) Although I ran into another issue, but I am not sure if they are related:

using Enzyme
using SparseArrays, LinearAlgebra
using Random 

N = 10
n = 2
rs = shuffle(1:N)[1:round(Int, N/n)]
cs = shuffle(1:N)[1:round(Int, N/n)]
_Q = sprand(N, N, 0.1)
_S(x) = sparse(rs, cs, x, N, N)

z = rand(N)
x = rand(round(Int, N/n))
_S(x)

F(x, z) = [ # some complicated sparse amenable function
    _S(x) * z 
    _Q * z.^2
]

F(x, z)

jac = Enzyme.jacobian(Forward, z -> F(x, z), z, Val(1))

causes this error:

julia> jac = Enzyme.jacobian(Forward, z -> F(x, z), z, Val(1))
ERROR: AssertionError: length(args) == length((collect(parameters(entry_f)))[1 + sret:end])
Stacktrace:
 [1] lower_convention(functy::Type, mod::LLVM.Module, entry_f::LLVM.Function, actualRetType::Type)
   @ Enzyme.Compiler C:\Users\stelmo\.julia\packages\Enzyme\Ctome\src\compiler.jl:3553
 [2] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(sparse), Tuple{Vector{Int64}, Vector{Int64}, Vector{Float64}, Int64, Int64}}}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, ctx::LLVM.Context, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
   @ Enzyme.Compiler C:\Users\stelmo\.julia\packages\Enzyme\Ctome\src\compiler.jl:3958
 [3] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(sparse), Tuple{Vector{Int64}, Vector{Int64}, Vector{Float64}, Int64, Int64}}})
   @ Enzyme.Compiler C:\Users\stelmo\.julia\packages\Enzyme\Ctome\src\compiler.jl:4397
 [4] cached_compilation(job::GPUCompiler.CompilerJob, key::UInt64, specid::UInt64)
   @ Enzyme.Compiler C:\Users\stelmo\.julia\packages\Enzyme\Ctome\src\compiler.jl:4435
 [5] #s501#108
   @ C:\Users\stelmo\.julia\packages\Enzyme\Ctome\src\compiler.jl:4495 [inlined]
 [6] var"#s501#108"(F::Any, DF::Any, A::Any, TT::Any, Mode::Any, ModifiedBetween::Any, width::Any, specid::Any, ::Any, f::Any, df::Any, #unused#::Type, tt::Any, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Any)
   @ Enzyme.Compiler .\none:0
 [7] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
   @ Core .\boot.jl:580
 [8] thunk
   @ C:\Users\stelmo\.julia\packages\Enzyme\Ctome\src\compiler.jl:4523 [inlined]
 [9] thunk(f::typeof(sparse), df::Nothing, ::Type{Duplicated}, tt::Type{Tuple{Const{Vector{Int64}}, Const{Vector{Int64}}, Const{Vector{Float64}}, Const{Int64}, Const{Int64}}}, ::Val{Enzyme.API.DEM_ForwardMode}, ::Val{0x0000000000000001})
   @ Enzyme.Compiler C:\Users\stelmo\.julia\packages\Enzyme\Ctome\src\compiler.jl:4516

Any suggestions would be appreciated :) I am on the main branch btw

stelmo avatar May 26 '22 12:05 stelmo

The first code now hits a:

julia: /buildworker/worker/package_linux64/build/src/llvm-late-gc-lowering.cpp:870: std::vector<int, std::allocator<int> > LateLowerGCFrame::NumberAllBase(State&, llvm::Value*): Assertion `Tracked.size() == BaseNumbers.size()' failed.

signal (6): Aborted
in expression starting at REPL[8]:1
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: 0x7f2c0a4d53f9)
__assert_fail at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
NumberAllBase at /buildworker/worker/package_linux64/build/src/llvm-late-gc-lowering.cpp:870
NumberBase at /buildworker/worker/package_linux64/build/src/llvm-late-gc-lowering.cpp:787
Number at /buildworker/worker/package_linux64/build/src/llvm-late-gc-lowering.cpp:804
LocalScan at /buildworker/worker/package_linux64/build/src/llvm-late-gc-lowering.cpp:1590
runOnFunction at /buildworker/worker/package_linux64/build/src/llvm-late-gc-lowering.cpp:2675
_ZN4llvm13FPPassManager13runOnFunctionERNS_8FunctionE at /mnt/Data/git/Enzyme.jl/julia-1.7.2/bin/../lib/julia/libLLVM-12jl.so (unknown line)

Your latest code hits:

julia> Enzyme.API.looseTypeAnalysis!(true)

julia> jac = Enzyme.jacobian(Forward, z -> F(x, z), z, Val(1))
┌ Warning: Returned rooting not fully handled, segfault likely
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/jVY4I/src/utils.jl:35
warning: binary operator is integer and constant:   %59 = add i64 %58, 1, !dbg !95
warning: binary operator is integer and constant:   %155 = add i64 %154, 1, !dbg !201
ERROR: UndefRefError: access to undefined reference
Stacktrace:
 [1] getindex(t::Tuple, i::Int64)
   @ Base ./tuple.jl:29

wsmoses avatar Sep 05 '22 16:09 wsmoses

First code remains a GC issue of vector mode bug (the same as above).

The second code is now a segfault:

julia> jac = Enzyme.jacobian(Forward, z -> F(x, z), z, Val(1))


┌ Warning: Returned rooting not fully handled, segfault likely
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/07qaN/src/utils.jl:35

signal (11): Segmentation fault
in expression starting at REPL[15]:1
unknown function (ip: 0x7ff5d9833410)
unknown function (ip: 0x7ff5d9833730)
Allocations: 71962723 (Pool: 71893824; Big: 68899); GC: 76
Segmentation fault (core dumped)

wsmoses avatar Oct 22 '22 20:10 wsmoses

Original code now fixed, upon use of 1.8.3 with the vector bug fixed in Julia proper, the second code now also appears to work, so closing.

wsmoses avatar Dec 15 '22 21:12 wsmoses

Thanks to both of you for fixing these issues. I really appreciate the effort you put into this package, it is great!!

stelmo avatar Dec 15 '22 21:12 stelmo