Segfault with `Assertion 'isa<LoadInst>(orig) == isa<LoadInst>(val)' failed`

[avik-pal]

# using ArrayInterface, NNlib, Strided, Enzyme, FastBroadcast

using NNlib, Enzyme, Strided, FastBroadcast

@inline function __fast_broadcast!(f::F, x, args...) where {F}
    # if ArrayInterface.fast_scalar_indexing(x)
        if maximum(length, (x, args...)) > 200_000
            # @. x = f(x, args...)
            @strided x .= f.(x, args...)
        else
            @.. x = f(x, args...)
            # @.. x = f(x, args...)
        end
    # elseif f === ComposedFunction(sigmoid_fast, +) && length(args) == 1
    #     y = first(args)
    #     @. x = sigmoid_fast(x + y) # Has GPU Compilation Problems
    # else
    #     @. x = f(x, args...)
    # end
    return x
end

x = randn(Float32, 125)

__fast_broadcast!(gelu, x)

Enzyme.gradient(Enzyme.Reverse, x -> sum(__fast_broadcast!(gelu, x)), x)

I tried minimizing this further (see the commented parts) but nothing else seems to segfault

julia: /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:839: virtual llvm::Value* GradientUtils::unwrapM(llvm::Value*, llvm::IRBuilder<>&, const ValueToValueMapTy&, UnwrapMode, llvm::BasicBlock*, bool): Assertion `isa<LoadInst>(orig) == isa<LoadInst>(val)' failed.

[2846259] signal (6.-6): Aborted
in expression starting at /home/avik-pal/perf/enzyme/simple_enzyme_segfault.jl:66
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: 0x7f1e44461728)
__assert_fail at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unwrapM at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:839
unwrapM at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:1329
unwrapM at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:1065
getSubLimits at /workspace/srcdir/Enzyme/enzyme/Enzyme/CacheUtility.cpp:1197
createCacheForScope at /workspace/srcdir/Enzyme/enzyme/Enzyme/CacheUtility.cpp:792
ensureLookupCached at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:2368 [inlined]
ensureLookupCached at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:2354
lookupM at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:7305
lookup at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:2198
recursivelyHandleSubfunction at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:5172
visitCallInst at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:6601
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:4376
recursivelyHandleSubfunction at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:5736
visitCallInst at /workspace/srcdir/Enzyme/enzyme/Enzyme/AdjointGenerator.h:6601
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:4376
EnzymeCreatePrimalAndGradient at /workspace/srcdir/Enzyme/enzyme/Enzyme/CApi.cpp:614
EnzymeCreatePrimalAndGradient at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/api.jl:154
unknown function (ip: 0x7f1dad35858b)
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
enzyme! at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/compiler.jl:3109
unknown function (ip: 0x7f1dad353ea8)
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
#codegen#487 at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/compiler.jl:4964
codegen at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/compiler.jl:4391 [inlined]
_thunk at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5646
_thunk at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5646 [inlined]
cached_compilation at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5680 [inlined]
#532 at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5746
#JuliaContext#149 at /home/avik-pal/.julia/packages/GPUCompiler/kqxyC/src/driver.jl:52
unknown function (ip: 0x7f1dae157916)
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
JuliaContext at /home/avik-pal/.julia/packages/GPUCompiler/kqxyC/src/driver.jl:42
#s1926#531 at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/compiler.jl:5698 [inlined]
#s1926#531 at ./none:0
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
GeneratedFunctionStub at ./boot.jl:602
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_call_staged at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/method.c:540
ijl_code_for_staged at /cache/build/builder-amdci5-1/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_ext at ./compiler/typeinfer.jl:1049
typeinf_ext_toplevel at ./compiler/typeinfer.jl:1082
typeinf_ext_toplevel at ./compiler/typeinfer.jl:1078
jfptr_typeinf_ext_toplevel_35691.1 at /home/avik-pal/.julia/juliaup/julia-1.10.2+0.x64.linux.gnu/lib/julia/sys.so (unknown line)
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
jl_type_infer at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:394
jl_generate_fptr_impl at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/jitlayers.cpp:502
jl_compile_method_internal at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2480 [inlined]
jl_compile_method_internal at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2368
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2886 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
autodiff at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/Enzyme.jl:270 [inlined]
autodiff at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/Enzyme.jl:287 [inlined]
gradient at /home/avik-pal/.julia/packages/Enzyme/MIIMf/src/Enzyme.jl:938
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
do_call at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/interpreter.c:126
eval_value at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/interpreter.c:223
eval_stmt_value at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/interpreter.c:174 [inlined]
eval_body at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/interpreter.c:617
jl_interpret_toplevel_thunk at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/interpreter.c:775
jl_toplevel_eval_flex at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/toplevel.c:934
jl_toplevel_eval_flex at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/toplevel.c:877
ijl_toplevel_eval_in at /cache/build/builder-amdci5-1/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-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
jl_f__call_latest at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/builtins.c:812
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
do_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/builtins.c:768
#invokelatest#2 at ./essentials.jl:892
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
do_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/builtins.c:768
invokelatest at ./essentials.jl:889
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
do_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/builtins.c:768
#inlineeval#76 at /home/avik-pal/.vscode-server/extensions/julialang.language-julia-1.76.2/scripts/packages/VSCodeServer/src/eval.jl:263
inlineeval at /home/avik-pal/.vscode-server/extensions/julialang.language-julia-1.76.2/scripts/packages/VSCodeServer/src/eval.jl:260
#69 at /home/avik-pal/.vscode-server/extensions/julialang.language-julia-1.76.2/scripts/packages/VSCodeServer/src/eval.jl:181
withpath at /home/avik-pal/.vscode-server/extensions/julialang.language-julia-1.76.2/scripts/packages/VSCodeServer/src/repl.jl:274
#68 at /home/avik-pal/.vscode-server/extensions/julialang.language-julia-1.76.2/scripts/packages/VSCodeServer/src/eval.jl:179
hideprompt at /home/avik-pal/.vscode-server/extensions/julialang.language-julia-1.76.2/scripts/packages/VSCodeServer/src/repl.jl:38
#67 at /home/avik-pal/.vscode-server/extensions/julialang.language-julia-1.76.2/scripts/packages/VSCodeServer/src/eval.jl:150
unknown function (ip: 0x7f1e2d274fb2)
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
with_logstate at ./logging.jl:515
with_logger at ./logging.jl:627 [inlined]
#66 at /home/avik-pal/.vscode-server/extensions/julialang.language-julia-1.76.2/scripts/packages/VSCodeServer/src/eval.jl:255
unknown function (ip: 0x7f1e2d274e82)
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
jl_f__call_latest at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/builtins.c:812
#invokelatest#2 at ./essentials.jl:892 [inlined]
invokelatest at ./essentials.jl:889
unknown function (ip: 0x7f1e2d219565)
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
do_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/builtins.c:768
#64 at /home/avik-pal/.vscode-server/extensions/julialang.language-julia-1.76.2/scripts/packages/VSCodeServer/src/eval.jl:34
unknown function (ip: 0x7f1e2d2307b2)
_jl_invoke at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:2894 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/gf.c:3076
jl_apply at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/julia.h:1982 [inlined]
start_task at /cache/build/builder-amdci5-1/julialang/julia-release-1-dot-10/src/task.c:1238
Allocations: 50462671 (Pool: 50398715; Big: 63956); GC: 68
Aborted (core dumped)

Enzyme version is 0.12

[wsmoses]

Technically an assertion error not segfault, but in any case reducing:

using Enzyme, Strided, FastBroadcast

const gelu_2λ = √(8 / π)

function gelu(x)
    oftype(float(x), y)
end

@inline function __fast_broadcast!(x)
    # if ArrayInterface.fast_scalar_indexing(x)
        if length(x) > 200_000
            p6 = Strided.maybestrided(x)
            p10 = Strided.maybestrided(x)
            p11 = Base.broadcasted(gelu, p10)
            bc = Base.Broadcast.combine_styles(p6, p11)
            # inst = Base.Broadcast.instantiate(Base.Broadcast.Broadcasted(bc.style, bc.f, bc.args, axes(p6)))
            Base.materialize!(Base.Broadcast.combine_styles(p6, p11), p6, p11)
        end
    return x
end

x = randn(Float32, 125)

# Enzyme.gradient(Enzyme.Reverse, x -> sum(__fast_broadcast!(gelu, x)), x)

f(x) = sum(__fast_broadcast!(x))

dx = Enzyme.make_zero(x)
Enzyme.autodiff(Enzyme.Reverse, f, Active, Duplicated(x, dx))

[wsmoses]

Will be fixed by https://github.com/EnzymeAD/Enzyme/pull/1862 but will require a jll bump