AssertionError: useidx !== nothing

[willtebbutt]

There is a known issue with Core.Compiler.populate_use_def_map!, as discussed here: JuliaLang/julia#56193

This can result in problems like that encountered by @penelopeysm in this PR: TuringLang/Bijectors.jl#338

For posterity, here is the kind of error that you see if you encounter it:

  Got exception outside of a @test
  AssertionError: useidx !== nothing
  Stacktrace:
    [1] kill_def_use!(tpdum::Core.Compiler.TwoPhaseDefUseMap, def::Int64, use::Int64)
      @ Core.Compiler ./compiler/inferencestate.jl:147
    [2] kill_def_use!
      @ ./compiler/inferencestate.jl:158 [inlined]
    [3] reprocess_instruction!(interp::Core.Compiler.NativeInterpreter, inst::Core.Compiler.Instruction, idx::Int64, bb::Nothing, irsv::Core.Compiler.IRInterpretationState)
      @ Core.Compiler ./compiler/ssair/irinterp.jl:119
    [4] _ir_abstract_constant_propagation(interp::Core.Compiler.NativeInterpreter, irsv::Core.Compiler.IRInterpretationState; externally_refined::Nothing)
      @ Core.Compiler ./compiler/ssair/irinterp.jl:391
    [5] _ir_abstract_constant_propagation
      @ ./compiler/ssair/irinterp.jl:280 [inlined]
    [6] __infer_ir!
      @ ~/.julia/packages/Mooncake/LKJK9/src/interpreter/ir_utils.jl:117 [inlined]
    [7] optimise_ir!(ir::Core.Compiler.IRCode; show_ir::Bool, do_inline::Bool)
      @ Mooncake ~/.julia/packages/Mooncake/LKJK9/src/interpreter/ir_utils.jl:150
    [8] optimise_ir!
      @ ~/.julia/packages/Mooncake/LKJK9/src/interpreter/ir_utils.jl:139 [inlined]
    [9] build_rrule(interp::Mooncake.MooncakeInterpreter{Mooncake.DefaultCtx}, sig_or_mi::Type; debug_mode::Bool, silence_debug_messages::Bool)
      @ Mooncake ~/.julia/packages/Mooncake/LKJK9/src/interpreter/s2s_reverse_mode_ad.jl:925
   [10] build_rrule
      @ ~/.julia/packages/Mooncake/LKJK9/src/interpreter/s2s_reverse_mode_ad.jl:861 [inlined]
   [11] #build_rrule#210
      @ ~/.julia/packages/Mooncake/LKJK9/src/interpreter/s2s_reverse_mode_ad.jl:841 [inlined]
   [12] build_rrule
      @ ~/.julia/packages/Mooncake/LKJK9/src/interpreter/s2s_reverse_mode_ad.jl:839 [inlined]

As you can see from the stack trace, the problem occurs in _ir_abstract_constant_propagation, which is the function which performs type inference, performs constant propagation, and a few other things (I think).

A minimal reproducer on Julia 1.11.1 is:

using Pkg
Pkg.activate(; temp=true)
Pkg.@pkg_str"add [email protected] [email protected] [email protected] [email protected]"
using Bijectors, Distributions, LinearAlgebra, Mooncake

function f(θ)
    layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5])
    flow = transformed(MvNormal(zeros(2), I), layer)
    x = θ[6:7]
    return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)
end

rule = Mooncake.build_rrule(Tuple{typeof(f), Vector{Float64}})

We could do with a general strategy for working around this whenever it is encountered, or fixing it by patching the compiler internals until it is fixed more widely. Unfortunately, just fixing CC.populate_use_def_map! does not appear to resolve the problem, so I'm assuming that this is the surface of the problem. In slack, Keno mentioned that the code path on which this occurs is likely under-tested :(

edit: I've managed to find a more minimal reproducer:

using Pkg
Pkg.activate(; temp=true)
Pkg.@pkg_str"add [email protected]"
using Mooncake

function h(θ)
    d = [0.0, 0.0]
    x = θ[1:2]
    return d
end
Mooncake.build_rrule(Tuple{typeof(h), Vector{Float64}})