Forward over Reverse Crash

[michel2323]

The following code crashes when calling forward_over_reverse(...). The code and the Manifest.toml is at https://github.com/exanauts/EnzymeDiff.jl. Let me know what I should throw out further.

using Enzyme
using LinearAlgebra
Enzyme.API.printall!(true)
Enzyme.API.printtype!(true)

n = 2
function speelpenning(y, x)
    y .= x .* x
    return nothing
end

function reverse(y::VT, x::VT) where {VT}
    FT = eltype(x)
    rx = convert(VT, zeros(FT,n))
    ry = convert(VT, ones(FT,n))

    _x = Duplicated(x, rx)
    _y = Duplicated(y, ry)

    autodiff(speelpenning, _y, _x)
    return rx
end

function forward_over_reverse(y::VT, x::VT) where {VT}
    FT = eltype(x)
    dx = convert(VT, ones(FT,n)); rx = convert(VT, zeros(FT,n)); drx = convert(VT, zeros(FT,n))
    dy = convert(VT, zeros(FT,n)); ry = convert(VT, ones(FT,n)); dry = convert(VT, zeros(FT,n))

    function foo(y, x)
        autodiff_deferred(speelpenning, Const, y, x)
        return nothing
    end

    _x = Duplicated(Duplicated(x,rx), Duplicated(dx,drx))
    _y = Duplicated(Duplicated(y,ry), Duplicated(dy,dry))

    fwddiff(foo, _y, _x)
    return rx, drx
end

x = [i/(1.0+i) for i in 1:n]
y = zeros(n)
speelpenning(y,x)

g = reverse(y, x)

# Crashes
g1, g2 = forward_over_reverse(y,x)
# This test should pass
@show all(g .== g1)

Output:
log.tar.gz
Julia version and dependencies:

Julia Version 1.8.0-beta3
Commit 3e092a2521 (2022-03-29 15:42 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: 144 × Intel(R) Xeon(R) Gold 6140 CPU @ 2.30GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-13.0.1 (ORCJIT, skylake-avx512)
  Threads: 1 on 144 virtual cores

(EnzymeDiff.jl) pkg> st
Status `/scratch/mschanen/git/EnzymeDiff.jl/Project.toml`
  [7da242da] Enzyme v0.10.0-dev `/scratch/mschanen/julia_depot/dev/Enzyme`

[wsmoses]

@vchuravy this appears to be a bug in an Julia optimization pass:

Before:

L101.i:                                           ; preds = %L91.i
  %40 = call nonnull {}* @julia.pointer_from_objref({} addrspace(11)* %37) #13, !dbg !187
  %41 = bitcast {}* %40 to i8**, !dbg !187
  %42 = load i8*, i8** %41, align 8, !dbg !187, !tbaa !70, !nonnull !4
  %43 = call nonnull {}* @julia.pointer_from_objref({} addrspace(11)* %36) #13, !dbg !187
  %44 = bitcast {}* %43 to i8**, !dbg !187
  %45 = load i8*, i8** %44, align 8, !dbg !187, !tbaa !70, !nonnull !4
  %.not = icmp eq i8* %42, %45, !dbg !193
  store i8 2, i8* %_cache7, align 1, !dbg !198, !invariant.group !246
  br i1 %.not, label %L113.i, label %L118.i, !dbg !198

after

L101.i.i:                                         ; preds = %L91.i.i
  %45 = bitcast {} addrspace(10)* %0 to i8* addrspace(10)*, !dbg !117
  %46 = addrspacecast i8* addrspace(10)* %45 to i8**, !dbg !117
  %47 = load i8*, i8** %46, align 8, !dbg !117, !tbaa !124, !nonnull !4
  %48 = bitcast {} addrspace(10)* %2 to i8* addrspace(10)*, !dbg !117
  %49 = addrspacecast i8* addrspace(10)* %48 to i8**, !dbg !117
  %50 = load i8*, i8** %49, align 8, !dbg !117, !tbaa !124, !nonnull !4
  %.not.i = icmp eq i8* %47, %50, !dbg !126
  br i1 %.not.i, label %L113.i.i, label %L118.i.i, !dbg !131