Illegal loop cache type crash in `cacheForReverse`

[penelopeysm]

Full error message: https://gist.github.com/penelopeysm/b106d30daf38a536913f69aa113532cc

This one also only crashes with reverse-mode on 1.11. It runs with either forward-mode and/or 1.10.

struct IL{T}
    idxs::T
end
struct VN{sym,T<:Union{typeof(identity),IL}}
    inner::T
    function VN{sym}(inner) where {sym}
        return new{sym,typeof(inner)}(inner)
    end
end
function Base.:(==)(x::VN{sym1}, y::VN{sym2}) where {sym1, sym2}
    return sym1 == sym2 && x.inner == y.inner
end
struct VI{Tmd,Tlogp}
    metadata::Tmd
    logp::Tlogp
end
struct TSVI{V}
    varinfo::V
end
struct MyMetadata{
    TIdcs<:Dict{<:VN,Int},
    TVal<:AbstractVector{<:Real},
}
    idcs::TIdcs
    vals::TVal
end
function g(vi, z)
    s = Vector{Float64}(undef, length(z))
    for i in eachindex(s)
        s[i] = vi.varinfo.metadata.s.vals[1]
    end
    return VI(vi.varinfo.metadata, 1.0)
end
params = [0.0]
function f(x)
    md_s = MyMetadata(
        Dict{VN,Int}(VN{:s}(IL((1,))) => 1),
        x,
    )
    vi = TSVI(VI((; s=md_s), 0.0))
    vi = g(vi, [1.5, 2.0])
    return vi.logp
end
f(params)

import Enzyme: Enzyme, Reverse, set_runtime_activity
Enzyme.gradient(set_runtime_activity(Reverse), f, params)

[penelopeysm]

Minimised a bit more now. I'm not sure where to go from here, the next thing I tried was to replace == with === in the Base.:(==) method (https://github.com/JuliaLang/julia/blob/9615af0f269df4d371b8010e9507ed5bae86103b/base/Base.jl#L212) but replacing either of the == causes it to work.

struct E{T}
    e::T
end
struct D{sym,T<:E}
    d::T
end
function Base.:(==)(d1::D{sym1}, d2::D{sym2}) where {sym1,sym2}
    # changing either == to === makes it no longer crash
    return sym1 == sym2 && d1.d == d2.d
end
function f(x)  # x is unused now
    Dict{D,Int}(D{:s,E{Int}}(E(1)) => 1)
    return 1.0
end
@show f([0.0])
import Enzyme: Enzyme, Reverse, set_runtime_activity
@show Enzyme.gradient(set_runtime_activity(Reverse), f, [0.0])

[penelopeysm]

I've inlined the relevant parts of Base.Dict (taken from 1.11.6) and trimmed it further, but I think I may have to leave the rest up to you :)

mutable struct MyDict{K,V}
    slots::Memory{UInt8}
    keys::Memory{K}
    vals::Memory{V}
    idxfloor::Int
    maxprobe::Int
    function MyDict{K,V}() where V where K
        slots = Memory{UInt8}(undef, 0)
        fill!(slots, 0x0)
        new(slots, Memory{K}(undef, 0), Memory{V}(undef, 0), 1, 0)
    end
end
function my_rehash!(h::MyDict{K,V}) where V where K
    newsz = 16
    h.idxfloor = 1
    slots = Memory{UInt8}(undef, newsz)
    fill!(slots, 0x0)
    h.slots = slots
    h.keys = Memory{K}(undef, newsz)
    h.vals = Memory{V}(undef, newsz)
    h.maxprobe = 0
    return h
end
function ht_keyindex2_shorthash!(h::MyDict{K,V}, key) where V where K
    sz = length(h.keys)
    if sz == 0
        my_rehash!(h)
        index, sh = Base.hashindex(key, length(h.keys))
        return -index, sh
    end
    index, sh = Base.hashindex(key, sz)
    keys = h.keys
    while true
        if h.slots[index] == 0x00
            return -index, sh
        end
        if h.slots[index] == sh
            k = keys[index]
            if key === k || isequal(key, k)
                return index, sh
            end
        end
        1 > h.maxprobe && break
    end
    return 0, sh
end
function my_setindex!(h::MyDict{K,V}, v, key::K) where V where K
    index, sh = ht_keyindex2_shorthash!(h, key)
    h.slots[-index] = sh
    h.keys[-index] = key
    h.vals[-index] = v
    return h
end

struct E{T}
    e::T
end
struct D{sym,T<:E}
    d::T
end
function Base.:(==)(d1::D{sym1}, d2::D{sym2}) where {sym1,sym2}
    # changing either == to === makes it no longer crash
    return sym1 == sym2 && d1.d == d2.d
end
function f(x)  # x is unused now
    d = MyDict{D,Int}()
    my_setindex!(d, 1, D{:s,E{Int}}(E(1)))
    return 1.0
end
@show f([0.0])
import Enzyme: Enzyme, Reverse, set_runtime_activity
@show Enzyme.gradient(set_runtime_activity(Reverse), f, [0.0])