search

EnzymeAD / Enzyme.jl

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

Enzyme compilation error when using a sparse cholesky #997

Closed ptiede closed 9 months ago

ptiede commented 1 year ago

I get an error using Enzyme when there is a sparse Cholesky back substitution. My guess is that this is just a missing rule? A mwe is

using Enzyme
using SparseArrays, LinearAlgebra

# Otherwise it complains about the activity of C even though it is Const
Enzyme.API.runtimeActivity!(true)

f(x, C) = sum(C\x)

N = 32
σ = sprand(N, N, 0.05)
Σ = (σ + σ') + Diagonal(fill(5.0, N))
C = cholesky(Σ)

x = randn(N)
f(x, C)

dx = zero(x)
autodiff(Reverse, f, Active, Duplicated(x, dx), Const(C))

When I run this on Julia 1.10-beta1 I get the following stack trace (I have omitted the beginning print out from Enzyme)

No augmented forward pass found for cholmod_l_start
declare i32 @cholmod_l_start(i64) local_unnamed_addr

Stacktrace:
  [1] julia_error(cstr::Cstring, val::Ptr{LLVM.API.LLVMOpaqueValue}, errtype::Enzyme.API.ErrorType, data::Ptr{Nothing}, data2::Ptr{LLVM.API.LLVMOpaqueValue}, B::Ptr{LLVM.API.LLVMOpaqueBuilder})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:5697
  [2] EnzymeCreateAugmentedPrimal(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnUsed::Bool, shadowReturnUsed::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, forceAnonymousTape::Bool, width::Int64, atomicAdd::Bool)
    @ Enzyme.API ~/.julia/packages/Enzyme/ZH7JA/src/api.jl:164
  [3] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{…}, returnPrimal::Bool, jlrules::Vector{…}, expectedTapeType::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:7452
  [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/packages/Enzyme/ZH7JA/src/compiler.jl:9028
  [5] codegen
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:8634 [inlined]
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool) (repeats 2 times)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9548
  [7] cached_compilation
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9582 [inlined]
  [8] (::Enzyme.Compiler.var"#463#464"{DataType, DataType, DataType, Enzyme.API.CDerivativeMode, NTuple{4, Bool}, Int64, Bool, Bool, UInt64, DataType})(ctx::LLVM.Context)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9645
  [9] JuliaContext(f::Enzyme.Compiler.var"#463#464"{DataType, DataType, DataType, Enzyme.API.CDerivativeMode, NTuple{4, Bool}, Int64, Bool, Bool, UInt64, DataType})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:47
 [10] #s292#462
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9600 [inlined]
 [11] 
    @ Enzyme.Compiler ./none:0
 [12] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any})
    @ Core ./boot.jl:600
 [13] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(get!), df::Nothing, primal_1::typeof(SparseArrays.CHOLMOD.newcommon_l), shadow_1_1::Nothing, primal_2::IdDict{…}, shadow_2_1::IdDict{…}, primal_3::Symbol, shadow_3_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:1354
 [14] getcommon
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:170 [inlined]
 [15] getcommon
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:173 [inlined]
 [16] allocate_dense
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:448
 [17] Dense
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:793
 [18] \
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:1680 [inlined]
 [19] f
    @ ~/Research/Enzyme/mwe_sparse_chol.jl:7 [inlined]
 [20] f
    @ ~/Research/Enzyme/mwe_sparse_chol.jl:0 [inlined]
 [21] diffejulia_f_2414_inner_1wrap
    @ ~/Research/Enzyme/mwe_sparse_chol.jl:0
 [22] macro expansion
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9496 [inlined]
 [23] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Type{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Const{…}, ::Float64)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9174
 [24] (::Enzyme.Compiler.CombinedAdjointThunk{Ptr{…}, Const{…}, Active{…}, Tuple{…}, Val{…}, Val{…}()})(::Const{typeof(f)}, ::Duplicated{Vector{…}}, ::Vararg{Any})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9137
 [25] autodiff(::ReverseMode{false, FFIABI}, ::Const{typeof(f)}, ::Type{Active}, ::Duplicated{Vector{Float64}}, ::Vararg{Any})
    @ Enzyme ~/.julia/packages/Enzyme/ZH7JA/src/Enzyme.jl:213
 [26] autodiff(::ReverseMode{false, FFIABI}, ::typeof(f), ::Type, ::Duplicated{Vector{Float64}}, ::Vararg{Any})
    @ Enzyme ~/.julia/packages/Enzyme/ZH7JA/src/Enzyme.jl:222
 [27] top-level scope
    @ ~/Research/Enzyme/mwe_sparse_chol.jl:18
Some type information was truncated. Use `show(err)` to see complete types.

My Project is

Status `~/Research/Enzyme/Project.toml`
  [7da242da] Enzyme v0.11.7 `https://github.com/EnzymeAD/Enzyme.jl.git#main`
  [37e2e46d] LinearAlgebra
  [2f01184e] SparseArrays v1.10.0
julia> versioninfo()
Julia Version 1.10.0-beta1
Commit 6616549950e (2023-07-25 17:43 UTC)
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 20 × 12th Gen Intel(R) Core(TM) i7-12700H
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, alderlake)
  Threads: 1 on 20 virtual cores
Environment:
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 1
wsmoses commented 1 year ago

Haven't seen that function before. In any case, retry on latest main and it should give you a nicer stacktrace.

I'll look into adding support/a rule for that soon.

On Fri, Aug 11, 2023 at 11:40 AM Paul Tiede @.***> wrote:

I get an error using Enzyme when there is a sparse Cholesky back substitution. My guess is that this is just a missing rule? A mwe is

using Enzymeusing SparseArrays, LinearAlgebra

Otherwise it complains about the activity of C even though it is Const

Enzyme.API.runtimeActivity!(true) f(x, C) = sum(C\x)

N = 32 σ = sprand(N, N, 0.05) Σ = (σ + σ') + Diagonal(fill(5.0, N)) C = cholesky(Σ)

x = randn(N)f(x, C)

dx = zero(x)autodiff(Reverse, f, Active, Duplicated(x, dx), Const(C))

When I run this on Julia 1.10-beta1 I get the following stack trace (I have omitted the beginning print out from Enzyme)

No augmented forward pass found for cholmod_l_start declare i32 @cholmod_l_start(i64) local_unnamed_addr

Stacktrace: [1] julia_error(cstr::Cstring, val::Ptr{LLVM.API.LLVMOpaqueValue}, errtype::Enzyme.API.ErrorType, data::Ptr{Nothing}, data2::Ptr{LLVM.API.LLVMOpaqueValue}, B::Ptr{LLVM.API.LLVMOpaqueBuilder}) @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:5697 [2] EnzymeCreateAugmentedPrimal(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnUsed::Bool, shadowReturnUsed::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, forceAnonymousTape::Bool, width::Int64, atomicAdd::Bool) @ Enzyme.API ~/.julia/packages/Enzyme/ZH7JA/src/api.jl:164 [3] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{…}, returnPrimal::Bool, jlrules::Vector{…}, expectedTapeType::Type) @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:7452 [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/packages/Enzyme/ZH7JA/src/compiler.jl:9028 [5] codegen @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:8634 [inlined] [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool) (repeats 2 times) @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9548 [7] cached_compilation @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9582 [inlined] [8] (::Enzyme.Compiler.var"#463#464"{DataType, DataType, DataType, Enzyme.API.CDerivativeMode, NTuple{4, Bool}, Int64, Bool, Bool, UInt64, DataType})(ctx::LLVM.Context) @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9645 [9] JuliaContext(f::Enzyme.Compiler.var"#463#464"{DataType, DataType, DataType, Enzyme.API.CDerivativeMode, NTuple{4, Bool}, Int64, Bool, Bool, UInt64, DataType}) @ GPUCompiler ~/.julia/packages/GPUCompiler/YO8Uj/src/driver.jl:47 [10] #s292#462 @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9600 [inlined] [11] @ Enzyme.Compiler ./none:0 [12] (::Core.GeneratedFunctionStub)(::UInt64, ::LineNumberNode, ::Any, ::Vararg{Any}) @ Core ./boot.jl:600 [13] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(get!), df::Nothing, primal_1::typeof(SparseArrays.CHOLMOD.newcommon_l), shadow_1_1::Nothing, primal_2::IdDict{…}, shadow_2_1::IdDict{…}, primal_3::Symbol, shadow_3_1::Nothing) @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:1354 [14] getcommon @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:170 [inlined] [15] getcommon @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:173 [inlined] [16] allocate_dense @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:448 [17] Dense @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:793 [18] \ @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:1680 [inlined] [19] f @ ~/Research/Enzyme/mwe_sparse_chol.jl:7 [inlined] [20] f @ ~/Research/Enzyme/mwe_sparse_chol.jl:0 [inlined] [21] diffejulia_f_2414_inner_1wrap @ ~/Research/Enzyme/mwe_sparse_chol.jl:0 [22] macro expansion @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9496 [inlined] [23] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Type{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Const{…}, ::Float64) @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9174 [24] (::Enzyme.Compiler.CombinedAdjointThunk{Ptr{…}, Const{…}, Active{…}, Tuple{…}, Val{…}, Val{…}()})(::Const{typeof(f)}, ::Duplicated{Vector{…}}, ::Vararg{Any}) @ Enzyme.Compiler ~/.julia/packages/Enzyme/ZH7JA/src/compiler.jl:9137 [25] autodiff(::ReverseMode{false, FFIABI}, ::Const{typeof(f)}, ::Type{Active}, ::Duplicated{Vector{Float64}}, ::Vararg{Any}) @ Enzyme ~/.julia/packages/Enzyme/ZH7JA/src/Enzyme.jl:213 [26] autodiff(::ReverseMode{false, FFIABI}, ::typeof(f), ::Type, ::Duplicated{Vector{Float64}}, ::Vararg{Any}) @ Enzyme ~/.julia/packages/Enzyme/ZH7JA/src/Enzyme.jl:222 [27] top-level scope @ ~/Research/Enzyme/mwe_sparse_chol.jl:18 Some type information was truncated. Use show(err) to see complete types.

My Project is

Status ~/Research/Enzyme/Project.toml [7da242da] Enzyme v0.11.7 https://github.com/EnzymeAD/Enzyme.jl.git#main [37e2e46d] LinearAlgebra [2f01184e] SparseArrays v1.10.0

julia> versioninfo() Julia Version 1.10.0-beta1 Commit 6616549950e (2023-07-25 17:43 UTC) Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 20 × 12th Gen Intel(R) Core(TM) i7-12700H WORD_SIZE: 64 LIBM: libopenlibm LLVM: libLLVM-15.0.7 (ORCJIT, alderlake) Threads: 1 on 20 virtual cores Environment: JULIA_EDITOR = code JULIA_NUM_THREADS = 1

— Reply to this email directly, view it on GitHub https://github.com/EnzymeAD/Enzyme.jl/issues/997, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXG3OXYTSDR5N3SEBH3XUZHAXANCNFSM6AAAAAA3NC3BEI . You are receiving this because you are subscribed to this thread.Message ID: @.***>

ptiede commented 1 year ago

Ok on Enzyme#main I now get

warning: didn't implement memmove, using memcpy as fallback which can result in errors
ERROR: Enzyme execution failed.
Enzyme: Not yet implemented augmented forward for jl_eqtable_get
Stacktrace:
 [1] get!
   @ ./iddict.jl:179
 [2] get!
   @ ./iddict.jl:0

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/LOvs5/src/compiler.jl:2986
  [2] get!
    @ ./iddict.jl:179 [inlined]
  [3] get!
    @ ./iddict.jl:0 [inlined]
  [4] augmented_julia_get__5506_inner_1wrap
    @ ./iddict.jl:0
  [5] macro expansion
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/LOvs5/src/compiler.jl:9619 [inlined]
  [6] enzyme_call
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/LOvs5/src/compiler.jl:9297 [inlined]
  [7] AugmentedForwardThunk
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/LOvs5/src/compiler.jl:9271 [inlined]
  [8] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(get!), df::Nothing, primal_1::typeof(SparseArrays.CHOLMOD.newcommon_l), shadow_1_1::Nothing, primal_2::IdDict{…}, shadow_2_1::IdDict{…}, primal_3::Symbol, shadow_3_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/LOvs5/src/compiler.jl:1365
  [9] getcommon
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:170 [inlined]
 [10] getcommon
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:173 [inlined]
 [11] allocate_dense
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:448
 [12] Dense
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:793
 [13] \
    @ ~/.julia/juliaup/julia-1.10.0-beta1+0.x64.linux.gnu/share/julia/stdlib/v1.10/SparseArrays/src/solvers/cholmod.jl:1680 [inlined]
 [14] f
    @ ~/Research/Enzyme/mwe_sparse_chol.jl:7 [inlined]
 [15] f
    @ ~/Research/Enzyme/mwe_sparse_chol.jl:0 [inlined]
 [16] diffejulia_f_3285_inner_1wrap
    @ ~/Research/Enzyme/mwe_sparse_chol.jl:0
 [17] macro expansion
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/LOvs5/src/compiler.jl:9619 [inlined]
 [18] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Type{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Const{…}, ::Float64)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/LOvs5/src/compiler.jl:9297
 [19] (::Enzyme.Compiler.CombinedAdjointThunk{Ptr{…}, Const{…}, Active{…}, Tuple{…}, Val{…}, Val{…}()})(::Const{typeof(f)}, ::Duplicated{Vector{…}}, ::Vararg{Any})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/LOvs5/src/compiler.jl:9260
 [20] autodiff(::ReverseMode{false, FFIABI}, ::Const{typeof(f)}, ::Type{Active}, ::Duplicated{Vector{Float64}}, ::Vararg{Any})
    @ Enzyme ~/.julia/packages/Enzyme/LOvs5/src/Enzyme.jl:213
 [21] autodiff(::ReverseMode{false, FFIABI}, ::typeof(f), ::Type, ::Duplicated{Vector{Float64}}, ::Vararg{Any})
    @ Enzyme ~/.julia/packages/Enzyme/LOvs5/src/Enzyme.jl:222
 [22] top-level scope
    @ ~/Research/Enzyme/mwe_sparse_chol.jl:18
Some type information was truncated. Use `show(err)` to see complete types.
wsmoses commented 9 months ago

Interestingly enough, trying this now it appears to succeed. Reopen if it persists?