Huge LLVM dump+LLVM error when mixed activity error expected

[danielwe]

This is a funny issue because this code is supposed to error, it's just not supposed to error like this.

The issue arises when a custom rule is defined for a function where some methods have MixedDuplicated return and are thus incompatible with custom rules. I'm expecting these cases to throw a stock EnzymeRuntimeException citing mixed activity. However, in some cases they dump thousands of lines of LLVM and throw an LLVMException instead, with the message "function failed verification (4)".

Here's an MWE. The issue seems to be triggered by the handle_infinities function, which I've copied almost verbatim from QuadGK.jl. See the full output at https://gist.github.com/danielwe/50fc3bee1f39620eef3d20339b6dad48

using Enzyme, LinearAlgebra

function handle_infinities(workfunc, f, s)
    s1, s2 = first(s), last(s)
    inf1, inf2 = isinf(s1), isinf(s2)
    if inf1 || inf2
        if inf1 && inf2 # x = t / (1 - t^2)
            return workfunc(
                function (t)
                    t2 = t * t
                    den = 1 / (1 - t2)
                    return f(oneunit(s1) * t * den) * (1 + t2) * den * den * oneunit(s1)
                end,
                map(s) do x
                    isinf(x) ? copysign(one(x), x) : 2x / (oneunit(x) + hypot(oneunit(x), 2x))
                end,
                t -> oneunit(s1) * t / (1 - t^2),
            )
        else
            (s0, si) = inf1 ? (s2, s1) : (s1, s2)
            if si < zero(si) # x = s0 - t / (1 - t)
                return workfunc(
                    function (t)
                        den = 1 / (1 - t)
                        return f(s0 - oneunit(s1) * t * den) * den * den * oneunit(s1)
                    end,
                    reverse(map(s) do x
                        1 / (1 + oneunit(x) / (s0 - x))
                    end),
                    t -> s0 - oneunit(s1) * t / (1 - t),
                )
            else # x = s0 + t / (1 - t)
                return workfunc(
                    function (t)
                        den = 1 / (1 - t)
                        return f(s0 + oneunit(s1) * t * den) * den * den * oneunit(s1)
                    end,
                    map(s) do x
                        1 / (1 + oneunit(x) / (x - s0))
                    end,
                    t -> s0 + oneunit(s1) * t / (1 - t),
                )
            end
        end
    end
    return workfunc(f, s, identity)
end

outer(f, xs...) = handle_infinities((f_, xs_, _) -> inner(f_, xs_), f, xs)

function inner(f::F, xs) where {F}  # remove type annotation => problem solved
    s = sum(f, xs)
    return (s, norm(s))
end

function EnzymeRules.augmented_primal(
    config::EnzymeRules.RevConfig, ::Const{typeof(inner)}, ::Type, f, xs
)
    true_primal = inner(f.val, xs.val)
    primal = EnzymeRules.needs_primal(config) ? true_primal : nothing
    shadow = if EnzymeRules.needs_shadow(config)
        if EnzymeRules.width(config) == 1
            make_zero(true_primal)
        else
            ntuple(_ -> make_zero(true_primal), Val(EnzymeRules.width(config)))
        end
    else
        nothing
    end
    return EnzymeRules.AugmentedReturn(primal, shadow, nothing)
end

function EnzymeRules.reverse(
    ::EnzymeRules.RevConfig, ::Const{typeof(inner)}, shadow::Active, tape, f, xs
)
    return ((f isa Active) ? f : nothing, (xs isa Active) ? xs : nothing)
end

# F_good(x) = outer(y -> [cos(x * y)], 0.0, 1.0)[1][1]
# autodiff(Reverse, F_good, Active(0.3))

F_bad(x) = outer(y -> [cos(y)], 0.0, x)[1][1]
autodiff(Reverse, F_bad, Active(0.3))

Output:

┌ Warning: Using fallback BLAS replacements for (["dasum_64_"]), performance may be degraded
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/Nxf8r/src/utils.jl:59
[...] # 1031 lines of LLVM

┌ Warning: Using fallback BLAS replacements for (["dasum_64_"]), performance may be degraded
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/Nxf8r/src/utils.jl:59
[...]. # 1031 lines of LLVM

┌ Warning: Using fallback BLAS replacements for (["dasum_64_"]), performance may be degraded
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/Nxf8r/src/utils.jl:59
[...] # 1031 lines of LLVM

ERROR: LoadError: LLVM error: function failed verification (4)
Stacktrace:
 [1] handle_error(reason::Cstring)
   @ LLVM ~/.julia/packages/LLVM/b3kFs/src/core/context.jl:194
in expression starting at /tmp/bad.jl:4

Some observations

• If the function arg type annotation in the inner definition is removed, the issue disappears and we get the expected error. This suggests to me that the issue is somehow related to specialization on function args and/or inlining, considering Julia's default non-specialization on function arguments.

• If the lines with F_good are uncommented and F_bad commented we get the expected error. The key difference seems to be that F_bad is differentiating wrt. a variable that handle_infinities branches on, while F_good does not. (However, the actual branch taken is always the trivial one, there are no infinities in these examples.)

  Here's the output when differentiating F_good:

┌ Warning: Using fallback BLAS replacements for (["dasum_64_"]), performance may be degraded
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/Nxf8r/src/utils.jl:59
ERROR: LoadError: Enzyme execution failed.
Enzyme: Return type Tuple{Vector{Float64}, Float64} has mixed internal activity types in evaluation of custom rule for MethodInstance for inner(::var"#23#24"{Float64}, ::Tuple{Float64, Float64}). See https://enzyme.mit.edu/julia/stable/faq/#Mixed-activity for more information

Stacktrace:
  [1] #19
    @ /tmp/mwe.jl:49 [inlined]
  [2] handle_infinities
    @ /tmp/mwe.jl:46
  [3] outer
    @ /tmp/mwe.jl:49 [inlined]
  [4] F_good
    @ /tmp/mwe.jl:79 [inlined]
  [5] diffejulia_F_good_111wrap
    @ /tmp/mwe.jl:0
  [6] macro expansion
    @ ~/.julia/packages/Enzyme/R6sE8/src/compiler.jl:5340 [inlined]
  [7] enzyme_call
    @ ~/.julia/packages/Enzyme/R6sE8/src/compiler.jl:4878 [inlined]
  [8] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/R6sE8/src/compiler.jl:4750 [inlined]
  [9] autodiff
    @ ~/.julia/packages/Enzyme/R6sE8/src/Enzyme.jl:503 [inlined]
 [10] autodiff
    @ ~/.julia/packages/Enzyme/R6sE8/src/Enzyme.jl:544 [inlined]
 [11] autodiff(mode::ReverseMode{false, false, FFIABI, false, false}, f::typeof(F_good), args::Active{Float64})
    @ Enzyme ~/.julia/packages/Enzyme/R6sE8/src/Enzyme.jl:516
 [12] top-level scope
    @ /tmp/mwe.jl:80
in expression starting at /tmp/mwe.jl:80