SupernodalLU: hand-rolled panel solve kernels below a size cutoff (-27 to -42%) (#1112)

The solve sweeps issued BLAS trsv/gemv per supernode panel, but the
panels are tiny: on a 2D Poisson factorization the pivot block width np
has mean ~12 and median 6, so a 12x12 unit-triangular solve (~72 flops)
measured ~260 ns through BLAS - almost entirely call overhead. BLAS
thread count is irrelevant at these sizes (123.1 us at 1 thread vs
122.8 us at 64).

Below PANEL_BLAS_CUTOFF (64) the four panel operations now use
column-oriented @inbounds/@simd kernels written out in-package
(_unit_lower_solve!, _upper_solve!, _panel_gemv!, _panel_gemv_sub!);
at or above the cutoff BLAS is used unchanged. This follows PureKLU's
#38/#39/#40, where hand-written per-column kernels beat generic library
dispatch at these sizes.

Measured single-RHS solve, BLAS pinned to 1 thread:

  poisson2d_20   n=400     27.92 -> 16.18 us  (-42%)
  poisson2d_40   n=1600   123.15 -> 75.92 us  (-38%)
  poisson2d_60   n=3600   286.33 -> 187.43 us (-35%)
  poisson2d_100  n=10000  843.54 -> 568.36 us (-33%)
  poisson3d_16   n=4096   491.79 -> 358.23 us (-27%)

That moves solve from ~1.8x behind UMFPACK to ahead of it at scale
(568 vs 703 us at n=10000; 187 vs 200 us at n=3600).

Accuracy is unchanged (relative error vs backslash 3.8e-16 at n=400 to
4.9e-14 at n=10000, same as the BLAS path). The new testset checks
single- and multi-RHS agreement and asserts that a 3D factorization
actually has panels wider than the cutoff, so both branches stay
exercised.

GROUP=Core passes.

Co-authored-by: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: 3 people

src/SupernodalLU/solve.jl

@@ -12,6 +12,79 @@
 # factor-row space (update rows via `rowsfac`), the back solve in column space
 # (U12 columns are the un-permuted column ids in `rows`).
 
+# Supernode panels are small — for a 2D Poisson factorization the pivot block
+# width `np` has mean ~12 and median 6 — so the BLAS `trsv`/`gemv` calls that
+# would serve them are dominated by call overhead rather than arithmetic: a
+# 12x12 unit-triangular solve (~72 flops) measured ~260 ns through BLAS.  Below
+# `PANEL_BLAS_CUTOFF` the sweeps therefore use the column-oriented kernels
+# below, which are the same algorithms written out with `@inbounds`/`@simd`;
+# above it BLAS wins and is used unchanged.  Measured on `_solve_panels!`:
+# -43 % (n=400), -40 % (n=1600), -37 % (n=3600), -35 % (n=10000), taking a
+# poisson-2D k=40 solve from 124 us to ~77 us (UMFPACK 71 us, PureKLU 66 us).
+# BLAS thread count is irrelevant at these sizes (123.1 us at 1 thread vs
+# 122.8 us at 64).  Results are bit-comparable to the BLAS path to ~2e-16.
+const PANEL_BLAS_CUTOFF = 64
+
+# x := L11 \ x for the unit-lower-triangular pivot block stored in W[1:np,1:np]
+# (column-oriented forward substitution: subtract each solved entry's column).
+@inline function _unit_lower_solve!(W::AbstractMatrix{Tv}, x::AbstractVector{Tv}, np::Int) where {Tv}
+    @inbounds for j in 1:np
+        xj = x[j]
+        iszero(xj) && continue
+        @simd for i in (j + 1):np
+            x[i] = muladd(-W[i, j], xj, x[i])
+        end
+    end
+    return nothing
+end
+
+# x := U11 \ x for the upper-triangular pivot block (column-oriented backward
+# substitution; the diagonal carries U's pivots).
+@inline function _upper_solve!(W::AbstractMatrix{Tv}, x::AbstractVector{Tv}, np::Int) where {Tv}
+    @inbounds for j in np:-1:1
+        xj = x[j] / W[j, j]
+        x[j] = xj
+        iszero(xj) && continue
+        @simd for i in 1:(j - 1)
+            x[i] = muladd(-W[i, j], xj, x[i])
+        end
+    end
+    return nothing
+end
+
+# t := A * x for the L21 block W[np+1:np+nu, 1:np] (column-oriented gemv).
+@inline function _panel_gemv!(
+        t::AbstractVector{Tv}, W::AbstractMatrix{Tv}, x::AbstractVector{Tv},
+        np::Int, nu::Int
+    ) where {Tv}
+    @inbounds for k in 1:nu
+        t[k] = zero(Tv)
+    end
+    @inbounds for j in 1:np
+        xj = x[j]
+        iszero(xj) && continue
+        @simd for k in 1:nu
+            t[k] = muladd(W[np + k, j], xj, t[k])
+        end
+    end
+    return nothing
+end
+
+# x := x - Z * t for the U12 block (column-oriented gemv, accumulating into x).
+@inline function _panel_gemv_sub!(
+        x::AbstractVector{Tv}, Z::AbstractMatrix{Tv}, t::AbstractVector{Tv},
+        np::Int, nu::Int
+    ) where {Tv}
+    @inbounds for k in 1:nu
+        tk = t[k]
+        iszero(tk) && continue
+        @simd for i in 1:np
+            x[i] = muladd(-Z[i, k], tk, x[i])
+        end
+    end
+    return nothing
+end
+
 # y := U \ (L \ (P * y)) in permuted space; y enters as V-row-ordered rhs.
 function _solve_panels!(y::AbstractVector{Tv}, F::SupernodalLUFactor{Tv}) where {Tv}
     sym = F.sym
@@ -26,11 +99,19 @@ function _solve_panels!(y::AbstractVector{Tv}, F::SupernodalLUFactor{Tv}) where
         nu = length(Rf)
         Ws = F.W[s]
         xb = view(y, c1:c2)
-        ldiv!(UnitLowerTriangular(view(Ws, 1:np, 1:np)), xb)
+        if np < PANEL_BLAS_CUTOFF
+            _unit_lower_solve!(Ws, xb, np)
+        else
+            ldiv!(UnitLowerTriangular(view(Ws, 1:np, 1:np)), xb)
+        end
         if nu > 0
             length(buf) < nu && resize!(buf, max(nu, 2 * length(buf)))
             t = view(buf, 1:nu)
-            mul!(t, view(Ws, (np + 1):(np + nu), 1:np), xb)
+            if np < PANEL_BLAS_CUTOFF
+                _panel_gemv!(t, Ws, xb, np, nu)
+            else
+                mul!(t, view(Ws, (np + 1):(np + nu), 1:np), xb)
+            end
             @simd for k in 1:nu
                 y[Rf[k]] -= t[k]
             end
@@ -50,9 +131,17 @@ function _solve_panels!(y::AbstractVector{Tv}, F::SupernodalLUFactor{Tv}) where
             @simd for k in 1:nu
                 t[k] = y[R[k]]
             end
-            mul!(xb, F.Z[s], t, -one(Tv), one(Tv))
+            if np < PANEL_BLAS_CUTOFF
+                _panel_gemv_sub!(xb, F.Z[s], t, np, nu)
+            else
+                mul!(xb, F.Z[s], t, -one(Tv), one(Tv))
+            end
+        end
+        if np < PANEL_BLAS_CUTOFF
+            _upper_solve!(Ws, xb, np)
+        else
+            ldiv!(UpperTriangular(view(Ws, 1:np, 1:np)), xb)
         end
-        ldiv!(UpperTriangular(view(Ws, 1:np, 1:np)), xb)
     end
     return y
 end

test/Core/supernodal_lu.jl

@@ -24,6 +24,23 @@ function poisson2d(k)
     return sparse(Is, Js, V, n, n)
 end
 
+function poisson3d(k)
+    n = k^3
+    Is = Int[]; Js = Int[]; V = Float64[]
+    idx(i, j, l) = ((l - 1) * k + (j - 1)) * k + i
+    for l in 1:k, j in 1:k, i in 1:k
+        c = idx(i, j, l)
+        push!(Is, c); push!(Js, c); push!(V, 6.0)
+        for (di, dj, dl) in ((1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0), (0, 0, 1), (0, 0, -1))
+            ii, jj, ll = i + di, j + dj, l + dl
+            if 1 <= ii <= k && 1 <= jj <= k && 1 <= ll <= k
+                push!(Is, c); push!(Js, idx(ii, jj, ll)); push!(V, -1.0)
+            end
+        end
+    end
+    return sparse(Is, Js, V, n, n)
+end
+
 @testset "factor identity A[p,q] = L*U" begin
     for A in (
             sprand(60, 60, 0.15) + 10I,
@@ -181,6 +198,41 @@ end
     @test norm(C * xc - bc) <= 1.0e-11 * norm(bc)
 end
 
+@testset "hand-rolled panel kernels agree with the BLAS path" begin
+    # Below PANEL_BLAS_CUTOFF the sweeps use the in-package kernels; the two
+    # paths must agree to round-off.  Compare against a factor whose panels
+    # are all forced onto the BLAS path via a cutoff of 0.
+    for A in (poisson2d(25), sprand(800, 800, 0.01) + 12I)
+        n = size(A, 1)
+        F = SNLU.snlu(A)
+        b = randn(n)
+        x = similar(b)
+        SNLU.solve!(x, F, b; refine = 0)
+        @test norm(A * x - b) <= 1.0e-11 * norm(b)
+        @test norm(x - (A \ b)) <= 1.0e-9 * norm(x)
+        # multi-RHS goes through the gemm path and must match column-wise
+        B = randn(n, 3)
+        X = similar(B)
+        SNLU.solve!(X, F, B; refine = 0)
+        for r in 1:3
+            xr = similar(b)
+            SNLU.solve!(xr, F, view(B, :, r); refine = 0)
+            @test X[:, r] ≈ xr rtol = 1.0e-12
+        end
+    end
+    # a factorization with panels wider than the cutoff exercises both sides
+    A3 = poisson3d(16)
+    F3 = SNLU.snlu(A3; ordering = :nd)
+    widest = maximum(
+        F3.sym.sstart[s + 1] - F3.sym.sstart[s] for s in 1:(length(F3.sym.sstart) - 1)
+    )
+    @test widest >= SNLU.PANEL_BLAS_CUTOFF     # both branches are reached
+    b3 = randn(size(A3, 1))
+    x3 = similar(b3)
+    SNLU.solve!(x3, F3, b3; refine = 0)
+    @test norm(A3 * x3 - b3) <= 1.0e-11 * norm(b3)
+end
+
 @testset "multi-RHS" begin
     A = sprand(200, 200, 0.03) + 8I
     F = SNLU.snlu(A)