StaticArray block vec' * mat segfault on 0.5.11

[marius311]

The following throws an error on Zygote 0.5.10, but segfaults on Zygote 0.5.11 & beyond (with Julia 1.5.3 / ChainRules v0.7.33)

julia> using Zygote, StaticArrays, LinearAlgebra

julia> f = Diagonal([1,2])
2×2 Diagonal{Int64,Array{Int64,1}}:
 1  ⋅
 ⋅  2

julia> gradient(f -> sum(sum(@SMatrix[f f] * @SMatrix[f f; f f])), f)
Internal error: encountered unexpected error in runtime:
BoundsError(a=Array{Any, (2,)}[
  Core.Compiler.VarState(typ=Zygote.Pullback{Tuple{typeof(StaticArrays.mul_unrolled), StaticArrays.Size{(1, 2)}, StaticArrays.Size{(2, 2)}, StaticArrays.SArray{Tuple{1, 2}, LinearAlgebra.Diagonal{Int64, Array{Int64, 1}}, 2, 2}, StaticArrays.SArray{Tuple{2, 2}, LinearAlgebra.Diagonal{Int64, Array{Int64, 1}}, 2, 4}}, Any}, undef=false),
  Core.Compiler.VarState(typ=FillArrays.Fill{FillArrays.Fill{Int64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, undef=false)], i=(3,))

From the v0.5.10...v0.5.11 diff, it seems likely its some difference in the ChainRules adjoint which was switched to there.

An edge-case no doubt and it wasn't even working before (but at least no segfault), but thought I'd point out in case its a sign of a more general thing that could be worth fixing.

Full stack trace

jl_bounds_error_ints at /home/marius/src/julia-1.5/src/rtutils.c:183
setindex! at ./essentials.jl:452 [inlined]
stupdate! at ./compiler/typelattice.jl:245
typeinf_local at ./compiler/abstractinterpretation.jl:1299
typeinf_nocycle at ./compiler/abstractinterpretation.jl:1326
typeinf at ./compiler/typeinfer.jl:12
typeinf_edge at ./compiler/typeinfer.jl:484
abstract_call_method at ./compiler/abstractinterpretation.jl:419
abstract_call_gf_by_type at ./compiler/abstractinterpretation.jl:111
abstract_call at ./compiler/abstractinterpretation.jl:924
abstract_call at ./compiler/abstractinterpretation.jl:911
abstract_eval at ./compiler/abstractinterpretation.jl:1005
typeinf_local at ./compiler/abstractinterpretation.jl:1270
typeinf_nocycle at ./compiler/abstractinterpretation.jl:1326
typeinf at ./compiler/typeinfer.jl:12
typeinf_edge at ./compiler/typeinfer.jl:484
abstract_call_method at ./compiler/abstractinterpretation.jl:419
abstract_call_gf_by_type at ./compiler/abstractinterpretation.jl:111
abstract_call at ./compiler/abstractinterpretation.jl:924
abstract_call at ./compiler/abstractinterpretation.jl:911
abstract_eval at ./compiler/abstractinterpretation.jl:1005
typeinf_local at ./compiler/abstractinterpretation.jl:1270
typeinf_nocycle at ./compiler/abstractinterpretation.jl:1326
typeinf at ./compiler/typeinfer.jl:12
typeinf_edge at ./compiler/typeinfer.jl:484
abstract_call_method at ./compiler/abstractinterpretation.jl:419
abstract_call_gf_by_type at ./compiler/abstractinterpretation.jl:111
abstract_call at ./compiler/abstractinterpretation.jl:924
abstract_call at ./compiler/abstractinterpretation.jl:911
abstract_eval at ./compiler/abstractinterpretation.jl:1005
typeinf_local at ./compiler/abstractinterpretation.jl:1270
typeinf_nocycle at ./compiler/abstractinterpretation.jl:1326
typeinf at ./compiler/typeinfer.jl:12
typeinf_edge at ./compiler/typeinfer.jl:484
abstract_call_method at ./compiler/abstractinterpretation.jl:419
abstract_call_gf_by_type at ./compiler/abstractinterpretation.jl:111
abstract_call at ./compiler/abstractinterpretation.jl:924
abstract_call at ./compiler/abstractinterpretation.jl:911
abstract_eval at ./compiler/abstractinterpretation.jl:1005
typeinf_local at ./compiler/abstractinterpretation.jl:1270
typeinf_nocycle at ./compiler/abstractinterpretation.jl:1326
typeinf at ./compiler/typeinfer.jl:12
typeinf_ext at ./compiler/typeinfer.jl:570
typeinf_ext at ./compiler/typeinfer.jl:601
jfptr_typeinf_ext_23397 at /home/marius/src/julia-1.5/usr/lib/julia/sys.so (unknown line)
jl_apply at /home/marius/src/julia-1.5/src/julia.h:1690 [inlined]
jl_type_infer at /home/marius/src/julia-1.5/src/gf.c:296
jl_generate_fptr at /home/marius/src/julia-1.5/src/jitlayers.cpp:290
jl_compile_method_internal at /home/marius/src/julia-1.5/src/gf.c:1964
jl_compile_method_internal at /home/marius/src/julia-1.5/src/gf.c:1931 [inlined]
_jl_invoke at /home/marius/src/julia-1.5/src/gf.c:2224 [inlined]
jl_apply_generic at /home/marius/src/julia-1.5/src/gf.c:2398
gradient at /home/marius/.julia/packages/Zygote/xBjHw/src/compiler/interface.jl:54
unknown function (ip: 0x7f7594a20e15)
jl_apply at /home/marius/src/julia-1.5/src/julia.h:1690 [inlined]
do_call at /home/marius/src/julia-1.5/src/interpreter.c:117
eval_value at /home/marius/src/julia-1.5/src/interpreter.c:206
eval_stmt_value at /home/marius/src/julia-1.5/src/interpreter.c:157 [inlined]
eval_body at /home/marius/src/julia-1.5/src/interpreter.c:548
jl_interpret_toplevel_thunk at /home/marius/src/julia-1.5/src/interpreter.c:660
top-level scope at REPL[3]:1
jl_toplevel_eval_flex at /home/marius/src/julia-1.5/src/toplevel.c:840
jl_toplevel_eval_flex at /home/marius/src/julia-1.5/src/toplevel.c:790
jl_toplevel_eval_in at /home/marius/src/julia-1.5/src/toplevel.c:883
eval at ./boot.jl:331
eval_user_input at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:134
repl_backend_loop at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:195
start_repl_backend at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:180
#run_repl#37 at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:292
run_repl at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:288
#806 at ./client.jl:399
jfptr_YY.806_48609 at /home/marius/src/julia-1.5/usr/lib/julia/sys.so (unknown line)
jl_apply at /home/marius/src/julia-1.5/src/julia.h:1690 [inlined]
do_apply at /home/marius/src/julia-1.5/src/builtins.c:655
jl_f__apply_latest at /home/marius/src/julia-1.5/src/builtins.c:705
#invokelatest#1 at ./essentials.jl:710 [inlined]
invokelatest at ./essentials.jl:709 [inlined]
run_main_repl at ./client.jl:383
exec_options at ./client.jl:313
_start at ./client.jl:506
jfptr__start_56184 at /home/marius/src/julia-1.5/usr/lib/julia/sys.so (unknown line)
unknown function (ip: 0x556eba3fe9d9)
unknown function (ip: 0x556eba3fe5a6)
__libc_start_main at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x556eba3fe64d)

signal (11): Segmentation fault
in expression starting at REPL[3]:1
_ZN4llvm3Use3setEPNS_5ValueE at /home/marius/src/julia-1.5/usr/bin/../lib/libLLVM-9jl.so (unknown line)
_ZN4llvm9StoreInstC1EPNS_5ValueES2_bjNS_14AtomicOrderingEhPNS_11InstructionE at /home/marius/src/julia-1.5/usr/bin/../lib/libLLVM-9jl.so (unknown line)
_ZN4llvm9StoreInstC2EPNS_5ValueES2_bjPNS_11InstructionE at /home/marius/src/julia-1.5/usr/bin/../lib/libLLVM-9jl.so (unknown line)
CreateStore at /home/marius/src/julia-1.5/usr/include/llvm/IR/IRBuilder.h:1606
emit_varinfo_assign at /home/marius/src/julia-1.5/src/codegen.cpp:3402
emit_assignment at /home/marius/src/julia-1.5/src/codegen.cpp:3441 [inlined]
emit_expr at /home/marius/src/julia-1.5/src/codegen.cpp:3654
emit_ssaval_assign at /home/marius/src/julia-1.5/src/codegen.cpp:3339
emit_stmtpos at /home/marius/src/julia-1.5/src/codegen.cpp:3503
emit_function at /home/marius/src/julia-1.5/src/codegen.cpp:6085
jl_emit_code at /home/marius/src/julia-1.5/src/codegen.cpp:6431
jl_emit_codeinst at /home/marius/src/julia-1.5/src/codegen.cpp:6465
_jl_compile_codeinst at /home/marius/src/julia-1.5/src/jitlayers.cpp:97
jl_generate_fptr_for_unspecialized at /home/marius/src/julia-1.5/src/jitlayers.cpp:339
jl_compile_method_internal at /home/marius/src/julia-1.5/src/gf.c:1970
jl_compile_method_internal at /home/marius/src/julia-1.5/src/gf.c:1931 [inlined]
_jl_invoke at /home/marius/src/julia-1.5/src/gf.c:2224 [inlined]
jl_apply_generic at /home/marius/src/julia-1.5/src/gf.c:2398
macro expansion at /home/marius/.julia/packages/StaticArrays/LJQEe/src/matrix_multiply.jl:72 [inlined]
_mul at /home/marius/.julia/packages/StaticArrays/LJQEe/src/matrix_multiply.jl:67 [inlined]
Pullback at /home/marius/.julia/packages/Zygote/xBjHw/src/compiler/interface2.jl:0
* at /home/marius/.julia/packages/StaticArrays/LJQEe/src/matrix_multiply.jl:9 [inlined]
Pullback at /home/marius/.julia/packages/Zygote/xBjHw/src/compiler/interface2.jl:0
#1 at ./REPL[3]:1 [inlined]
Pullback at /home/marius/.julia/packages/Zygote/xBjHw/src/compiler/interface2.jl:0
#41 at /home/marius/.julia/packages/Zygote/xBjHw/src/compiler/interface.jl:45
gradient at /home/marius/.julia/packages/Zygote/xBjHw/src/compiler/interface.jl:54
unknown function (ip: 0x7f7594a20e15)
jl_apply at /home/marius/src/julia-1.5/src/julia.h:1690 [inlined]
do_call at /home/marius/src/julia-1.5/src/interpreter.c:117
eval_value at /home/marius/src/julia-1.5/src/interpreter.c:206
eval_stmt_value at /home/marius/src/julia-1.5/src/interpreter.c:157 [inlined]
eval_body at /home/marius/src/julia-1.5/src/interpreter.c:548
jl_interpret_toplevel_thunk at /home/marius/src/julia-1.5/src/interpreter.c:660
jl_toplevel_eval_flex at /home/marius/src/julia-1.5/src/toplevel.c:840
jl_toplevel_eval_flex at /home/marius/src/julia-1.5/src/toplevel.c:790
jl_toplevel_eval_in at /home/marius/src/julia-1.5/src/toplevel.c:883
eval at ./boot.jl:331
eval_user_input at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:134
repl_backend_loop at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:195
start_repl_backend at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:180
#run_repl#37 at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:292
run_repl at /home/marius/src/julia-1.5/usr/share/julia/stdlib/v1.5/REPL/src/REPL.jl:288
#806 at ./client.jl:399
jfptr_YY.806_48609 at /home/marius/src/julia-1.5/usr/lib/julia/sys.so (unknown line)
jl_apply at /home/marius/src/julia-1.5/src/julia.h:1690 [inlined]
do_apply at /home/marius/src/julia-1.5/src/builtins.c:655
jl_f__apply_latest at /home/marius/src/julia-1.5/src/builtins.c:705
#invokelatest#1 at ./essentials.jl:710 [inlined]
invokelatest at ./essentials.jl:709 [inlined]
run_main_repl at ./client.jl:383
exec_options at ./client.jl:313
_start at ./client.jl:506
jfptr__start_56184 at /home/marius/src/julia-1.5/usr/lib/julia/sys.so (unknown line)
unknown function (ip: 0x556eba3fe9d9)
unknown function (ip: 0x556eba3fe5a6)
__libc_start_main at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x556eba3fe64d)
Allocations: 79488667 (Pool: 79474354; Big: 14313); GC: 62
Segmentation fault (core dumped)

[simeonschaub]

This is because the rules in ChainRules.jl restrict the element types of the multiplicands to just complex and real numbers, whereas the previous rule in Zygote caught all eltypes. The segfault is from Zygote then trying to recurse into StaticArray's * implementation and this seems to somehow fail, probably because of its heavy use of generated function. (That might actually be a more serious underlying issue in Zygote, I think generated functions should generally be supported) @oxinabox Should we perhaps just loosen the signature for this rule in ChainRules, or would that break other stuff?

[DhairyaLGandhi]

I don't see why generated functions would be a massive issue, but worth fixing in ChainRules. Perhaps we might be able to have a stub in Zygote too for specific type signatures, but that seems a bit too narrow at first glance.

[oxinabox]

I am not sure that this is a ChainRules problem, so much as a problem with the previous invalid rule in Zygote causing a different error to occur before the segfault could.
At least re: the segfault.

But here is a similar example that does show that a rule is now missing.
Zygorte 0.5.11:

julia> gradient(f -> sum(sum(reshape([f,f],(1,2)) * reshape([f,f,f,f], (2,2)))), f)
ERROR: Mutating arrays is not supported

Zygote 0.5.10

julia> gradient(f -> sum(sum(reshape([f,f],(1,2)) * reshape([f,f,f,f], (2,2)))), f)
([8 12; 12 16],)

I think we should accept real complex or AbstractArrays, thought it it always a bit problematic to be too wide.
Still better than erroring.
Better to be too wide in ChainRules where we can later fix it than to have extra rules special for Zygote.
I think.