handle the getrs as well and support more lapack installations

Author: ChrisRackauckas

ext/LinearSolveMKLExt.jl

@@ -24,7 +24,7 @@ function getrf!(A::AbstractMatrix{<:Float64}; ipiv = similar(A, BlasInt, min(siz
             Ref{BlasInt}, Ptr{BlasInt}, Ptr{BlasInt}),
             m, n, A, lda, ipiv, info)
     chkargsok(info[])
-    A, ipiv, info[] #Error code is stored in LU factorization type
+    A, ipiv, info[], info #Error code is stored in LU factorization type
 end
 
 default_alias_A(::MKLLUFactorization, ::Any, ::Any) = false
@@ -33,7 +33,7 @@ default_alias_b(::MKLLUFactorization, ::Any, ::Any) = false
 function LinearSolve.init_cacheval(alg::MKLLUFactorization, A, b, u, Pl, Pr,
     maxiters::Int, abstol, reltol, verbose::Bool,
     assumptions::OperatorAssumptions)
-    ArrayInterface.lu_instance(convert(AbstractMatrix, A))
+    ArrayInterface.lu_instance(convert(AbstractMatrix, A)), Ref{BlasInt}()
 end
 
 function SciMLBase.solve!(cache::LinearCache, alg::MKLLUFactorization;
@@ -42,11 +42,12 @@ function SciMLBase.solve!(cache::LinearCache, alg::MKLLUFactorization;
     A = convert(AbstractMatrix, A)
     if cache.isfresh
         cacheval = @get_cacheval(cache, :MKLLUFactorization)
-        fact = LU(getrf!(A; ipiv = cacheval.ipiv)...)
+        res = getrf!(A; ipiv = cacheval[1].ipiv, info = cacheval[2])
+        fact = LU(res[1:3]...), res[4]
         cache.cacheval = fact
         cache.isfresh = false
     end
-    y = ldiv!(cache.u, @get_cacheval(cache, :MKLLUFactorization), cache.b)
+    y = ldiv!(cache.u, @get_cacheval(cache, :MKLLUFactorization)[1], cache.b)
     SciMLBase.build_linear_solution(alg, y, nothing, cache)
 end
 

src/appleaccelerate.jl

@@ -41,7 +41,27 @@ function aa_getrf!(A::AbstractMatrix{<:Float64}; ipiv = similar(A, Cint, min(siz
             Ref{Cint}, Ptr{Cint}, Ptr{Cint}),
             m, n, A, lda, ipiv, info)
     info[] < 0 && throw(ArgumentError("Invalid arguments sent to LAPACK dgetrf_"))
-    A, Vector{BlasInt}(ipiv), BlasInt(info[]) #Error code is stored in LU factorization type
+    A, ipiv, BlasInt(info[]), info #Error code is stored in LU factorization type
+end
+
+function aa_getrs!(trans::AbstractChar, A::AbstractMatrix{<:Float64}, ipiv::AbstractVector{Cint}, B::AbstractVecOrMat{<:Float64}; info = Ref{Cint}())
+    require_one_based_indexing(A, ipiv, B)
+    LinearAlgebra.LAPACK.chktrans(trans)
+    chkstride1(A, B, ipiv)
+    n = LinearAlgebra.checksquare(A)
+    if n != size(B, 1)
+        throw(DimensionMismatch("B has leading dimension $(size(B,1)), but needs $n"))
+    end
+    if n != length(ipiv)
+        throw(DimensionMismatch("ipiv has length $(length(ipiv)), but needs to be $n"))
+    end
+    nrhs = size(B, 2)
+    ccall(("dgetrs_", libacc), Cvoid,
+          (Ref{UInt8}, Ref{Cint}, Ref{Cint}, Ptr{Float64}, Ref{Cint},
+           Ptr{Cint}, Ptr{Float64}, Ref{Cint}, Ptr{Cint}, Clong),
+          trans, n, size(B,2), A, max(1,stride(A,2)), ipiv, B, max(1,stride(B,2)), info, 1)
+    LinearAlgebra.LAPACK.chklapackerror(BlasInt(info[]))
+    B
 end
 
 default_alias_A(::AppleAccelerateLUFactorization, ::Any, ::Any) = false
@@ -50,7 +70,8 @@ default_alias_b(::AppleAccelerateLUFactorization, ::Any, ::Any) = false
 function LinearSolve.init_cacheval(alg::AppleAccelerateLUFactorization, A, b, u, Pl, Pr,
     maxiters::Int, abstol, reltol, verbose::Bool,
     assumptions::OperatorAssumptions)
-    ArrayInterface.lu_instance(convert(AbstractMatrix, A))
+    luinst = ArrayInterface.lu_instance(convert(AbstractMatrix, A))
+    LU(luinst.factors,similar(A, Cint, 0), luinst.info), Ref{Cint}()
 end
 
 function SciMLBase.solve!(cache::LinearCache, alg::AppleAccelerateLUFactorization;
@@ -59,10 +80,23 @@ function SciMLBase.solve!(cache::LinearCache, alg::AppleAccelerateLUFactorizatio
     A = convert(AbstractMatrix, A)
     if cache.isfresh
         cacheval = @get_cacheval(cache, :AppleAccelerateLUFactorization)
-        fact = LU(aa_getrf!(A; ipiv = cacheval.ipiv)...)
+        res = aa_getrf!(A; ipiv = cacheval[1].ipiv, info = cacheval[2])
+        fact = LU(res[1:3]...), res[4]
         cache.cacheval = fact
         cache.isfresh = false
     end
-    y = ldiv!(cache.u, @get_cacheval(cache, :AppleAccelerateLUFactorization), cache.b)
-    SciMLBase.build_linear_solution(alg, y, nothing, cache)
+
+    A, info = @get_cacheval(cache, :AppleAccelerateLUFactorization)
+    LinearAlgebra.require_one_based_indexing(cache.u, cache.b)
+    m, n = size(A, 1), size(A, 2)
+    if m > n
+        Bc = copy(cache.b)
+        aa_getrs!('N', A.factors, A.ipiv, Bc; info)
+        return copyto!(cache.u, 1, Bc, 1, n)
+    else
+        copyto!(cache.u, cache.b)
+        aa_getrs!('N', A.factors, A.ipiv, cache.u; info)
+    end
+
+    SciMLBase.build_linear_solution(alg, cache.u, nothing, cache)
 end