JET integration #261

aviatesk · 2022-03-18T08:18:33Z

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

MathieuMorlighem · 2022-03-18T12:08:25Z

yes, that would be perfect!

vchuravy

This is great. Would be good to discuss the cost :)

On the other hand this is also a bit sensitive (but I don't know) if we can actually get it to match 100%.

julia> using Enzyme

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 1 method)

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

but since y is not active Enzyme can actually differentiate the code. On this branch:

┌ @ REPL[5]:1 Main.f(cond, y)
│┌ @ REPL[4]: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:

vchuravy · 2022-03-18T15:52:58Z

src/compiler.jl

+end
+
+function GPUCompiler.ci_cache_populate(interp::EnzymeInterpeter, cache, mt, mi, min_world, max_world)
+	jetresult = report_call(mi.specTypes; analyzer=EnzymeAnalyzer)


How bad is the overhead here? Is this causing an inference using JETAbstractInterpreter?

Even beyond that, we should be able to handle unions if the two datatypes are pointers, yes?

I'm inclined to have this be something that runs optionally right now (and perhaps we get it to properly match the "what is handled" as time goes on).

A second avenue would be instead of:

Illegal updateAnalysis prev:{[-1]:Float@double} new: {[-1]:Integer} val: %4 = bitcast double %1 to i64, !dbg !5 origin= %.sroa.02.0 = select i1 %.not, i64 0, i64 %4, !dbg !5 julia: /workspace/srcdir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp:643: void TypeAnalyzer::updateAnalysis(llvm::Value*, TypeTree, llvm::Value*): Assertion `0 && "Performed illegal updateAnalysis"' failed. signal (6): Aborted in expression starting at REPL[5]:1 __pthread_kill_implementation at /usr/bin/../lib/libc.so.6 (unknown line) raise at /usr/bin/../lib/libc.so.6 (unknown line) abort at /usr/bin/../lib/libc.so.6 (unknown line) __assert_fail_base.cold at /usr/bin/../lib/libc.so.6 (unknown line) __assert_fail at /usr/bin/../lib/libc.so.6 (unknown line) updateAnalysis at /workspace/srcdir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp:643 visitSelectInst at /workspace/srcdir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp:1777 run at /workspace/srcdir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp:1087

Hitting an assertion error to firstly call julia_error and secondly to pass the module and instruction as well so that we can decode the !dbg !5 chain and report a better error.

(Even if we don't call julia_error we should print a stack trace on assertion that is in user code not just system code.)

How bad is the overhead here? Is this causing an inference using JETAbstractInterpreter?

The analysis uses JET's own AbstractInterpreter (not integrated with Enzyme's AbstractInterpreter), and it runs fresh analysis. So there would be duplicated work in Julia-level inference.

In terms of correctness, this non-composed approach is okay for now since EnzymeInterpreter doesn't change anything about abstract interpretation (it does tweak the optimization by chaning inlining heuristics, but it doesn't matter for JET's type-level analysis).

I'm inclined to have this be something that runs optionally right now (and perhaps we get it to properly match the "what is handled" as time goes on).

Yeah, sure. Maybe it would be reasonable to setup an optional keyword for autodiff(... ; check::Bool=false)?

Yeah, sure. Maybe it would be reasonable to setup an optional keyword for autodiff(... ; check::Bool=false)?

I think I would prefer an @report_enzyme or something.

Done: please see docstrings for the analysis APIs.
https://github.com/EnzymeAD/Enzyme.jl/pull/261/files?diff=unified&w=1

aviatesk · 2022-03-19T04:28:27Z

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.

Enzyme.jl/src/compiler.jl

Lines 2248 to 2261 in 8e3b84d

    
           # check if this return value is used in the caller and its type is inferred as Union: 
        
           # XXX `Core.Compiler.call_result_unused` is a very inaccurate model of Enzyme's activity analysis, 
        
           # and so this active Union return check might be very incomplete 
        
           function (::EnzymeAnalysisPass)(::Type{UnionResultReport}, analyzer::EnzymeAnalyzer, frame::Core.Compiler.InferenceState) 
        
               parent = frame.parent 
        
               if isa(parent, Core.Compiler.InferenceState) && !(Core.Compiler.call_result_unused(parent)) 
        
                   rt = frame.bestguess 
        
                   if isa(rt, Union) 
        
                       add_new_report!(analyzer, frame.result, UnionResultReport(frame.linfo, rt)) 
        
                       return true 
        
                   end 
        
               end 
        
               return false 
        
           end

aviatesk mentioned this pull request Mar 18, 2022

Illegal updateAnalysis prev:{[-1]:Integer} new: {[-1]:Pointer} #212

Closed

vchuravy reviewed Mar 18, 2022

View reviewed changes

aviatesk force-pushed the avi/jet branch from 895d0a4 to 8e3b84d Compare March 19, 2022 04:24

aviatesk force-pushed the avi/jet branch from 8e3b84d to 7472e41 Compare March 22, 2022 06:35

JET integration

c4245ab

aviatesk force-pushed the avi/jet branch from 7472e41 to c4245ab Compare March 22, 2022 06:35

aviatesk changed the title ~~wip: JET integration~~ JET integration Mar 22, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

JET integration #261

JET integration #261

aviatesk commented Mar 18, 2022

MathieuMorlighem commented Mar 18, 2022

vchuravy left a comment

vchuravy Mar 18, 2022

wsmoses Mar 18, 2022

vchuravy Mar 18, 2022

aviatesk Mar 19, 2022

vchuravy Mar 20, 2022

aviatesk Mar 22, 2022

aviatesk commented Mar 19, 2022 •

edited

Loading

JET integration #261

Are you sure you want to change the base?

JET integration #261

Conversation

aviatesk commented Mar 18, 2022

MathieuMorlighem commented Mar 18, 2022

vchuravy left a comment

Choose a reason for hiding this comment

vchuravy Mar 18, 2022

Choose a reason for hiding this comment

wsmoses Mar 18, 2022

Choose a reason for hiding this comment

vchuravy Mar 18, 2022

Choose a reason for hiding this comment

aviatesk Mar 19, 2022

Choose a reason for hiding this comment

vchuravy Mar 20, 2022

Choose a reason for hiding this comment

aviatesk Mar 22, 2022

Choose a reason for hiding this comment

aviatesk commented Mar 19, 2022 • edited Loading

aviatesk commented Mar 19, 2022 •

edited

Loading