Special handling for staticarrays

Author: avik-pal

Project.toml

@@ -59,6 +59,7 @@ LinearSolvePardisoExt = "Pardiso"
 LinearSolveRecursiveArrayToolsExt = "RecursiveArrayTools"
 
 [compat]
+AllocCheck = "0.1"
 Aqua = "0.8"
 ArrayInterface = "7.4.11"
 BandedMatrices = "1"
@@ -109,6 +110,7 @@ UnPack = "1"
 julia = "1.9"
 
 [extras]
+AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
@@ -133,4 +135,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "Enzyme", "FiniteDiff", "BandedMatrices", "FastAlmostBandedMatrices", "StaticArrays"]
+test = ["Aqua", "Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "Enzyme", "FiniteDiff", "BandedMatrices", "FastAlmostBandedMatrices", "StaticArrays", "AllocCheck"]

src/common.jl

@@ -169,31 +169,9 @@ function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
     isfresh = true
     Tc = typeof(cacheval)
 
-    cache = LinearCache{
-        typeof(A),
-        typeof(b),
-        typeof(u0_),
-        typeof(p),
-        typeof(alg),
-        Tc,
-        typeof(Pl),
-        typeof(Pr),
-        typeof(reltol),
-        typeof(assumptions.issq),
-    }(A,
-        b,
-        u0_,
-        p,
-        alg,
-        cacheval,
-        isfresh,
-        Pl,
-        Pr,
-        abstol,
-        reltol,
-        maxiters,
-        verbose,
-        assumptions)
+    cache = LinearCache{typeof(A), typeof(b), typeof(u0_), typeof(p), typeof(alg), Tc,
+        typeof(Pl), typeof(Pr), typeof(reltol), typeof(assumptions.issq)}(A, b, u0_,
+        p, alg, cacheval, isfresh, Pl, Pr, abstol, reltol, maxiters, verbose, assumptions)
     return cache
 end
 
@@ -210,3 +188,33 @@ end
 function SciMLBase.solve!(cache::LinearCache, args...; kwargs...)
     solve!(cache, cache.alg, args...; kwargs...)
 end
+
+# Special Case for StaticArrays
+const StaticLinearProblem = LinearProblem{uType, iip, <:SMatrix,
+    <:Union{<:SMatrix, <:SVector}} where {uType, iip}
+
+function SciMLBase.solve(prob::StaticLinearProblem, args...; kwargs...)
+    return SciMLBase.solve(prob, nothing, args...; kwargs...)
+end
+
+function SciMLBase.solve(prob::StaticLinearProblem,
+        alg::Union{Nothing, SciMLLinearSolveAlgorithm}, args...; kwargs...)
+    if alg === nothing || alg isa DirectLdiv!
+        u = prob.A \ prob.b
+    elseif alg isa LUFactorization
+        u = lu(prob.A) \ prob.b
+    elseif alg isa QRFactorization
+        u = qr(prob.A) \ prob.b
+    elseif alg isa CholeskyFactorization
+        u = cholesky(prob.A) \ prob.b
+    elseif alg isa NormalCholeskyFactorization
+        u = cholesky(Symmetric(prob.A' * prob.A)) \ (prob.A' * prob.b)
+    elseif alg isa SVDFactorization
+        u = svd(prob.A) \ prob.b
+    else
+        # Slower Path but handles all cases
+        cache = init(prob, alg, args...; kwargs...)
+        return solve!(cache)
+    end
+    return SciMLBase.build_linear_solution(alg, u, nothing, prob)
+end

src/factorization.jl

@@ -975,11 +975,17 @@ function init_cacheval(alg::NormalCholeskyFactorization,
     ArrayInterface.cholesky_instance(convert(AbstractMatrix, A))
 end
 
+function init_cacheval(alg::NormalCholeskyFactorization, A::SMatrix, b, u, Pl, Pr,
+    maxiters::Int, abstol, reltol, verbose::Bool,
+    assumptions::OperatorAssumptions)
+    return cholesky(Symmetric((A)' * A))
+end
+
 function init_cacheval(alg::NormalCholeskyFactorization, A, b, u, Pl, Pr,
     maxiters::Int, abstol, reltol, verbose::Bool,
     assumptions::OperatorAssumptions)
     A_ = convert(AbstractMatrix, A)
-    ArrayInterface.cholesky_instance(Symmetric((A)' * A, :L), alg.pivot)
+    return ArrayInterface.cholesky_instance(Symmetric((A)' * A), alg.pivot)
 end
 
 function init_cacheval(alg::NormalCholeskyFactorization,
@@ -993,17 +999,20 @@ function SciMLBase.solve!(cache::LinearCache, alg::NormalCholeskyFactorization;
     A = cache.A
     A = convert(AbstractMatrix, A)
     if cache.isfresh
-        if A isa SparseMatrixCSC || A isa GPUArraysCore.AbstractGPUArray
-            fact = cholesky(Symmetric((A)' * A, :L); check = false)
+        if A isa SparseMatrixCSC || A isa GPUArraysCore.AbstractGPUArray || A isa SMatrix
+            fact = cholesky(Symmetric((A)' * A); check = false)
         else
-            fact = cholesky(Symmetric((A)' * A, :L), alg.pivot; check = false)
+            fact = cholesky(Symmetric((A)' * A), alg.pivot; check = false)
         end
         cache.cacheval = fact
         cache.isfresh = false
     end
     if A isa SparseMatrixCSC
         cache.u .= @get_cacheval(cache, :NormalCholeskyFactorization) \ (A' * cache.b)
         y = cache.u
+    elseif A isa StaticArray
+        cache.u = @get_cacheval(cache, :NormalCholeskyFactorization) \ (A' * cache.b)
+        y = cache.u
     else
         y = ldiv!(cache.u, @get_cacheval(cache, :NormalCholeskyFactorization), A' * cache.b)
     end

test/static_arrays.jl

@@ -1,13 +1,30 @@
 using LinearSolve, StaticArrays, LinearAlgebra, Test
+using AllocCheck
 
 A = SMatrix{5, 5}(Hermitian(rand(5, 5) + I))
 b = SVector{5}(rand(5))
 
+@check_allocs __solve_no_alloc(A, b, alg) = solve(LinearProblem(A, b), alg)
+
+function __non_native_static_array_alg(alg)
+    return alg isa SVDFactorization || alg isa KrylovJL
+end
+
 for alg in (nothing, LUFactorization(), SVDFactorization(), CholeskyFactorization(),
-    KrylovJL_GMRES())
+    NormalCholeskyFactorization(), KrylovJL_GMRES())
     sol = solve(LinearProblem(A, b), alg)
     @inferred solve(LinearProblem(A, b), alg)
     @test norm(A * sol .- b) < 1e-10
+
+    if __non_native_static_array_alg(alg)
+        @test_broken __solve_no_alloc(A, b, alg)
+    else
+        @test_nowarn __solve_no_alloc(A, b, alg)
+    end
+
+    cache = init(LinearProblem(A, b), alg)
+    sol = solve!(cache)
+    @test norm(A * sol .- b) < 1e-10
 end
 
 A = SMatrix{7, 5}(rand(7, 5))