Fix dense LU cache prototypes for wrapper matrices

ext/LinearSolveBLISExt.jl

@@ -268,10 +268,9 @@ function LinearSolve.init_cacheval(
         maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool},
         assumptions::OperatorAssumptions
     )
-    # Build the instance from a 0×0 array of `A`'s own type so the cacheval slot
-    # matches the factorization `solve!` stores (built from `A`). This handles
-    # dense CPU arrays that aren't `Base.Array`, e.g. `FixedSizeArray`.
-    return ArrayInterface.lu_instance(similar(A, 0, 0)), Ref{BlasInt}()
+    # Ask `lu_instance` about `A` itself so wrapper-specific dispatch can choose
+    # the factorization container that `solve!` will store.
+    return ArrayInterface.lu_instance(A), Ref{BlasInt}()
 end
 
 function SciMLBase.solve!(

src/appleaccelerate.jl

@@ -300,12 +300,10 @@ function LinearSolve.init_cacheval(
         maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool},
         assumptions::OperatorAssumptions
     )
-    # Build the instance from a 0×0 array of `A`'s own type so the cacheval slot
-    # matches the factorization `solve!` stores (built from `A`). This handles
-    # dense CPU arrays that aren't `Base.Array`, e.g. `FixedSizeArray`.
-    A0 = similar(A, 0, 0)
-    luinst = ArrayInterface.lu_instance(A0)
-    return LU(luinst.factors, similar(A0, Cint, 0), luinst.info), Ref{Cint}()
+    # Ask `lu_instance` about `A` itself so wrapper-specific dispatch can choose
+    # the factorization container that `solve!` will store.
+    luinst = ArrayInterface.lu_instance(A)
+    return LU(luinst.factors, similar(luinst.ipiv, Cint, 0), luinst.info), Ref{Cint}()
 end
 
 function SciMLBase.solve!(

src/mkl.jl

@@ -291,10 +291,9 @@ function LinearSolve.init_cacheval(
         maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool},
         assumptions::OperatorAssumptions
     )
-    # Build the instance from a 0×0 array of `A`'s own type so the cacheval slot
-    # matches the factorization `solve!` stores (built from `A`). This handles
-    # dense CPU arrays that aren't `Base.Array`, e.g. `FixedSizeArray`.
-    return ArrayInterface.lu_instance(similar(A, 0, 0)), Ref{BlasInt}()
+    # Ask `lu_instance` about `A` itself so wrapper-specific dispatch can choose
+    # the factorization container that `solve!` will store.
+    return ArrayInterface.lu_instance(A), Ref{BlasInt}()
 end
 
 function SciMLBase.solve!(

src/openblas.jl

@@ -312,10 +312,9 @@ function LinearSolve.init_cacheval(
         maxiters::Int, abstol, reltol, verbose::Union{LinearVerbosity, Bool},
         assumptions::OperatorAssumptions
     )
-    # Build the instance from a 0×0 array of `A`'s own type so the cacheval slot
-    # matches the factorization `solve!` stores (built from `A`). This handles
-    # dense CPU arrays that aren't `Base.Array`, e.g. `FixedSizeArray`.
-    return ArrayInterface.lu_instance(similar(A, 0, 0)), Ref{BlasInt}()
+    # Ask `lu_instance` about `A` itself so wrapper-specific dispatch can choose
+    # the factorization container that `solve!` will store.
+    return ArrayInterface.lu_instance(A), Ref{BlasInt}()
 end
 
 function SciMLBase.solve!(

test/Core/componentarrays.jl

@@ -0,0 +1,63 @@
+using LinearSolve, ComponentArrays, LinearAlgebra, Test
+using LinearSolve: OperatorAssumptions
+
+@testset "BLAS LU cache tracks ComponentMatrix factors" begin
+    u0 = ComponentArray(a = ones(4, 2), b = ones(4, 2))
+    ax = only(ComponentArrays.getaxes(u0))
+    A = ComponentArray(Matrix{Float64}(I, length(u0), length(u0)), (ax, ax))
+    b = copy(u0)
+    assump = OperatorAssumptions(true)
+
+    @test A isa DenseMatrix
+    @test !(A isa Matrix)
+    @test similar(A, 0, 0) isa Matrix
+    @test lu!(copy(A)).factors isa ComponentMatrix
+
+    for alg in (MKLLUFactorization(), OpenBLASLUFactorization())
+        slot = LinearSolve.init_cacheval(
+            alg,
+            A,
+            b,
+            b,
+            nothing,
+            nothing,
+            0,
+            0.0,
+            0.0,
+            true,
+            assump,
+        )
+        @test slot[1].factors isa ComponentMatrix
+    end
+
+    slot = LinearSolve.init_cacheval(
+        AppleAccelerateLUFactorization(),
+        A,
+        b,
+        b,
+        nothing,
+        nothing,
+        0,
+        0.0,
+        0.0,
+        true,
+        assump,
+    )
+    @test slot[1].factors isa ComponentMatrix
+    @test slot[1].ipiv isa Vector{Cint}
+end
+
+@testset "dense ComponentMatrix default solve" begin
+    u0 = ComponentArray(a = ones(4, 2), b = ones(4, 2))
+    ax = only(ComponentArrays.getaxes(u0))
+    n = length(u0)
+    M = Matrix{Float64}(I, n, n) .+ 0.01 .* reshape(1:(n * n), n, n)
+    A = ComponentArray(M, (ax, ax))
+    b = copy(u0)
+
+    sol = solve(LinearProblem(A, b))
+    @test norm(M * Array(sol.u) - Array(b)) < 1.0e-10
+
+    sol = solve(LinearProblem(A, b), OpenBLASLUFactorization())
+    @test norm(M * Array(sol.u) - Array(b)) < 1.0e-10
+end

test/runtests.jl

@@ -71,6 +71,7 @@ else
             @time @safetestset "Nonstructural Zeros" include("Core/nonstructural_zeros.jl")
             @time @safetestset "Default Alg Tests" include("Core/default_algs.jl")
             @time @safetestset "FixedSizeArrays" include("Core/fixedsizearrays.jl")
+            @time @safetestset "ComponentArrays" include("Core/componentarrays.jl")
             @time @safetestset "Adjoint Sensitivity" include("Core/adjoint.jl")
             @time @safetestset "ForwardDiff Overloads" include("Core/forwarddiff_overloads.jl")
             @time @safetestset "Traits" include("Core/traits.jl")