Enzyme.jl icon indicating copy to clipboard operation
Enzyme.jl copied to clipboard
verified publisher icon EnzymeAD

JET integration

Open aviatesk opened this issue 3 years ago • 2 comments

Just scratched -- I will work on this again tomorrow.

Better than segfault?

[ Info: Setting up Base development utilities ...
julia> using Enzyme

julia> f(cond, x) = cond ? x : 0
f (generic function with 1 method)

julia> g(cond, x) = f(cond,x)*x
g (generic function with 1 method)

julia> autodiff(g, Active, true, Active(1.0))
═════ 1 possible error found ═════
┌ @ REPL[3]:1 Main.f(cond, x)
│┌ @ REPL[2]:1 f(::Bool, ::Float64)
││ Union result (`Union{Float64, Int64}`) detected: f(::Bool, ::Float64)
│└─────────────
ERROR: "Enzyme can't differentiate this program: fix the Union result"
Stacktrace:
  [1] ci_cache_populate(interp::Enzyme.Compiler.EnzymeInterpeter, cache::GPUCompiler.CodeCache, mt::Core.MethodTable, mi::Core.MethodInstance, min_world::UInt64, max_world::Int32)
    @ Enzyme.Compiler ~/julia/packages/Enzyme/src/compiler.jl:2255
  [2] compile_method_instance(job::GPUCompiler.CompilerJob, method_instance::Core.MethodInstance; ctx::LLVM.Context)
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/jlgen.jl:335
  [3] macro expansion
    @ ~/.julia/packages/TimerOutputs/nDhDw/src/TimerOutput.jl:252 [inlined]
  [4] irgen(job::GPUCompiler.CompilerJob, method_instance::Core.MethodInstance; ctx::LLVM.Context)
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/irgen.jl:5
  [5] macro expansion
    @ ~/julia/packages/GPUCompiler/src/driver.jl:205 [inlined]
  [6] macro expansion
    @ ~/.julia/packages/TimerOutputs/nDhDw/src/TimerOutput.jl:252 [inlined]
  [7] macro expansion
    @ ~/julia/packages/GPUCompiler/src/driver.jl:204 [inlined]
  [8] emit_llvm(job::GPUCompiler.CompilerJob, method_instance::Any; libraries::Bool, deferred_codegen::Bool, optimize::Bool, only_entry::Bool, ctx::LLVM.Context)
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/utils.jl:64
  [9] codegen(output::Symbol, job::GPUCompiler.CompilerJob; libraries::Bool, deferred_codegen::Bool, optimize::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing, ctx::LLVM.Context)
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/driver.jl:108
 [10] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(g), Tuple{Bool, Float64}}}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, ctx::LLVM.Context, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/julia/packages/Enzyme/src/compiler.jl:3138
 [11] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(g), Tuple{Bool, Float64}}})
    @ Enzyme.Compiler ~/julia/packages/Enzyme/src/compiler.jl:3707
 [12] cached_compilation(cache::Dict{UInt64, Any}, job::GPUCompiler.CompilerJob, compiler::typeof(Enzyme.Compiler._thunk), linker::typeof(Enzyme.Compiler._link))
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/cache.jl:90
 [13] thunk(f::typeof(g), df::Nothing, ::Type{Active}, tt::Type{Tuple{Const{Bool}, Active{Float64}}}, ::Val{Enzyme.API.DEM_ReverseModeCombined})
    @ Enzyme.Compiler ~/julia/packages/Enzyme/src/compiler.jl:3760
 [14] autodiff(::typeof(g), ::Type{Active}, ::Bool, ::Vararg{Any})
    @ Enzyme ~/julia/packages/Enzyme/src/Enzyme.jl:197
 [15] top-level scope
    @ REPL[4]:1

aviatesk avatar Mar 18 '22 08:03 aviatesk

yes, that would be perfect!

mmorligh avatar Mar 18 '22 12:03 mmorligh

but since y is not active Enzyme can actually differentiate the code.

Now this PR handles this very simple case with a very simple activity analysis (one that just checks if a return value is used within a caller):

julia> using Enzyme

julia> f(cond, x) = cond ? x : 0
f (generic function with 1 method)

julia> g(cond, x) = f(cond,x)*x
g (generic function with 1 method)

julia> autodiff(g, Active, true, Active(1.0))
ERROR: Enzyme can't differentiate `g(::Bool, ::Float64)`. Fix the `Union` result:
═════ 1 possible error found ═════
┌ @ REPL[5]:1 Main.f(cond, x)
│┌ @ REPL[4]:1 f(::Bool, ::Float64)
││ active Union result detected: f(::Bool, ::Float64)::Union{Float64, Int64}
│└─────────────

Stacktrace:
  [1] ci_cache_populate(interp::Enzyme.Compiler.EnzymeInterpeter, cache::GPUCompiler.CodeCache, mt::Core.MethodTable, mi::Core.MethodInstance, min_world::UInt64, max_world::Int32)
    @ Enzyme.Compiler ~/julia/packages/Enzyme/src/compiler.jl:2266
  [2] compile_method_instance(job::GPUCompiler.CompilerJob, method_instance::Core.MethodInstance; ctx::LLVM.Context)
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/jlgen.jl:335
  [3] macro expansion
    @ ~/.julia/packages/TimerOutputs/nDhDw/src/TimerOutput.jl:252 [inlined]
  [4] irgen(job::GPUCompiler.CompilerJob, method_instance::Core.MethodInstance; ctx::LLVM.Context)
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/irgen.jl:5
  [5] macro expansion
    @ ~/julia/packages/GPUCompiler/src/driver.jl:205 [inlined]
  [6] macro expansion
    @ ~/.julia/packages/TimerOutputs/nDhDw/src/TimerOutput.jl:252 [inlined]
  [7] macro expansion
    @ ~/julia/packages/GPUCompiler/src/driver.jl:204 [inlined]
  [8] emit_llvm(job::GPUCompiler.CompilerJob, method_instance::Any; libraries::Bool, deferred_codegen::Bool, optimize::Bool, only_entry::Bool, ctx::LLVM.Context)
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/utils.jl:64
  [9] codegen(output::Symbol, job::GPUCompiler.CompilerJob; libraries::Bool, deferred_codegen::Bool, optimize::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing, ctx::LLVM.Context)
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/driver.jl:108
 [10] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(g), Tuple{Bool, Float64}}}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, ctx::LLVM.Context, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/julia/packages/Enzyme/src/compiler.jl:3161
 [11] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(g), Tuple{Bool, Float64}}})
    @ Enzyme.Compiler ~/julia/packages/Enzyme/src/compiler.jl:3730
 [12] cached_compilation(cache::Dict{UInt64, Any}, job::GPUCompiler.CompilerJob, compiler::typeof(Enzyme.Compiler._thunk), linker::typeof(Enzyme.Compiler._link))
    @ GPUCompiler ~/julia/packages/GPUCompiler/src/cache.jl:90
 [13] thunk(f::typeof(g), df::Nothing, ::Type{Active}, tt::Type{Tuple{Const{Bool}, Active{Float64}}}, ::Val{Enzyme.API.DEM_ReverseModeCombined})
    @ Enzyme.Compiler ~/julia/packages/Enzyme/src/compiler.jl:3783
 [14] autodiff(::typeof(g), ::Type{Active}, ::Bool, ::Vararg{Any})
    @ Enzyme ~/julia/packages/Enzyme/src/Enzyme.jl:197
 [15] top-level scope
    @ REPL[6]:1

julia> f(cond, x) = cond ? x : 0
f (generic function with 1 method)

julia> g(cond, x, y) = begin; f(cond,y); x end
g (generic function with 2 methods)

julia> autodiff(g, Active, Const(true), Active(1.0), Const(2.0))
(1.0,)

But as I said in the following comment, this activity analysis might be very incomplete in terms of covering Enzyme's auto-differentiability. https://github.com/EnzymeAD/Enzyme.jl/blob/8e3b84d9c103f66e75dd23ed2da812a160d9dc2b/src/compiler.jl#L2248-L2261

aviatesk avatar Mar 19 '22 04:03 aviatesk