Merge pull request #470 from SciML/ap/overhead

[Do Not Merge] Preallocate more caches

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "LinearSolve"
 uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 authors = ["SciML"]
-version = "2.23.2"
+version = "2.23.3"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/common.jl

@@ -90,7 +90,8 @@ function Base.setproperty!(cache::LinearCache, name::Symbol, x)
         update_cacheval!(cache, :b, x)
     elseif name === :cacheval && cache.alg isa DefaultLinearSolver
         @assert cache.cacheval isa DefaultLinearSolverInit
-        return setfield!(cache.cacheval, Symbol(cache.alg.alg), x)
+        return __setfield!(cache.cacheval, cache.alg, x)
+        # return setfield!(cache.cacheval, Symbol(cache.alg.alg), x)
     end
     setfield!(cache, name, x)
 end

src/default.jl

@@ -24,6 +24,23 @@ mutable struct DefaultLinearSolverInit{T1, T2, T3, T4, T5, T6, T7, T8, T9, T10,
     KrylovJL_LSMR::T21
 end
 
+@generated function __setfield!(cache::DefaultLinearSolverInit, alg::DefaultLinearSolver, v)
+    ex = :()
+    for alg in first.(EnumX.symbol_map(DefaultAlgorithmChoice.T))
+        newex = quote
+            setfield!(cache, $(Meta.quot(alg)), v)
+        end
+        alg_enum = getproperty(LinearSolve.DefaultAlgorithmChoice, alg) 
+        ex = if ex == :()
+            Expr(:elseif, :(alg.alg == $(alg_enum)), newex,
+                :(error("Algorithm Choice not Allowed")))
+        else
+            Expr(:elseif, :(alg.alg == $(alg_enum)), newex, ex)
+        end
+    end
+    ex = Expr(:if, ex.args...)
+end
+
 # Legacy fallback
 # For SciML algorithms already using `defaultalg`, all assume square matrix.
 defaultalg(A, b) = defaultalg(A, b, OperatorAssumptions(true))
@@ -159,7 +176,7 @@ function defaultalg(A, b, assump::OperatorAssumptions{Bool})
                (__conditioning(assump) === OperatorCondition.IllConditioned ||
                 __conditioning(assump) === OperatorCondition.WellConditioned)
                 if length(b) <= 10
-                    DefaultAlgorithmChoice.GenericLUFactorization
+                    DefaultAlgorithmChoice.RFLUFactorization
                 elseif appleaccelerate_isavailable()
                     DefaultAlgorithmChoice.AppleAccelerateLUFactorization
                 elseif (length(b) <= 100 || (isopenblas() && length(b) <= 500) ||
@@ -345,11 +362,12 @@ end
                 retcode = sol.retcode,
                 iters = sol.iters, stats = sol.stats)
         end
+        alg_enum = getproperty(LinearSolve.DefaultAlgorithmChoice, alg) 
         ex = if ex == :()
-            Expr(:elseif, :(Symbol(alg.alg) === $(Meta.quot(alg))), newex,
+            Expr(:elseif, :(alg.alg == $(alg_enum)), newex,
                 :(error("Algorithm Choice not Allowed")))
         else
-            Expr(:elseif, :(Symbol(alg.alg) === $(Meta.quot(alg))), newex, ex)
+            Expr(:elseif, :(alg.alg == $(alg_enum)), newex, ex)
         end
     end
     ex = Expr(:if, ex.args...)

src/factorization.jl

@@ -173,18 +173,24 @@ function init_cacheval(alg::QRFactorization, A, b, u, Pl, Pr,
     ArrayInterface.qr_instance(convert(AbstractMatrix, A), alg.pivot)
 end
 
+const PREALLOCATED_QR_ColumnNorm = ArrayInterface.qr_instance(rand(1, 1), ColumnNorm())
+
+function init_cacheval(alg::QRFactorization{ColumnNorm}, A::Matrix{Float64}, b, u, Pl, Pr,
+    maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions)
+    return PREALLOCATED_QR_ColumnNorm
+end
+
 function init_cacheval(alg::QRFactorization, A::Union{<:Adjoint, <:Transpose}, b, u, Pl, Pr,
     maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions)
     A isa GPUArraysCore.AnyGPUArray && return qr(A)
     return qr(A, alg.pivot)
 end
 
-const PREALLOCATED_QR = ArrayInterface.qr_instance(rand(1, 1))
+const PREALLOCATED_QR_NoPivot = ArrayInterface.qr_instance(rand(1, 1))
 
 function init_cacheval(alg::QRFactorization{NoPivot}, A::Matrix{Float64}, b, u, Pl, Pr,
-    maxiters::Int, abstol, reltol, verbose::Bool,
-    assumptions::OperatorAssumptions)
-    PREALLOCATED_QR
+    maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions)
+    return PREALLOCATED_QR_NoPivot
 end
 
 function init_cacheval(alg::QRFactorization, A::AbstractSciMLOperator, b, u, Pl, Pr,
@@ -1013,6 +1019,14 @@ function init_cacheval(alg::NormalCholeskyFactorization, A, b, u, Pl, Pr,
     return ArrayInterface.cholesky_instance(Symmetric(Matrix{eltype(A)}(undef,0,0)), alg.pivot)
 end
 
+const PREALLOCATED_NORMALCHOLESKY_SYMMETRIC = ArrayInterface.cholesky_instance(
+    Symmetric(rand(1, 1)), NoPivot())
+
+function init_cacheval(alg::NormalCholeskyFactorization, A::Matrix{Float64}, b, u, Pl, Pr,
+    maxiters::Int, abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions)
+    return PREALLOCATED_NORMALCHOLESKY_SYMMETRIC
+end
+
 function init_cacheval(alg::NormalCholeskyFactorization,
     A::Union{Diagonal, AbstractSciMLOperator}, b, u, Pl, Pr,
     maxiters::Int, abstol, reltol, verbose::Bool,

test/default_algs.jl

@@ -1,6 +1,6 @@
 using LinearSolve, LinearAlgebra, SparseArrays, Test, JET
 @test LinearSolve.defaultalg(nothing, zeros(3)).alg ===
-      LinearSolve.DefaultAlgorithmChoice.GenericLUFactorization
+      LinearSolve.DefaultAlgorithmChoice.RFLUFactorization
 prob = LinearProblem(rand(3, 3), rand(3))
 solve(prob)