two errors on julia 1.12 for Flux layers #2998
Description
The following two examples give similar errors on julia 1.12
# Enzyme tries to differentiate through Julia's task CPU-time tracking
# (record_running_time! in task.jl), which uses atomic cmpxchg instructions.
# This is triggered indirectly via BLAS threading (wait is a yield point
# inside multithreaded BLAS calls).
#
# Error:
# WARNING: active variables passed by value to jl_new_task are not yet supported
# EnzymeInternalError: Illegal replace ficticious phi for:
# %_replacementA = phi { i64, i1 } ,
# of: %14 = cmpxchg ptr %6, i64 %12, i64 %13 monotonic monotonic, ...
#
# Simplified stacktrace:
# [1] modifyproperty! @ Base_compiler.jl:94
# [2] record_running_time! @ task.jl:1244
# [3] wait @ task.jl:1225
using Flux, Enzyme
using Statistics: mean
loss(f, x) = mean(f(x))
# --- Case 1: Bilinear ---
model = Flux.Bilinear((2, 2) => 3)
x = randn(Float32, 2, 1)
Enzyme.autodiff(
Enzyme.set_runtime_activity(ReverseWithPrimal), Const(loss), Active,
Duplicated(model, Enzyme.make_zero(model)),
Duplicated(x, Enzyme.make_zero(x)),
)
# --- Case 2: MultiHeadAttention (same error) ---
model2 = first ∘ MultiHeadAttention(16)
x2 = randn(Float32, 16, 20, 2)
Enzyme.autodiff(
Enzyme.set_runtime_activity(ReverseWithPrimal), Const(loss), Active,
Duplicated(model2, Enzyme.make_zero(model2)),
Duplicated(x2, Enzyme.make_zero(x2)),
)Stacktrace:
┌ Warning: active variables passed by value to jl_new_task are not yet supported
└ @ Enzyme.Compiler ~/.julia/packages/Enzyme/g8JAv/src/rules/parallelrules.jl:821
┌ Warning: active variables passed by value to jl_new_task are not yet supported
└ @ Enzyme.Compiler ~/.julia/packages/Enzyme/g8JAv/src/rules/parallelrules.jl:821
┌ Warning: active variables passed by value to jl_new_task are not yet supported
└ @ Enzyme.Compiler ~/.julia/packages/Enzyme/g8JAv/src/rules/parallelrules.jl:821
ERROR: LoadError: EnzymeInternalError: Enzyme compilation failed due to an internal error.
Please open an issue with the code to reproduce and full error log on github.com/EnzymeAD/Enzyme.jl
To toggle more information for debugging (needed for bug reports), set Enzyme.Compiler.VERBOSE_ERRORS[] = true (default false)
Illegal replace ficticious phi for: %_replacementA = phi { i64, i1 } , !dbg !188 of %14 = cmpxchg ptr %6, i64 %12, i64 %13 monotonic monotonic, align 8, !dbg !178, !tbaa !166, !alias.scope !169, !noalias !170
Stacktrace:
[1] modifyproperty!
@ ./Base_compiler.jl:94
[2] record_running_time!
@ ./task.jl:1244
[3] wait
@ ./task.jl:1225
Stacktrace:
[1] julia_error(msg::String, val::Ptr{…}, errtype::Enzyme.API.ErrorType, data::Ptr{…}, data2::Ptr{…}, B::Ptr{…})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/g8JAv/src/errors.jl:1281
[2] julia_error(cstr::Cstring, val::Ptr{…}, errtype::Enzyme.API.ErrorType, data::Ptr{…}, data2::Ptr{…}, B::Ptr{…})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/g8JAv/src/errors.jl:997
[3] EnzymeCreatePrimalAndGradient(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnValue::Bool, dretUsed::Bool, mode::Enzyme.API.CDerivativeMode, runtimeActivity::Bool, strongZero::Bool, width::Int64, additionalArg::Ptr{…}, forceAnonymousTape::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, augmented::Ptr{…}, atomicAdd::Bool)
@ Enzyme.API ~/.julia/packages/Enzyme/g8JAv/src/api.jl:270
[4] macro expansion
@ ~/.julia/packages/Enzyme/g8JAv/src/compiler.jl:2772 [inlined]
[5] macro expansion
@ ~/.julia/packages/LLVM/fEIbx/src/base.jl:97 [inlined]
[6] enzyme!(job::GPUCompiler.CompilerJob{…}, interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{…}, 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{…}, removedRoots::Set{…})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/g8JAv/src/compiler.jl:2645
[7] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob{…})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/g8JAv/src/compiler.jl:5145
[8] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
@ GPUCompiler ~/.julia/packages/GPUCompiler/OCZFZ/src/driver.jl:67
[9] compile
@ ~/.julia/packages/GPUCompiler/OCZFZ/src/driver.jl:55 [inlined]
[10] _thunk(job::GPUCompiler.CompilerJob{…}, postopt::Bool)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/g8JAv/src/compiler.jl:6745
[11] _thunk
@ ~/.julia/packages/Enzyme/g8JAv/src/compiler.jl:6743 [inlined]
[12] cached_compilation
@ ~/.julia/packages/Enzyme/g8JAv/src/compiler.jl:6801 [inlined]
[13] 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/g8JAv/src/compiler.jl:6917
[14] 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/g8JAv/src/compiler.jl:7061
[15] autodiff(::ReverseMode{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Duplicated{…})
@ Enzyme ~/.julia/packages/Enzyme/g8JAv/src/Enzyme.jl:509
[16] top-level scope
@ ~/.julia/dev/Flux/mwe_enzyme_internal_error.jl:29
[17] include(mapexpr::Function, mod::Module, _path::String)
@ Base ./Base.jl:307
[18] top-level scope
@ REPL[2]:1
in expression starting at /Users/carlo/.julia/dev/Flux/mwe_enzyme_internal_error.jl:29
Some type information was truncated. Use `show(err)` to see complete types.
``