help with R²→R implicit rule

[tpapp]

This is neither a bug (I think) nor a feature request, I need a bit of help implementing forward- and reverse rules for an implicitly calculated quantity. Specifically, let p be calculated as the solution to

$$p = H((1-p) \cdot A + p \cdot B)$$

where $H$ is the cdf of the $\mathrm{Normal}(0, 1)$ distribution. Note that everything is a scalar.

This defines the primal and a utility function to calculate the partials:

using Roots, Distributions, Enzyme
using .EnzymeRules

function solvep(A::T, B::T) where {T<:AbstractFloat}
    condition(p) = cdf(Normal(), p * A + (1 - p) * B) - p
    find_zero(condition, (0.0, 1.0))
end

function partialsp(A, B, p)
    H′ = pdf(Normal(), p * A + (1 - p) * B)
    C = H′ / (1 - H′ * (B - A))
    (∂p∂A = C * (1 - p), ∂p∂B = C * p)
end

and I thought this would work:

function EnzymeRules.forward(config::FwdConfig, func::Const{typeof(solvep)},
                             ::Type{<:Duplicated}, A::Duplicated, B::Duplicated)
    println("using custom rule")
    p = solvep(A.val, B.val)
    (; ∂p∂A, ∂p∂B) = partialsp(A.val, B.val, p)
    Duplicated(p, ∂p∂A * A.dval + ∂p∂B * B.dval)
end

but

autodiff(ForwardWithPrimal, solvep, Duplicated(1.0, 2.0), Duplicated(1.0, 1.0))

gives me a long machine instriction dump with a stacktrace

Stacktrace:
  [1] post_optimze!(mod::LLVM.Module, tm::LLVM.TargetMachine, machine::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/7Bzkr/src/compiler/optimize.jl:744
  [2] post_optimze!
    @ ~/.julia/packages/Enzyme/7Bzkr/src/compiler/optimize.jl:722 [inlined]
  [3] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/7Bzkr/src/compiler.jl:5722
  [4] _thunk
    @ ~/.julia/packages/Enzyme/7Bzkr/src/compiler.jl:5695 [inlined]
  [5] cached_compilation
    @ ~/.julia/packages/Enzyme/7Bzkr/src/compiler.jl:5749 [inlined]
  [6] 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, StrongZero::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/7Bzkr/src/compiler.jl:5863
  [7] thunk_generator(world::UInt64, source::Union{…}, 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, StrongZero::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, strongzero::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/7Bzkr/src/compiler.jl:6056
  [8] autodiff
    @ ~/.julia/packages/Enzyme/7Bzkr/src/Enzyme.jl:654 [inlined]
  [9] autodiff
    @ ~/.julia/packages/Enzyme/7Bzkr/src/Enzyme.jl:558 [inlined]
 [10] autodiff(::ForwardMode{true, FFIABI, false, false, false}, ::typeof(solvep), ::Duplicated{Float64}, ::Duplicated{Float64})
    @ Enzyme ~/.julia/packages/Enzyme/7Bzkr/src/Enzyme.jl:530
 [11] top-level scope
    @ REPL[13]:1
Some type information was truncated. Use `show(err)` to see complete types.

so probably I am making a mistake. I am using

  [31c24e10] Distributions v0.25.120
  [7da242da] Enzyme v0.13.67
  [57b37032] ImplicitDifferentiation v0.9.0
  [f2b01f46] Roots v2.2.8

and Julia 1.11.6.

[tpapp]

On Julia 1.12-rc1, I get

julia> autodiff(ForwardWithPrimal, solvep, Duplicated(1.0, 2.0), Duplicated(1.0, 1.0))
ERROR: UndefVarError: `_any` not defined in `Compiler`
Suggestion: check for spelling errors or missing imports.
Stacktrace:
  [1] getproperty
    @ ./Base_compiler.jl:47 [inlined]
  [2] isapplicable(f::Any, TT::Any; world::UInt64, method_table::Compiler.InternalMethodTable, caller::Nothing)
    @ EnzymeCore.EnzymeRules ~/.julia/packages/EnzymeCore/lmG5F/src/rules.jl:215
  [3] isapplicable
    @ ~/.julia/packages/EnzymeCore/lmG5F/src/rules.jl:197 [inlined]
  [4] #is_inactive_from_sig#8
    @ ~/.julia/packages/EnzymeCore/lmG5F/src/rules.jl:256 [inlined]
  [5] (::Enzyme.Compiler.Interpreter.var"#2#3"{Type{Tuple{typeof(Core.Typeof), Float64}}})(ret::Compiler.CallMeta, interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{Nothing}, sv::Compiler.InferenceState)
    @ Enzyme.Compiler.Interpreter ~/.julia/packages/Enzyme/7Bzkr/src/compiler/interpreter.jl:378
  [6] (::Compiler.var"#72#73"{Enzyme.Compiler.Interpreter.var"#2#3"{Type{Tuple{…}}}, Compiler.Future{Compiler.CallMeta}, Base.RefValue{Compiler.CallMeta}})(interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{Nothing}, sv::Compiler.InferenceState)
    @ Compiler ./../usr/share/julia/Compiler/src/inferencestate.jl:1183
  [7] doworkloop(interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{Nothing}, sv::Compiler.InferenceState)
    @ Compiler ./../usr/share/julia/Compiler/src/inferencestate.jl:1205
  [8] typeinf(interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{Nothing}, frame::Compiler.InferenceState)
    @ Compiler ./../usr/share/julia/Compiler/src/abstractinterpretation.jl:4434
  [9] typeinf_ext(interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{Nothing}, mi::Core.MethodInstance, source_mode::UInt8)
    @ Compiler ./../usr/share/julia/Compiler/src/typeinfer.jl:1259
 [10] typeinf_type
    @ ./../usr/share/julia/Compiler/src/typeinfer.jl:1281 [inlined]
 [11] return_type
    @ ~/.julia/packages/Enzyme/7Bzkr/src/typeutils/inference.jl:12 [inlined]
 [12] primal_return_type_world(mode::Mode, world::UInt64, mi::Core.MethodInstance)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/7Bzkr/src/typeutils/inference.jl:82
 [13] primal_return_type_generator(world::UInt64, source::Any, self::Any, mode::Type, ft::Type, tt::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/7Bzkr/src/typeutils/inference.jl:150
 [14] autodiff
    @ ~/.julia/packages/Enzyme/7Bzkr/src/Enzyme.jl:552 [inlined]
 [15] autodiff(::ForwardMode{true, FFIABI, false, false, false}, ::typeof(solvep), ::Duplicated{Float64}, ::Duplicated{Float64})
    @ Enzyme ~/.julia/packages/Enzyme/7Bzkr/src/Enzyme.jl:530
 [16] top-level scope
    @ REPL[9]:1
Some type information was truncated. Use `show(err)` to see complete types.

but I don't know if that is supported yet.

[tpapp]

attaching the full log (on Julia 1.11)

log.txt

[wsmoses]

cc @oscardssmith / @gbaraldi re 1.12

[wsmoses]

in 1.11 the error is

ERROR: LLVM error: broken gc calling conv fix
inlinable function call in a function with debug info must have a !dbg location
  %5 = call double @julia_nextfloat_16047(double noundef 0.000000e+00) #23
inlinable function call in a function with debug info must have a !dbg location
  %6 = call double @julia_nextfloat_16047(double noundef 0.000000e+00) #23

which is ....odd, will take a look [though this also reminds me of a case where the Julia compiler didn't emit this properly, so might be a bug in julia itself]

[wsmoses]

the bug comes from....llvm

#5  (anonymous namespace)::SimplifyCFGOpt::SpeculativelyExecuteBB (BI=BI@entry=0x628b497e3330, ThenBB=ThenBB@entry=0x628b48805930, TTI=..., this=<optimized out>)
    at /home/wmoses/git/julia11/deps/srccache/llvm-julia-16.0.6-3/llvm/lib/Transforms/Utils/SimplifyCFG.cpp:3037
3037	        I.setDebugLoc(DebugLoc());
(rr) l
3032	  // hoisting above.
3033	  for (auto &I : make_early_inc_range(*ThenBB)) {
3034	    if (!SpeculatedStoreValue || &I != SpeculatedStore) {
3035	      // Don't update the DILocation of dbg.assign intrinsics.
3036	      if (!isa<DbgAssignIntrinsic>(&I))
3037	        I.setDebugLoc(DebugLoc());
3038	    }
3039	    I.dropUndefImplyingAttrsAndUnknownMetadata();
3040	
3041	    // Drop ephemeral values.

[wsmoses]

we need to backport llvm/llvm-project#158448