AD for sparse matrix in DJUICE #2565
Description
Hi @wsmoses,
Could you please help us with the AD for backslash in sparse matrix in DJUICE.jl
The version is: https://github.com/DJ4Earth/DJUICE.jl/tree/sparsematrix
I have a mini example:
using Enzyme
using DJUICE
n = 5
A = DJUICE.IssmMatrix(n,n)
for i = 1:n
DJUICE.AddValues!(A, 1, [i], 1,[i], ones(1,1))
end
DJUICE.Assemble!(A)
b = DJUICE.IssmVector(n)
x = DJUICE.IssmVector(n)
function mybackslash(x, A, b)
x.vector = A.matrix \ b.vector
norm(x.vector)
end
∂z_∂A = DJUICE.IssmMatrix(n,n)
∂z_∂b = DJUICE.IssmVector(n)
∂z_∂x = DJUICE.IssmVector(n)
Enzyme.autodiff(set_runtime_activity(Enzyme.Reverse), mybackslash,Active, Duplicated(x, ∂z_∂x), Duplicated(A, ∂z_∂A), Duplicated(b, ∂z_∂b))
The error message on my end is:
Stacktrace:
[1] should_recurse(typ2::Any, arg_t::LLVM.LLVMType, byref::GPUCompiler.ArgumentCC, dl::LLVM.DataLayout)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/ez9it/src/absint.jl:238
[2] abs_typeof(arg::LLVM.Value, partial::Bool, seenphis::Set{LLVM.PHIInst})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/ez9it/src/absint.jl:684
[3] abs_typeof
@ ~/.julia/packages/Enzyme/ez9it/src/absint.jl:333 [inlined]
[4] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/ez9it/src/compiler.jl:4655
[5] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
@ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:67
[6] compile
@ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:55 [inlined]
[7] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}}, postopt::Bool)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/ez9it/src/compiler.jl:5786
[8] _thunk
@ ~/.julia/packages/Enzyme/ez9it/src/compiler.jl:5784 [inlined]
[9] cached_compilation
@ ~/.julia/packages/Enzyme/ez9it/src/compiler.jl:5841 [inlined]
[10] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, edges::Vector{…})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/ez9it/src/compiler.jl:5955
[11] thunk_generator(world::UInt64, source::Union{…}, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::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, strongzero::Type)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/ez9it/src/compiler.jl:6148
[12] autodiff_thunk
@ ~/.julia/packages/Enzyme/ez9it/src/Enzyme.jl:997 [inlined]
[13] autodiff
@ ~/.julia/packages/Enzyme/ez9it/src/Enzyme.jl:391 [inlined]
[14] autodiff(::ReverseMode{false, true, false, FFIABI, false, false}, ::typeof(mybackslash), ::Type{Active}, ::Duplicated{DJUICE.IssmVector}, ::Duplicated{DJUICE.IssmMatrix}, ::Duplicated{DJUICE.IssmVector})
@ Enzyme ~/.julia/packages/Enzyme/ez9it/src/Enzyme.jl:538
[15] top-level scope
@ ~/Research/DJUICE/examples/test_sparse.jl:25
[16] include(fname::String)
@ Main ./sysimg.jl:38
[17] top-level scope
@ REPL[1]:1
in expression starting at /Users/chenggong/Research/DJUICE/examples/test_sparse.jl:25
Some type information was truncated. Use `show(err)` to see complete types.