much better defaults

src/LinearSolve.jl

@@ -39,22 +39,11 @@ include("factorization.jl")
 include("iterative_wrappers.jl")
 include("preconditioners.jl")
 include("default.jl")
+include("init.jl")
 
 const IS_OPENBLAS = Ref(true)
 isopenblas() = IS_OPENBLAS[]
 
-function __init__()
-  @static if VERSION < v"1.7beta"
-    blas = BLAS.vendor()
-    IS_OPENBLAS[] = blas == :openblas64 || blas == :openblas
-  else
-    IS_OPENBLAS[] = occursin("openblas", BLAS.get_config().loaded_libs[1].libname)
-  end
-
-  @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
-  @require Pardiso="46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2" include("pardiso.jl")
-end
-
 export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization,
        RFLUFactorization, UMFPACKFactorization, KLUFactorization
 export KrylovJL, KrylovJL_CG, KrylovJL_GMRES, KrylovJL_BICGSTAB, KrylovJL_MINRES,

src/cuda.jl

@@ -1,3 +1,26 @@
+gpu_or_cpu(x::CUDA.CuArray) = CUDA.CuArray
+gpu_or_cpu(x::Transpose{<:Any,<:CUDA.CuArray}) = CUDA.CuArray
+gpu_or_cpu(x::Adjoint{<:Any,<:CUDA.CuArray}) = CUDA.CuArray
+isgpu(::CUDA.CuArray) = true
+isgpu(::Transpose{<:Any,<:CUDA.CuArray}) = true
+isgpu(::Adjoint{<:Any,<:CUDA.CuArray}) = true
+ifgpufree(x::CUDA.CuArray) = CUDA.unsafe_free!(x)
+ifgpufree(x::Transpose{<:Any,<:CUDA.CuArray}) = CUDA.unsafe_free!(x.parent)
+ifgpufree(x::Adjoint{<:Any,<:CUDA.CuArray}) = CUDA.unsafe_free!(x.parent)
+
+@require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" begin
+    TrackedArray = Tracker.TrackedArray
+    gpu_or_cpu(x::TrackedArray{<:Any,<:Any,<:CUDA.CuArray}) = CUDA.CuArray
+    gpu_or_cpu(x::Adjoint{<:Any,TrackedArray{<:Any,<:Any,<:CUDA.CuArray}}) = CUDA.CuArray
+    gpu_or_cpu(x::Transpose{<:Any,TrackedArray{<:Any,<:Any,<:CUDA.CuArray}}) = CUDA.CuArray
+    isgpu(::Adjoint{<:Any,TrackedArray{<:Any,<:Any,<:CUDA.CuArray}}) = true
+    isgpu(::TrackedArray{<:Any,<:Any,<:CUDA.CuArray}) = true
+    isgpu(::Transpose{<:Any,TrackedArray{<:Any,<:Any,<:CUDA.CuArray}}) = true
+    ifgpufree(x::TrackedArray{<:Any,<:Any,<:CUDA.CuArray}) = CUDA.unsafe_free!(x.data)
+    ifgpufree(x::Adjoint{<:Any,TrackedArray{<:Any,<:Any,<:CUDA.CuArray}}) = CUDA.unsafe_free!((x.data).parent)
+    ifgpufree(x::Transpose{<:Any,TrackedArray{<:Any,<:Any,<:CUDA.CuArray}}) = CUDA.unsafe_free!((x.data).parent)
+end
+
 struct GPUOffloadFactorization <: AbstractFactorization end
 
 function SciMLBase.solve(cache::LinearCache, alg::GPUOffloadFactorization; kwargs...)

src/default.jl

@@ -9,7 +9,7 @@ function defaultalg(A,b)
     # Special case on Arrays: avoid BLAS for RecursiveFactorization.jl when
     # it makes sense according to the benchmarks, which is dependent on
     # whether MKL or OpenBLAS is being used
-    if A === nothing || A isa Matrix
+    if (A === nothing && !isgpu(b)) || A isa Matrix
         if (A === nothing || eltype(A) <: Union{Float32,Float64,ComplexF32,ComplexF64}) &&
                     ArrayInterface.can_setindex(b) && (length(b) <= 100 ||
                                               (isopenblas() && length(b) <= 500)
@@ -30,18 +30,15 @@ function defaultalg(A,b)
 
     # This catches the cases where a factorization overload could exist
     # For example, BlockBandedMatrix
-    elseif ArrayInterface.isstructured(A)
+    elseif A !== nothing && ArrayInterface.isstructured(A)
         alg = GenericFactorization()
 
     # This catches the case where A is a CuMatrix
     # Which does not have LU fully defined
-    elseif !(A isa AbstractDiffEqOperator)
+    elseif isgpu(A) || isgpu(b)
         alg = QRFactorization(false)
 
     # Not factorizable operator, default to only using A*x
-    # IterativeSolvers is faster on CPU but not GPU-compatible
-    elseif cache.u isa Array
-        alg = IterativeSolversJL_GMRES()
     else
         alg = KrylovJL_GMRES()
     end
@@ -92,15 +89,12 @@ function SciMLBase.solve(cache::LinearCache, alg::Nothing,
 
     # This catches the case where A is a CuMatrix
     # Which does not have LU fully defined
-    elseif !(A isa AbstractDiffEqOperator)
+    elseif isgpu(A)
         alg = QRFactorization(false)
         SciMLBase.solve(cache, alg, args...; kwargs...)
 
     # Not factorizable operator, default to only using A*x
     # IterativeSolvers is faster on CPU but not GPU-compatible
-    elseif cache.u isa Array
-        alg = IterativeSolversJL_GMRES()
-        SciMLBase.solve(cache, alg, args...; kwargs...)
     else
         alg = KrylovJL_GMRES()
         SciMLBase.solve(cache, alg, args...; kwargs...)
@@ -147,15 +141,12 @@ function init_cacheval(alg::Nothing, A, b, u, Pl, Pr, maxiters, abstol, reltol,
 
     # This catches the case where A is a CuMatrix
     # Which does not have LU fully defined
-    elseif !(A isa AbstractDiffEqOperator)
+    elseif isgpu(A)
         alg = QRFactorization(false)
         init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
 
     # Not factorizable operator, default to only using A*x
     # IterativeSolvers is faster on CPU but not GPU-compatible
-    elseif u isa Array
-        alg = IterativeSolversJL_GMRES()
-        init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)
     else
         alg = KrylovJL_GMRES()
         init_cacheval(alg, A, b, u, Pl, Pr, maxiters, abstol, reltol, verbose)

src/init.jl

@@ -0,0 +1,13 @@
+isgpu(x) = false
+ifgpufree(x) = nothing
+function __init__()
+    @static if VERSION < v"1.7beta"
+      blas = BLAS.vendor()
+      IS_OPENBLAS[] = blas == :openblas64 || blas == :openblas
+    else
+      IS_OPENBLAS[] = occursin("openblas", BLAS.get_config().loaded_libs[1].libname)
+    end
+
+    @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
+    @require Pardiso="46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2" include("pardiso.jl")
+end