EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
458 stars 66 forks source link

Error moving array off GPU #2139

Open jgreener64 opened 2 hours ago

jgreener64 commented 2 hours ago

I am on Enzyme main (30b6b2d93d8ef1bdfb9f628e8c111d123cc4595e), CUDA master (7ff012) and Julia 1.10.6.

using Enzyme, CUDA
function f(a, b)
    c = a .+ b
    Array(c)[1]
end
a = CuArray(rand(5))
b = CuArray(rand(5))
da = zero(a)
db = zero(b)
f(a, b) # Works
autodiff(set_runtime_activity(Reverse), f, Active, Duplicated(a, da), Duplicated(b, db))
ERROR: AssertionError: Base.isconcretetype(typ)
Stacktrace:
  [1] abs_typeof(arg::LLVM.Value, partial::Bool, seenphis::Set{LLVM.PHIInst})
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/absint.jl:581
  [2] abs_typeof(arg::LLVM.Value, partial::Bool, seenphis::Set{LLVM.PHIInst})
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/absint.jl:692
  [3] abs_typeof
    @ ~/.julia/dev/Enzyme/src/absint.jl:283 [inlined]
  [4] codegen(output::Symbol, job::GPUCompiler.CompilerJob{…}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:7227
  [5] codegen
    @ ~/.julia/dev/Enzyme/src/compiler.jl:6166 [inlined]
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:8531
  [7] _thunk
    @ ~/.julia/dev/Enzyme/src/compiler.jl:8531 [inlined]
  [8] cached_compilation
    @ ~/.julia/dev/Enzyme/src/compiler.jl:8572 [inlined]
  [9] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::Tuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:8685
 [10] thunk_generator(world::UInt64, source::LineNumberNode, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::Tuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::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)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:8837
 [11] autodiff
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:485 [inlined]
 [12] autodiff(::ReverseMode{…}, ::typeof(f), ::Type{…}, ::Duplicated{…}, ::Duplicated{…})
    @ Enzyme ~/.julia/dev/Enzyme/src/Enzyme.jl:524
 [13] top-level scope
    @ REPL[8]:1
wsmoses commented 1 hour ago

potential fix: https://github.com/EnzymeAD/Enzyme.jl/pull/2140