From 0eda0bbf077c9c06a65688f1a5b9571edb24addb Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Fri, 20 Jun 2025 11:06:18 +0000 Subject: [PATCH 1/4] Fix CUDSS dispatches Fixes https://github.com/SciML/LinearSolve.jl/issues/625 --- ext/LinearSolveCUDAExt.jl | 4 ++++ ext/LinearSolveSparseArraysExt.jl | 36 +++++++++++++++++++++++++++++-- src/LinearSolve.jl | 1 + 3 files changed, 39 insertions(+), 2 deletions(-) diff --git a/ext/LinearSolveCUDAExt.jl b/ext/LinearSolveCUDAExt.jl index 94f6d7df0..291c8bcec 100644 --- a/ext/LinearSolveCUDAExt.jl +++ b/ext/LinearSolveCUDAExt.jl @@ -5,6 +5,10 @@ using LinearSolve using LinearSolve.LinearAlgebra, LinearSolve.SciMLBase, LinearSolve.ArrayInterface using SciMLBase: AbstractSciMLOperator +function LinearSolve.is_cusparse(A::Union{CUDA.CUSPARSE.CuSparseMatrixCSR, CUDA.CUSPARSE.CuSparseMatrixCSC}) + true +end + function LinearSolve.defaultalg(A::CUDA.CUSPARSE.CuSparseMatrixCSR{Tv, Ti}, b, assump::OperatorAssumptions{Bool}) where {Tv, Ti} if LinearSolve.cudss_loaded(A) diff --git a/ext/LinearSolveSparseArraysExt.jl b/ext/LinearSolveSparseArraysExt.jl index d04e8117f..d48079b66 100644 --- a/ext/LinearSolveSparseArraysExt.jl +++ b/ext/LinearSolveSparseArraysExt.jl @@ -100,7 +100,11 @@ function LinearSolve.init_cacheval( Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) where {T<:BLASELTYPES} - SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int64}(zero(Int64), zero(Int64), [Int64(1)], Int64[], T[])) + if is_cusparse(A) + ArrayInterface.lu_instance(A) + else + SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int64}(zero(Int64), zero(Int64), [Int64(1)], Int64[], T[])) + end end function LinearSolve.init_cacheval( @@ -108,7 +112,19 @@ function LinearSolve.init_cacheval( Pl, Pr, maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) where {T<:BLASELTYPES} - SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int32}(zero(Int32), zero(Int32), [Int32(1)], Int32[], T[])) + if LinearSolve.is_cusparse(A) + ArrayInterface.lu_instance(A) + else + SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC{T, Int32}(zero(Int32), zero(Int32), [Int32(1)], Int32[], T[])) + end +end + +function LinearSolve.init_cacheval( + alg::LUFactorization, A::LinearSolve.GPUArraysCore.AnyGPUArray, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Bool, assumptions::OperatorAssumptions) + ArrayInterface.lu_instance(A) end function LinearSolve.init_cacheval( @@ -120,6 +136,14 @@ function LinearSolve.init_cacheval( PREALLOCATED_UMFPACK end +function LinearSolve.init_cacheval( + alg::UMFPACKFactorization, A::LinearSolve.GPUArraysCore.AnyGPUArray, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Bool, assumptions::OperatorAssumptions) + nothing +end + function LinearSolve.init_cacheval( alg::UMFPACKFactorization, A::AbstractSparseArray{T, Int64}, b, u, Pl, Pr, @@ -191,6 +215,14 @@ function LinearSolve.init_cacheval( PREALLOCATED_KLU end +function LinearSolve.init_cacheval( + alg::KLUFactorization, A::LinearSolve.GPUArraysCore.AnyGPUArray, b, u, + Pl, Pr, + maxiters::Int, abstol, reltol, + verbose::Bool, assumptions::OperatorAssumptions) + nothing +end + function LinearSolve.init_cacheval( alg::KLUFactorization, A::AbstractSparseArray{Float64, Int32}, b, u, Pl, Pr, maxiters::Int, abstol, diff --git a/src/LinearSolve.jl b/src/LinearSolve.jl index e5673310c..270b64932 100644 --- a/src/LinearSolve.jl +++ b/src/LinearSolve.jl @@ -217,6 +217,7 @@ end ALREADY_WARNED_CUDSS = Ref{Bool}(false) error_no_cudss_lu(A) = nothing cudss_loaded(A) = false +is_cusparse(A) = false export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization, GenericLUFactorization, SimpleLUFactorization, RFLUFactorization, From f2d31e816ac1e55105fd4b7f3df5f1a562680354 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Fri, 20 Jun 2025 11:08:32 +0000 Subject: [PATCH 2/4] add CUDSS test --- test/gpu/Project.toml | 1 + test/gpu/cuda.jl | 15 +++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml index 5be1abdf4..4c304c0f2 100644 --- a/test/gpu/Project.toml +++ b/test/gpu/Project.toml @@ -1,6 +1,7 @@ [deps] BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" diff --git a/test/gpu/cuda.jl b/test/gpu/cuda.jl index 798813d5c..2b2dde30a 100644 --- a/test/gpu/cuda.jl +++ b/test/gpu/cuda.jl @@ -1,4 +1,5 @@ using LinearSolve, CUDA, LinearAlgebra, SparseArrays, StableRNGs +using CUDA.CUSPARSE, CUDSS using Test CUDA.allowscalar(false) @@ -91,3 +92,17 @@ prob2 = LinearProblem(transpose(A), b) sol = solve(prob2, alg; alias = LinearAliasSpecifier(alias_A = false)) @test norm(transpose(A) * sol.u .- b) < 1e-5 end + +@testset "CUDSS" begin + 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_csr = CuSparseMatrixCSR(A_cpu) + b_gpu = CuVector(b_cpu) + + prob = LinearProblem(A_gpu_csr, b_gpu) + sol = solve(prob) +end \ No newline at end of file From a082082e4596028e451593b4c5af5b131a53c942 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Fri, 20 Jun 2025 11:13:14 +0000 Subject: [PATCH 3/4] add missing import --- ext/LinearSolveSparseArraysExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/LinearSolveSparseArraysExt.jl b/ext/LinearSolveSparseArraysExt.jl index d48079b66..028dd84a0 100644 --- a/ext/LinearSolveSparseArraysExt.jl +++ b/ext/LinearSolveSparseArraysExt.jl @@ -3,7 +3,7 @@ module LinearSolveSparseArraysExt using LinearSolve, LinearAlgebra using SparseArrays using SparseArrays: AbstractSparseMatrixCSC, nonzeros, rowvals, getcolptr -using LinearSolve: BLASELTYPES, pattern_changed +using LinearSolve: BLASELTYPES, pattern_changed, ArrayInterface # Can't `using KLU` because cannot have a dependency in there without # requiring the user does `using KLU` From e07ed9f0d2aeb8c4c75626ffb933e11669c1c351 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Fri, 20 Jun 2025 12:13:53 +0000 Subject: [PATCH 4/4] Update pipeline.yml --- .buildkite/pipeline.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index fa895d8e0..b1931c260 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -12,6 +12,6 @@ steps: GROUP: 'LinearSolveCUDA' JULIA_PKG_SERVER: "" # it often struggles with our large artifacts # SECRET_CODECOV_TOKEN: "..." - timeout_in_minutes: 30 + timeout_in_minutes: 180 # Don't run Buildkite if the commit message includes the text [skip tests] if: build.message !~ /\[skip tests\]/