Distributions.MixtureModel lead to LLVM error: function failed verification

[scheidan]

Currently, Enzyme can not differentiate through the logpdf of a Distributions.MixtureModel:

using Distributions
import Enzyme

dd = MixtureModel([Normal(0,1), Normal(1,2)], [0.1, 0.9])

Enzyme.gradient(Enzyme.Forward, x -> logpdf(dd, x[1]), [1.1])
Enzyme.gradient(Enzyme.Reverse, x -> logpdf(dd, x[1]), [1.1])

Forward mode error message:

ERROR: LLVM error: function failed verification (4)
Stacktrace:
  [1] handle_error(reason::Cstring)
    @ LLVM ~/.julia/packages/LLVM/wMjUU/src/core/context.jl:194
  [2] EnzymeCreateForwardDiff(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnValue::Bool, mode::Enzyme.API.CDerivativeMode, runtimeActivity::Bool, width::Int64, additionalArg::Ptr{…}, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…})
    @ Enzyme.API ~/.julia/packages/Enzyme/QVjE5/src/api.jl:334
  [3] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{…} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:1631
  [4] codegen(output::Symbol, job::GPUCompiler.CompilerJob{…}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:4436
  [5] codegen
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:3239 [inlined]
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5288
  [7] _thunk
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5288 [inlined]
  [8] cached_compilation
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5340 [inlined]
  [9] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5451
 [10] thunk_generator(world::UInt64, source::LineNumberNode, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5636
 [11] runtime_generic_fwd(activity::Type{…}, runtimeActivity::Val{…}, width::Val{…}, RT::Val{…}, f::typeof(logpdf), df::Nothing, primal_1::MixtureModel{…}, shadow_1_1::Nothing, primal_2::Float64, shadow_2_1::Float64)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/rules/jitrules.jl:288
 [12] #47
    @ ./REPL[242]:1 [inlined]
 [13] fwddiffejulia__47_94077wrap
    @ ./REPL[242]:0
 [14] macro expansion
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5218 [inlined]
 [15] enzyme_call
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:4764 [inlined]
 [16] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:4652 [inlined]
 [17] autodiff
    @ ~/.julia/packages/Enzyme/QVjE5/src/Enzyme.jl:654 [inlined]
 [18] autodiff
    @ ~/.julia/packages/Enzyme/QVjE5/src/Enzyme.jl:524 [inlined]
 [19] macro expansion
    @ ~/.julia/packages/Enzyme/QVjE5/src/sugar.jl:675 [inlined]
 [20] gradient(::EnzymeCore.ForwardMode{…}, ::var"#47#48", ::Vector{…}; chunk::Nothing, shadows::Tuple{…})
    @ Enzyme ~/.julia/packages/Enzyme/QVjE5/src/sugar.jl:577
 [21] gradient(::EnzymeCore.ForwardMode{false, EnzymeCore.FFIABI, false, false}, ::var"#47#48", ::Vector{Float64})
    @ Enzyme ~/.julia/packages/Enzyme/QVjE5/src/sugar.jl:577
 [22] top-level scope
    @ REPL[242]:1
Some type information was truncated. Use `show(err)` to see complete types.

Reverse mode errormessage:

ERROR: LLVM error: function failed verification (2)
Stacktrace:
  [1] handle_error(reason::Cstring)
    @ LLVM ~/.julia/packages/LLVM/wMjUU/src/core/context.jl:194
  [2] EnzymeCreateAugmentedPrimal(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnUsed::Bool, shadowReturnUsed::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, forceAnonymousTape::Bool, runtimeActivity::Bool, width::Int64, atomicAdd::Bool)
    @ Enzyme.API ~/.julia/packages/Enzyme/QVjE5/src/api.jl:404
  [3] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{…} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:1501
  [4] codegen(output::Symbol, job::GPUCompiler.CompilerJob{…}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:4436
  [5] codegen
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:3239 [inlined]
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5288
  [7] _thunk
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5288 [inlined]
  [8] cached_compilation
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5340 [inlined]
  [9] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5451
 [10] thunk_generator(world::UInt64, source::LineNumberNode, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5636
 [11] runtime_generic_augfwd(activity::Type{…}, runtimeActivity::Val{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(logpdf), df::Nothing, primal_1::MixtureModel{…}, shadow_1_1::Nothing, primal_2::Float64, shadow_2_1::Base.RefValue{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/QVjE5/src/rules/jitrules.jl:465
 [12] #49
    @ ./REPL[243]:1 [inlined]
 [13] augmented_julia__49_94372wrap
    @ ./REPL[243]:0
 [14] macro expansion
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:5218 [inlined]
 [15] enzyme_call
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:4764 [inlined]
 [16] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/QVjE5/src/compiler.jl:4700 [inlined]
 [17] autodiff
    @ ~/.julia/packages/Enzyme/QVjE5/src/Enzyme.jl:396 [inlined]
 [18] autodiff
    @ ~/.julia/packages/Enzyme/QVjE5/src/Enzyme.jl:524 [inlined]
 [19] macro expansion
    @ ~/.julia/packages/Enzyme/QVjE5/src/sugar.jl:319 [inlined]
 [20] gradient(::EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, ::var"#49#50", ::Vector{Float64})
    @ Enzyme ~/.julia/packages/Enzyme/QVjE5/src/sugar.jl:258
 [21] top-level scope
    @ REPL[243]:1
Some type information was truncated. Use `show(err)` to see complete types.

Versions

 [31c24e10] Distributions v0.25.114
 [7da242da] Enzyme v0.13.23

Julia Version 1.11.1
Commit 8f5b7ca12ad (2024-10-16 10:53 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 8 × Intel(R) Core(TM) i7-8650U CPU @ 1.90GHz
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, skylake)
Threads: 8 default, 0 interactive, 4 GC (on 8 virtual cores)

[wsmoses]

The issue at hand should be resolved by #2211, but your code subsequently hits an issue where LLVM emits an i7...... which is weird