LinearSolve + Krylove + CUDA

[ASLabadi]

Hi guys,
i understood from the documentation that I can use LinearSolve with Krylov in both CPU and GPU.
i am trying to solve a sparse linear system Ax=b as the following:

Pl_cpu = ilu(A_cpu) # the preconditioner
prob_cpu = LinearProblem(A_cpu, b_cpu)
sol = solve(prob_cpu, KrylovJL_MINRES(rtol=1e-8),Pl =Pl_cpu, maxiters =200, verbose = true, log = true);

This version works super fine.

But when I try to perform it on the GPU, I convert A, b to CuSparseMatrixCSC.

P_gpul=ilu02(A_gpu)
prob_gpu=LinearProblem(A_gpu, b_gpu)
sol = solve(prob_gpu, KrylovJL_MINRES(rtol=1e-8),Pl =Pl_gpu, maxiters =200, verbose = true, log = true);

I get an error:

ERROR: MethodError: no method matching ldiv!(::CuSparseMatrixCSR{Float64, Int32}, ::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer})

Closest candidates are:
ldiv!(::AbstractVecOrMat, ::AbstractVecOrMat, ::SciMLOperators.AbstractSciMLOperator)
@ SciMLOperators C:\Users\ahmad.julia\packages\SciMLOperators\M8FAe\src\left.jl:27
ldiv!(::Any, ::Sparspak.SpkSparseSolver.SparseSolver{IT, FT}, ::Any) where {IT, FT}
@ Sparspak C:\Users\ahmad.julia\packages\Sparspak\oqBYl\src\SparseCSCInterface\SparseCSCInterface.jl:263
ldiv!(::Any, ::ChainRulesCore.AbstractThunk)
@ ChainRulesCore C:\Users\ahmad.julia\packages\ChainRulesCore\0t04l\src\tangent_types\thunks.jl:95
...

[ChrisRackauckas]

Sparse GPU is something that's still a bit in progress. The current state is SciML/OrdinaryDiffEq.jl#1566

[vpuri3]

1. does it work without preconditioner on gpu?
2. where is the error happening? can you share a little more of the stacktrace?

[ChrisRackauckas]

Last I checked, without a preconditioner should work SciML/OrdinaryDiffEq.jl#1566 (comment) . It's the preconditioners that cause an issue. They need ldiv! wrapped in CUDA.jl.

[amontoison]

If it can help, I explain how we can do the backward and forward sweeps with IC(0) and ILU(0) preconditioners here.

Note that the sparse CSR format is more relevant if you solve square system with a Krylov method.
A * v is faster with this format.

[ChrisRackauckas]

Hmm I wasn't able to get them to line up:

using LinearAlgebra, CUDA.CUSPARSE, CUDSS

T = Float32
n = 100
A_cpu = sprand(T, n, n, 0.05) + I
x_cpu = zeros(T, n)
b_cpu = rand(T, n)

A_gpu = CuSparseMatrixCSR(A_cpu)
b_gpu = CuVector(b_cpu)
prob = LinearProblem(A_gpu, b_gpu)

sol1 = solve(prob, LUFactorization())

function LinearAlgebra.ldiv!(y::AbstractArray, A::CuSparseMatrixCSR, x::AbstractArray)
    copyto!(y, x)
    ldiv!(UnitLowerTriangular(A), y)  # Forward substitution with L
    ldiv!(UpperTriangular(A), y)      # Backward substitution with U
    return y
end

CUDA.CUSPARSE.ilu02!(A_gpu, 'O')
sol2 = solve(prob, KrylovJL_GMRES(precs = (A, p) -> (CUDA.CUSPARSE.ilu02!(A, 'O'), I)))
sol1.u - sol2.u

That ldiv! doesn't seem right. @oscardssmith how with the current preconditioner interface would we get the right ldiv! in there?

[amontoison]

@ChrisRackauckas I defined kp_ilu0 and kp_ic0 in KrylovPreconditioners.jl with optimized operators for sparse matrix vector products and sparse triangular solves.

I implemented a ldiv! there for the preconditioners.
-> https://github.com/JuliaSmoothOptimizers/KrylovPreconditioners.jl/blob/main/ext/CUDA/ilu0.jl

[amontoison]

Off-topic: On all GPU backends, we need to perform an analysis phase to detect parallelism when a do sparse matrix vector products (and sparse triangular solves) as well as allocating a buffer.

We don't have an API for that with ldiv! and mul! so we need to reperform this kind of analysis / buffer allocation at each ldiv! / mul!.
With the small wrappers in KrylovPreconditioners.jl, the speed-up is between 10-30% by just wrapping the CSR matrix in KrylovOperator or TriangularOperator.

It could be relevant for LinearSolve.jl.

[oscardssmith]

@ChrisRackauckas I'm not really sure I understand the problem here. Also, by "current preconditioner interface" do you mean the old one or the new one?

[j-fu]

Isn't it possible to implement this behind the precs based interface? See https://docs.sciml.ai/LinearSolve/stable/basics/Preconditioners/#Specifying-Preconditioners and, for an example, https://github.com/JuliaLinearAlgebra/AlgebricMultigrid.jl/pull/116/files ?

An interesting point is here that this can be implemented without creating a dependency on LinearSolve.

[oscardssmith]

@ChrisRackauckas you were close (and the main problem was that Krylov docs have a typo for CSR). The following works:

struct ILU0{T}
    A::T
end

function LinearAlgebra.ldiv!(y::AbstractArray, A::ILU0, x::AbstractArray)
    copyto!(y, x)
    ldiv!(LowerTriangular(A.A), y)  # Forward substitution with L
    ldiv!(UnitUpperTriangular(A.A), y)      # Backward substitution with U
    return y
end

precs(A, p) = (ILU0(CUSPARSE.ilu02(A)))
sol2 = solve(prob, KrylovJL(;precs))

julia> norm(A_gpu*sol2.u-b_gpu)
2.1753882f-5

[amontoison]

The problem is bicgstab, not the preconditioner or the implementation in Krylov.jl.
It can breakdown at any moment.
Even if I highly recommend to use other Krylov methods, an alternative is to restart bicgstab with the current solution with the optional argument x0.
You can also try another initial second vector c used for bi-orthogonality.

[oscardssmith]

gmres also seemed unhappy with the original version.

[amontoison]

What does it mean "unhappy", does it not converge?
Can you display the output with verbose=1 ?
If you run the solver in single precision Float32, it's hard to expect good results.