Batch real matrix-RHS solves

src/Kernel.jl

@@ -202,6 +202,9 @@ mutable struct KLUNumeric{Tv, Ti <: Integer, Tr <: Real}
     Offx::Vector{Tv}
     nzoff::Ti
     Xwork::Vector{Tv}
+    # Lazy lane-major matrix-RHS workspace; private to this numeric factor.
+    # Reused across solves/refactors; every active lane is gathered before reading.
+    Xrhs::Matrix{Tv}
     # Allocated only when `common.scale > 0`; empty otherwise -- the
     # `Rs[k] = Rs[Pnum[k]]` shuffle at the tail of factor/refactor is
     # skipped entirely in the unscaled case.
@@ -228,7 +231,7 @@ KLUNumeric{Tv, Ti, Tr}() where {Tv, Ti <: Integer, Tr <: Real} = KLUNumeric{Tv,
     Tv[], Tr[],
     Ti[], Ti[], Tv[],
     Ti(0),
-    Tv[], Tr[],
+    Tv[], Matrix{Tv}(undef, 0, 0), Tr[],
     KernelWorkspace{Tv, Ti}(0),
     Ti[],
 )
@@ -345,6 +348,7 @@ function _alloc_numeric(
         Vector{Tv}(undef, max(nzoff, 0)),
         Sym.nzoff,
         Xwork,
+        Matrix{Tv}(undef, 0, 0),
         scale > 0 ? Vector{Tr}(undef, n) : Tr[],
         wk,
         Vector{Ti}(undef, maxblock),

src/PureKLU.jl

@@ -595,6 +595,12 @@ the corresponding adjoint and transpose systems.
 - `check = true`: throw for hard KLU errors. Numerical singularity remains a
   status on the factorization.
 
+# Workspace
+Real `Float32`/`Float64` matrix solves with at least eight RHS lazily allocate
+`n * min(64, nrhs)` elements of scratch. The workspace grows for wider solves
+(up to 64 columns), is retained across solves and refactorizations, and is reused
+without allocation while its capacity suffices. Smaller RHS do not shrink it.
+
 # Returns
 - The same array `B` after replacement by the solution.
 

src/Solve.jl

@@ -4,6 +4,10 @@
 # bit-for-bit, set `common.use_fma = false` (or pass `use_fma=false` to
 # `klu`/`klu_factor!`) -- see `Kernel.jl` for the full discussion.
 
+const _RHS_MIN_COLUMNS = 8
+const _RHS_TILE_WIDTH = 64
+const _RHS_TRANSFER_WIDTH = 16
+
 """
     klu_solve!(Sym, Num, B, common; conj_solve=false) -> B
 
@@ -40,6 +44,24 @@ function _klu_solve_impl!(
     size(B, 1) == n || throw(DimensionMismatch())
     nrhs = size(B, 2)
 
+    # Keep small, generic and vector solves on the original scalar path.
+    if Tv <: Union{Float32, Float64} && B isa StridedMatrix{Tv} && nrhs >= _RHS_MIN_COLUMNS
+        width = min(_RHS_TILE_WIDTH, nrhs)
+        if size(Num.Xrhs, 2) != n || size(Num.Xrhs, 1) < width
+            Num.Xrhs = Matrix{Tv}(undef, width, n)
+        end
+        for first_column in 1:width:nrhs
+            lanes = min(width, nrhs - first_column + 1)
+            _klu_solve_tile!(Num.Xrhs, B, first_column, lanes, Sym, Num, fma_val)
+        end
+        if was_singular
+            common.status = KLU_SINGULAR
+            common.numerical_rank = sing_rank
+            common.singular_col = sing_col
+        end
+        return B
+    end
+
     Q = Sym.Q
     R = Sym.R
     nblocks = Int(Sym.nblocks)
@@ -124,6 +146,119 @@ function _klu_solve_impl!(
     return B
 end
 
+# RHS lanes are contiguous in X. Sparse coefficients/indices are loaded once
+# per tile, without changing the arithmetic order within any RHS. Transfers use
+# smaller sub-blocks to limit the number of strided B columns touched at once.
+@inline function _submul_lanes!(X, target, source, coefficient, lanes, fma_val::Val)
+    @inbounds @simd for lane in 1:lanes
+        X[lane, target] = _mulsub(X[lane, target], coefficient, X[lane, source], fma_val)
+    end
+    return nothing
+end
+
+@inline function _divide_lanes!(X, row, diagonal, lanes, fma_val::Val)
+    @inbounds @simd for lane in 1:lanes
+        X[lane, row] = _cdiv(X[lane, row], diagonal, fma_val)
+    end
+    return nothing
+end
+
+function _klu_solve_tile!(
+        X::Matrix{Tv}, B::AbstractMatrix{Tv}, first_column::Int, lanes::Int,
+        Sym, Num, fma_val::Val
+    ) where {Tv}
+    n = Int(Sym.n)
+    nblocks = Int(Sym.nblocks)
+    Pnum, Q = Num.Pnum, Sym.Q
+    R = Sym.R
+    Lip, Llen = Num.Lip, Num.Llen
+    Uip, Ulen = Num.Uip, Num.Ulen
+    Udiag = Num.Udiag
+    Rs = Num.Rs
+    Offp, Offi, Offx = Num.Offp, Num.Offi, Num.Offx
+    scaled = !isempty(Rs)
+
+    @inbounds begin
+        # X = P * (R \ B), with Pnum zero-based as in the upstream kernel.
+        for first_lane in 1:_RHS_TRANSFER_WIDTH:lanes, k in 0:(n - 1)
+            source_row = Int(Pnum[k + 1]) + 1
+            @simd for lane in first_lane:min(first_lane + _RHS_TRANSFER_WIDTH - 1, lanes)
+                X[lane, k + 1] = B[source_row, first_column + lane - 1]
+            end
+        end
+        if scaled
+            for k in 1:n
+                scale = Rs[k]
+                @simd for lane in 1:lanes
+                    X[lane, k] = X[lane, k] / scale
+                end
+            end
+        end
+
+        # Preserve PureKLU's BTF order, block boundaries, pivoted L/U factors, and
+        # off-diagonal block scatter.  Only the RHS lane is moved to the innermost
+        # loop; arithmetic and traversal order within each lane are unchanged.
+        for block in nblocks:-1:1
+            k1 = Int(R[block])
+            k2 = Int(R[block + 1])
+            nk = k2 - k1
+            block_num = Num.LUbx[block]
+            if nk == 1
+                _divide_lanes!(X, k1 + 1, Udiag[k1 + 1], lanes, fma_val)
+            else
+                for local_k in 0:(nk - 1)
+                    row = k1 + local_k + 1
+                    lip = Int(Lip[row])
+                    len = Int(Llen[row])
+                    if len > 0
+                        for p in 0:(len - 1)
+                            index = lip + p + 1
+                            target = k1 + Int(block_num.Li[index]) + 1
+                            _submul_lanes!(X, target, row, block_num.Lx[index], lanes, fma_val)
+                        end
+                    end
+                end
+                for local_k in (nk - 1):-1:0
+                    row = k1 + local_k + 1
+                    _divide_lanes!(X, row, Udiag[row], lanes, fma_val)
+                    uip = Int(Uip[row])
+                    len = Int(Ulen[row])
+                    if len > 0
+                        for p in 0:(len - 1)
+                            index = uip + p + 1
+                            target = k1 + Int(block_num.Ui[index]) + 1
+                            _submul_lanes!(X, target, row, block_num.Ux[index], lanes, fma_val)
+                        end
+                    end
+                end
+            end
+            if block > 1
+                for k in k1:(k2 - 1)
+                    pstart = Int(Offp[k + 1])
+                    pend = Int(Offp[k + 2]) - 1
+                    if pstart <= pend
+                        for p in pstart:pend
+                            index = p + 1
+                            target = Int(Offi[index]) + 1
+                            _submul_lanes!(X, target, k + 1, Offx[index], lanes, fma_val)
+                        end
+                    end
+                end
+            end
+        end
+
+        # Scatter X = Q*x back into the caller's RHS.  The lane loop remains
+        # contiguous in X even when B has a larger column stride.
+        for first_lane in 1:_RHS_TRANSFER_WIDTH:lanes, k in 0:(n - 1)
+            destination_row = Int(Q[k + 1]) + 1
+            @simd for lane in first_lane:min(first_lane + _RHS_TRANSFER_WIDTH - 1, lanes)
+                B[destination_row, first_column + lane - 1] = X[lane, k + 1]
+            end
+        end
+    end
+    return nothing
+end
+
 function _klu_lsolve!(
         n::Int, Lip::AbstractVector{Ti}, Llen::AbstractVector{Ti},
         block::KLUNumericBlock{Tv, Ti}, X::AbstractVector{Tv},

test/test_matrix_rhs.jl

@@ -0,0 +1,116 @@
+using Test
+using PureKLU
+using SparseArrays
+using LinearAlgebra
+using Random
+
+# Compare against repeated vector solves using the very same numeric factors.
+function column_solve!(K, B)
+    for j in axes(B, 2)
+        solve!(K, view(B, :, j))
+    end
+    return B
+end
+
+@testset "Blocked matrix RHS" begin
+    rng = MersenneTwister(714)
+    # Two irreducible blocks plus a singleton and off-diagonal BTF couplings.
+    # Small first diagonal forces pivoting; row magnitudes exercise scaling.
+    A = sparse([0.01 2 3 0 1; 4 5 0 2 0; 0 0 6 1 2; 0 0 2 7 0; 0 0 0 0 8.0])
+    p = [3, 1, 5, 2, 4]
+    q = [2, 5, 1, 4, 3]
+    A = A[p, q]
+    configurations = (
+        (Float64, Int, true, 2), (Float64, Int, false, 0),
+        (Float32, Int32, true, 1), (Float32, Int32, false, 2),
+        (Float64, Int32, false, 1), (Float32, Int, true, 0),
+    )
+    for (Tv, Ti, fma, scale) in configurations
+        boundaries = (Tv, Ti, fma, scale) == first(configurations)
+        At = SparseMatrixCSC{Tv, Ti}(A)
+        K = klu(At; use_fma = fma, full_factor = false, detect_banded = false)
+        K.common.scale = scale
+        klu_factor!(K)
+        @test K.nblocks == 3
+        @test K.nzoff > 0
+        @test isempty(K.numeric.Xrhs)
+        counts = boundaries ? (0, 1, 7, 8, 16, 63, 64, 65, 129, 9, 256, 8) : (65,)
+        for nrhs in counts
+            B = randn(rng, Tv, 5, nrhs)
+            expected = column_solve!(K, copy(B))
+            X = copy(B)
+            solve!(K, X)
+            @test X == expected
+            @test At * X ≈ B
+        end
+        X = ones(Tv, 5, 65)
+        @test solve!(K, X) === X
+        @test (@allocated solve!(K, X)) == 0
+        boundaries || continue
+        scratch = K.numeric.Xrhs
+        K2 = deepcopy(K)
+        @test K2.numeric.Xrhs !== scratch
+        @test klu(At).numeric.Xrhs !== scratch
+        for refactor! in (klu!, (F, A) -> klu_factor!(F))
+            refactor!(K, At)
+            @test K.numeric.Xrhs === scratch
+            B = randn(rng, Tv, 5, 65)
+            @test solve!(K, copy(B)) == column_solve!(K2, copy(B))
+        end
+        changed = 2 * At
+        klu!(K, changed)
+        B = randn(rng, Tv, 5, 65)
+        @test changed * solve!(K, copy(B)) ≈ B
+        @test At * solve!(K2, copy(B)) ≈ B
+
+        storage = randn(rng, Tv, 7, 130)
+        X = view(storage, 2:6, 1:2:130)
+        expected = column_solve!(K, copy(X))
+        @test solve!(K, X) === X
+        @test X == expected
+        @test_throws ArgumentError solve!(K, view(storage, 1:2:7, :))
+        @test_throws DimensionMismatch solve!(K, zeros(Tv, 4, 65))
+    end
+end
+
+@testset "Matrix RHS fallbacks and singular status" begin
+    for Tv in (BigFloat, ComplexF64)
+        A = sparse(Tv[4 1; 2 3])
+        K = klu(A)
+        B = ones(Tv, 2, 65)
+        @test solve!(K, copy(B)) == column_solve!(K, copy(B))
+        @test isempty(K.numeric.Xrhs)
+        @test A' * solve!(K', copy(B)) ≈ B
+        @test transpose(A) * solve!(transpose(K), copy(B)) ≈ B
+    end
+    A = sparse([1.0 0; 0 0])
+    for check in (true, false)
+        K = klu(A; allowsingular = true)
+        rank, col = K.common.numerical_rank, K.common.singular_col
+        B = ones(2, 65)
+        expected = column_solve!(K, copy(B))
+        @test isequal(solve!(K, B; check), expected)
+        @test K.common.status == PureKLU.KLU_SINGULAR
+        @test (K.common.numerical_rank, K.common.singular_col) == (rank, col)
+    end
+end
+
+@testset "Matrix RHS workspace growth" begin
+    N = 40
+    T = spdiagm(-1 => -ones(N - 1), 0 => fill(2.0, N), 1 => -ones(N - 1))
+    E = spdiagm(0 => ones(N))
+    A = kron(E, T) + kron(T, E)
+    K = klu(A)
+    rng = MersenneTwister(715)
+    for (nrhs, capacity) in ((8, 8), (16, 16), (9, 16), (65, 64), (8, 64))
+        scratch = K.numeric.Xrhs
+        B = randn(rng, N^2, nrhs)
+        expected = column_solve!(K, copy(B))
+        @test solve!(K, B) == expected
+        @test size(K.numeric.Xrhs) == (capacity, N^2)
+        if size(scratch, 1) >= min(64, nrhs)
+            @test K.numeric.Xrhs === scratch
+        end
+        @test (@allocated solve!(K, B)) == 0
+    end
+end