EnzymeAD / Enzyme.jl

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

Error with LoopVectorization.jl #745

Closed jgreener64 closed 1 year ago

jgreener64 commented 1 year ago

I am on Julia 1.8.5, Enzyme main (4c30ed8b693a7a1d41f03bb4ef295201e26435bb) and LoopVectorization 0.12.157. This works:

using Enzyme, LoopVectorization

function f1(a, b)
    s = 0.0
    @inbounds @simd for i ∈ eachindex(a,b)
        s += a[i]*b[i]
    end
    s
end

a, b = rand(1000), rand(1000)
da, db = zero(a), zero(b)
autodiff(Reverse, f1, Active, Duplicated(a, da), Duplicated(b, db))

But this errors:

function f2(a, b)
    s = 0.0
    @turbo for i ∈ eachindex(a,b)
        s += a[i]*b[i]
    end
    s
end

da, db = zero(a), zero(b)
autodiff(Reverse, f2, Active, Duplicated(a, da), Duplicated(b, db))
ERROR: MethodError: no method matching length(::LLVM.VectorType)
Closest candidates are:
  length(::Union{Base.KeySet, Base.ValueIterator}) at abstractdict.jl:58
  length(::Union{ArrayInterfaceCore.BidiagonalIndex, ArrayInterfaceCore.TridiagonalIndex}) at ~/.julia/packages/ArrayInterfaceCore/o8YXM/src/ArrayInterfaceCore.jl:666
  length(::Union{ArrayInterface.BidiagonalIndex, ArrayInterface.TridiagonalIndex}) at ~/.julia/packages/ArrayInterface/RKU7b/src/ArrayInterface.jl:808
  ...
Stacktrace:
  [1] zero_single_allocation(builder::LLVM.IRBuilder, jlType::Type, LLVMType::LLVM.VectorType, nobj::LLVM.AllocaInst, zeroAll::Bool, idx::LLVM.ConstantInt, ctx::LLVM.Context)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:5579
  [2] zero_allocation(B::Ptr{LLVM.API.LLVMOpaqueBuilder}, LLVMType::Ptr{LLVM.API.LLVMOpaqueType}, obj::Ptr{LLVM.API.LLVMOpaqueValue}, isTape::UInt8)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:5546
  [3] 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/dev/Enzyme/src/api.jl:124
  [4] 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, Bool}, returnPrimal::Bool, jlrules::Vector{String}, expectedTapeType::Type)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:6710
  [5] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, ctx::LLVM.Context, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:7951
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, ctx::Nothing, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:8464
  [7] _thunk
    @ ~/.julia/dev/Enzyme/src/compiler.jl:8461 [inlined]
  [8] cached_compilation
    @ ~/.julia/dev/Enzyme/src/compiler.jl:8499 [inlined]
  [9] #s262#173
    @ ~/.julia/dev/Enzyme/src/compiler.jl:8557 [inlined]
 [10] var"#s262#173"(FA::Any, A::Any, TT::Any, Mode::Any, ModifiedBetween::Any, width::Any, ReturnPrimal::Any, ShadowInit::Any, World::Any, ::Any, ::Any, ::Any, ::Any, tt::Any, ::Any, ::Any, ::Any, ::Any, ::Any)
    @ Enzyme.Compiler ./none:0
 [11] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
    @ Core ./boot.jl:582
 [12] thunk(::Val{0x0000000000007f39}, ::Type{Const{typeof(f2)}}, ::Type{Active}, tt::Type{Tuple{Duplicated{Vector{Float64}}, Duplicated{Vector{Float64}}}}, ::Val{Enzyme.API.DEM_ReverseModeCombined}, ::Val{1}, ::Val{(false, false, false)}, ::Val{false})
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:8516
 [13] autodiff(::EnzymeCore.ReverseMode{false}, ::Const{typeof(f2)}, ::Type{Active}, ::Duplicated{Vector{Float64}}, ::Vararg{Duplicated{Vector{Float64}}})
    @ Enzyme ~/.julia/dev/Enzyme/src/Enzyme.jl:199
 [14] autodiff(::EnzymeCore.ReverseMode{false}, ::typeof(f2), ::Type, ::Duplicated{Vector{Float64}}, ::Vararg{Duplicated{Vector{Float64}}})
    @ Enzyme ~/.julia/dev/Enzyme/src/Enzyme.jl:214
 [15] top-level scope
    @ REPL[10]:1

Possibly related to https://github.com/EnzymeAD/Enzyme.jl/issues/98.

vchuravy commented 1 year ago

LoopVectorization will not be supported by Enzyme. It's compilation model breaks some of the assumption that Enzyme has (e.g. running before scheduling)

wsmoses commented 1 year ago

@vchuravy this bug is separate though and should be fixed regardless, however.