Matrix pullback

[Roger-luo]

the following doesn't give up, but hangs forever on master branch

julia> moo(x) = fill(x, 10, 10)
moo (generic function with 1 method)

julia> Enzyme.pullback(moo, 2.1)(rand(10, 10))
^C^C^C^C^C^CWARNING: Force throwing a SIGINT
ERROR: InterruptException:
Stacktrace:
 [1] top-level scope
   @ REPL[21]:1

[vchuravy]

On #87 it doesn't hang but it:

ulia> Enzyme.pullback(moo, 2.1)(rand(10, 10))
could not merge test  {[-1]:Float@double}
julia: /workspace/srcdir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeTree.h:655: TypeTree TypeTree::KeepMinusOne() const: Assertion `0 && "could not merge"' failed.

signal (6): Aborted
in expression starting at REPL[4]:1
gsignal at /usr/bin/../lib/libc.so.6 (unknown line)
abort at /usr/bin/../lib/libc.so.6 (unknown line)
__assert_fail_base.cold at /usr/bin/../lib/libc.so.6 (unknown line)
__assert_fail at /usr/bin/../lib/libc.so.6 (unknown line)
KeepMinusOne at /workspace/srcdir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeTree.h:655 [inlined]
visitGetElementPtrInst at /workspace/srcdir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp:1137
run at /workspace/srcdir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp:895
analyzeFunction at /workspace/srcdir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp:4055
CreatePrimalAndGradient at /workspace/srcdir/Enzyme/enzyme/Enzyme/EnzymeLogic.cpp:2767
EnzymeCreatePrimalAndGradient at /workspace/srcdir/Enzyme/enzyme/Enzyme/CApi.cpp:343
EnzymeCreatePrimalAndGradient at /home/vchuravy/src/Enzyme/src/api.jl:94
enzyme! at /home/vchuravy/src/Enzyme/src/compiler.jl:1160
unknown function (ip: 0x7fd9561f2cd5)
_jl_invoke at /buildworker/worker/package_linux64/build/src/gf.c:2237 [inlined]
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2419
#codegen#34 at /home/vchuravy/src/Enzyme/src/compiler.jl:1421
codegen##kw at /home/vchuravy/src/Enzyme/src/compiler.jl:1291 [inlined]
_thunk at /home/vchuravy/src/Enzyme/src/compiler.jl:1788
unknown function (ip: 0x7fd956193ac1)
_jl_invoke at /buildworker/worker/package_linux64/build/src/gf.c:2237 [inlined]
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2419
cached_compilation at /home/vchuravy/.julia/packages/GPUCompiler/wHjqZ/src/cache.jl:89
thunk at /home/vchuravy/src/Enzyme/src/compiler.jl:1837
autodiff at /home/vchuravy/src/Enzyme/src/Enzyme.jl:167 [inlined]
autodiff at /home/vchuravy/src/Enzyme/src/Enzyme.jl:186 [inlined]
gradient at /home/vchuravy/src/Enzyme/src/Enzyme.jl:252 [inlined]
#14 at /home/vchuravy/src/Enzyme/src/Enzyme.jl:258
unknown function (ip: 0x7fd956183bf7)
_jl_invoke at /buildworker/worker/package_linux64/build/src/gf.c:2237 [inlined]
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2419
jl_apply at /buildworker/worker/package_linux64/build/src/julia.h:1703 [inlined]
do_call at /buildworker/worker/package_linux64/build/src/interpreter.c:115
eval_value at /buildworker/worker/package_linux64/build/src/interpreter.c:204
eval_stmt_value at /buildworker/worker/package_linux64/build/src/interpreter.c:155 [inlined]
eval_body at /buildworker/worker/package_linux64/build/src/interpreter.c:562
jl_interpret_toplevel_thunk at /buildworker/worker/package_linux64/build/src/interpreter.c:670
jl_toplevel_eval_flex at /buildworker/worker/package_linux64/build/src/toplevel.c:877
jl_toplevel_eval_flex at /buildworker/worker/package_linux64/build/src/toplevel.c:825
jl_toplevel_eval_in at /buildworker/worker/package_linux64/build/src/toplevel.c:929
eval at ./boot.jl:360 [inlined]
eval_user_input at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/REPL/src/REPL.jl:139
repl_backend_loop at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/REPL/src/REPL.jl:200
start_repl_backend at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/REPL/src/REPL.jl:185
#run_repl#42 at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/REPL/src/REPL.jl:317
run_repl at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/REPL/src/REPL.jl:305
_jl_invoke at /buildworker/worker/package_linux64/build/src/gf.c:2237 [inlined]
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2419
#874 at ./client.jl:387
jfptr_YY.874_41532.clone_1 at /usr/lib/julia/sys.so (unknown line)
_jl_invoke at /buildworker/worker/package_linux64/build/src/gf.c:2237 [inlined]
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2419
jl_apply at /buildworker/worker/package_linux64/build/src/julia.h:1703 [inlined]
jl_f__call_latest at /buildworker/worker/package_linux64/build/src/builtins.c:714
#invokelatest#2 at ./essentials.jl:708 [inlined]
invokelatest at ./essentials.jl:706 [inlined]
run_main_repl at ./client.jl:372
exec_options at ./client.jl:302
_start at ./client.jl:485
jfptr__start_34289.clone_1 at /usr/lib/julia/sys.so (unknown line)
_jl_invoke at /buildworker/worker/package_linux64/build/src/gf.c:2237 [inlined]
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2419
jl_apply at /buildworker/worker/package_linux64/build/src/julia.h:1703 [inlined]
true_main at /buildworker/worker/package_linux64/build/src/jlapi.c:560
repl_entrypoint at /buildworker/worker/package_linux64/build/src/jlapi.c:702
main at julia (unknown line)
__libc_start_main at /usr/bin/../lib/libc.so.6 (unknown line)
unknown function (ip: 0x4007d8)
Allocations: 20394470 (Pool: 20385625; Big: 8845); GC: 23

cc: @wsmoses

[wsmoses]

No longer hangs or fails on master

julia> moo(x) = fill(x, 10, 10)
moo (generic function with 1 method)

julia> Enzyme.pullback(moo, 2.1)(rand(10, 10))
([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],

Diving a bit further here, I think the return activity guessing is nonintuitive here presently:

julia> Enzyme.autodiff(moo, Active(2.1))
(0.0,)

julia> Enzyme.autodiff(moo, Duplicated, Active(2.1))
ERROR: Unhandled return type Duplicated{Matrix{Float64}}

[wsmoses]

Reopening to preserve need to allow matrix duplicated pullback

[vchuravy]

I think we can handle this now much more gracefully.

julia> L!(x, y) = x .* y
L! (generic function with 1 method)
julia> x = [1.0, 2.0]
julia> dx = zeros(2)


julia> autodiff(L!, Duplicated(x, dx), Const([2.0, 3.0]))
ERROR: Duplicated Returns not yet handled
Stacktrace:
 [1] autodiff
   @ ~/src/Enzyme/src/Enzyme.jl:192 [inlined]
 [2] autodiff(::typeof(L!), ::Duplicated{Vector{Float64}}, ::Const{Vector{Float64}})
   @ Enzyme ~/src/Enzyme/src/Enzyme.jl:224
 [3] top-level scope
   @ REPL[5]:1

Right now the canonical way is to change L! to:

julia> function L!(out, x, y)
         out .= x .* y
         return nothing
       end
L! (generic function with 2 methods)

julia> out = zeros(2)

julia> dout = ones(2)

julia> autodiff(L!, Duplicated(out, dout), Duplicated(x, dx), Const([2.0, 3.0]))
()

Instead we could directly support:

julia> L!(x, y) = x .* y
L! (generic function with 1 method)
julia> x = [1.0, 2.0]
julia> dx = zeros(2)
julia> out = zeros(2)
julia> dout = ones(2)

julia> autodiff(L!, Duplicated(out, dout), Duplicated(x, dx), Const([2.0, 3.0]))

Since we now have return activity annotations.

[wsmoses]

I do like having the traditional "pullback" as the activity annotation. However, allowing that would cause confusion with the "default" activity deduction, since its not clear that's a return activity for a 2-arg call vs a default deduction for a 3-arg call.