Enzyme compilation failed on IPU

[giordano]

I'm playing with Enzyme.jl on an IPU. Autodiff of a few functions work out-of-the-box, which is awesome, but differentiation of log chokes:

f(x) = log(x)
f′(x) = first(first(autodiff_deferred(Reverse, f, Active(x))))

ERROR: LoadError: Enzyme compilation failed.
Current scope: 

Dump of LLVM IR (too long to post here).

It's not impossible I have something wrong in my GPUCompiler runtime.

[wsmoses]

This locally works, do you have the full config (it looks like this isn't base.log)

[giordano]

do you have the full config

What do you need exactly?

(it looks like this isn't base.log)

Uhm, why do you think so? I don't think I'm overriding log on this device, not intentionally at least, as far as I can tell log(::Float32) should be the same method as Base.log(::Float32), for example the definition of @julia__log_2972 in the above IR looks to me the same as code_llvm(Base.Math._log, (Float32, Val{:ℯ}, Symbol)).

I forgot to mention I'm on 1.9.0-rc1, if that makes a difference.

[wsmoses]

We explicitly mark Base.log as being log-like, if you call another log, we will not do that.

Hence wondering what method you call.

[wsmoses]

See

Enzyme.jl/src/compiler.jl

Line 57 in 3017330

Base.log => (:log, 1),

and

Enzyme.jl/src/compiler/interpreter.jl

Line 105 in 3017330

ft === typeof(Base.log) ||

If these wouldn't trigger on your function, that is the issue.

[giordano]

Uhm, I get the same error with

f(x) = Base.log(x)
f′(x) = first(first(autodiff_deferred(Reverse, f, Active(x))))

but not with

f′(x) = first(first(autodiff_deferred(Reverse, Base.log, Active(x))))

nor

f′(x) = first(first(autodiff_deferred(Reverse, log, Active(x))))

nor

const f = log
f′(x) = first(first(autodiff_deferred(Reverse, f, Active(x))))

😕

[wsmoses]

Maybe there's have been some overridden method in your IPU stuff.

Can you send me a file that I can run locally that errors?

[wsmoses]

try this PR maybe? #768

[giordano]

No, that doesn't seem to help unfortunately.

I'm very confused by the fact

f′(x) = first(first(autodiff_deferred(Reverse, log, Active(x))))

works though

[giordano]

I just realised I didn't share the stacktrace:

Stacktrace:
 [1] &
   @ ./int.jl:347
 [2] _log
   @ ./special/log.jl:324

Stacktrace:
  [1] julia_error(cstr::Cstring, val::Ptr{LLVM.API.LLVMOpaqueValue}, errtype::Enzyme.API.ErrorType, data::Ptr{Nothing})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/zewdu/src/compiler.jl:5314
  [2] EnzymeCreatePrimalAndGradient(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{Enzyme.API.CDIFFE_TYPE}, TA::Enzyme.TypeAnalysis, returnValue::Bool, dretUsed::Bool, mode::Enzyme.API.CDerivativeMode, width::Int64, additionalArg::Ptr{Nothing}, forceAnonymousTape::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{Bool}, augmented::Ptr{Nothing}, atomicAdd::Bool)
    @ Enzyme.API ~/.julia/packages/Enzyme/zewdu/src/api.jl:124
  [3] enzyme!(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::Tuple{Bool, Bool}, returnPrimal::Bool, jlrules::Vector{String}, expectedTapeType::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/zewdu/src/compiler.jl:6923
  [4] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, ctx::LLVM.ThreadSafeContext, strip::Bool, validate::Bool, only_entry::Bool, parent_job::GPUCompiler.CompilerJob{GPUCompiler.NativeCompilerTarget, IPUToolkit.IPUCompiler.IPUCompilerParams})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/zewdu/src/compiler.jl:8173
  [5] codegen
    @ ~/.julia/packages/Enzyme/zewdu/src/compiler.jl:7799 [inlined]

This is definitely entering inside Base.Math._log, so is that the problem? Maybe with

f(x) = log(x)
f′(x) = first(first(autodiff_deferred(Reverse, f, Active(x))))

the call to Base.log is inlined and it goes directly into Base.Math._log?

[wsmoses]

@vchuravy does this interpreter not abide by Enzyme's noinlining for custom functions?

[wsmoses]

Yeah I think this is equivalent to #683

[wsmoses]

Closing as equivalent