From 30d74ad44cc307a98b96219c731e2eb94ef3a44d Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sat, 8 Aug 2026 16:00:33 -0400 Subject: [PATCH 1/2] Replace the vector-RHS BLAS deferral with a rank-4 pure-Julia sweep Vector ldiv! previously ran the matrix kernels' M == 1 scalar remainder up to n = 128 and deferred to LinearAlgebra (BLAS trsv) above. Both lose to a plain column-oriented substitution sweep that LLVM vectorizes: outer-unrolled rank-4, it beats the old kernel path by 1.35-3.4x and trsv by 1.14-7.4x at every measured size (n = 4..2000, both triangles, unit/non-unit, Float32/64, AVX2), so the vector path is now never-BLAS with no swap logic. Inner-product forms cover row-contiguous strided parents. Vector methods now also throw DimensionMismatch for all four wrappers (the Lower ones previously read out of bounds on mismatched sizes). Matrix paths are unchanged; a naive per-column sweep lost to the blocked SIMD kernels in 17 of 18 tiny-size cells, so no matrix swapping. Co-Authored-By: Chris Rackauckas --- Project.toml | 2 +- src/TriangularSolve.jl | 209 ++++++++++++++++++++++++++++++++++------- test/runtests.jl | 87 +++++++++++++++-- 3 files changed, 255 insertions(+), 43 deletions(-) diff --git a/Project.toml b/Project.toml index a8869ee..fa9ea50 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "TriangularSolve" uuid = "d5829a12-d9aa-46ab-831f-fb7c9ab06edf" authors = ["chriselrod and contributors"] -version = "0.2.4" +version = "0.2.5" [deps] CloseOpenIntervals = "fb6a15b2-703c-40df-9091-08a04967cfa9" diff --git a/src/TriangularSolve.jl b/src/TriangularSolve.jl index ce614d2..fe13992 100644 --- a/src/TriangularSolve.jl +++ b/src/TriangularSolve.jl @@ -932,53 +932,192 @@ function ldiv!( return C end -# Vector right-hand sides reuse the matrix kernels: the left-division paths -# already operate on `transpose(A)`, and for a vector `b` the required 1×n -# transposed form is just `transpose(b)` — no reshape, no allocation. A single -# right-hand side runs the kernels' scalar row-remainder, which beats BLAS -# trsv up to the cutoff below (1.2-2.4x measured) but loses to trsv's blocked -# sweep once the triangle falls out of L2, so larger solves keep the -# LinearAlgebra path. -const VECTOR_RHS_CUTOFF = 128 +# Vector right-hand sides never defer to LinearAlgebra/BLAS and never touch +# the SIMD matrix kernels: the pure-Julia sweeps below beat the matrix +# drivers' M == 1 scalar remainder and beat BLAS trsv at every measured size +# (up to 7x below n = 256, 1.1-1.9x at n = 256..2000, AVX2). -for (wrap, dispatch, UNIT) in ( - (:LowerTriangular, :div_dispatch!, false), - (:UnitLowerTriangular, :div_dispatch!, true), - (:UpperTriangular, :div_dispatch_L!, false), - (:UnitUpperTriangular, :div_dispatch_L!, true) +# Column-oriented sweeps for parents whose columns are the contiguous +# direction, outer-unrolled rank-4 so each pass over `x` retires four +# columns: quartering the `x` store traffic measured ~1.5x faster than the +# rank-1 sweep. The unit-diagonal `x[j]` store discipline (elided forward, +# unconditional backward) also follows measurement; flipping either +# direction was up to 1.7x slower. +@inline function _naive_vsolve_fwd!(x, A, ::Val{UNIT}) where {UNIT} + N = length(x) + j = 1 + @inbounds while j < N - 2 + x0 = UNIT ? x[j] : x[j] / A[j, j] + UNIT || (x[j] = x0) + x1 = muladd(-x0, A[j+1, j], x[j+1]) + x1 = UNIT ? x1 : x1 / A[j+1, j+1] + x[j+1] = x1 + x2 = muladd(-x1, A[j+2, j+1], muladd(-x0, A[j+2, j], x[j+2])) + x2 = UNIT ? x2 : x2 / A[j+2, j+2] + x[j+2] = x2 + x3 = muladd( + -x2, + A[j+3, j+2], + muladd(-x1, A[j+3, j+1], muladd(-x0, A[j+3, j], x[j+3])) + ) + x3 = UNIT ? x3 : x3 / A[j+3, j+3] + x[j+3] = x3 + n0 = -x0 + n1 = -x1 + n2 = -x2 + n3 = -x3 + @simd ivdep for i = (j+4):N + x[i] = muladd( + n0, + A[i, j], + muladd( + n1, + A[i, j+1], + muladd(n2, A[i, j+2], muladd(n3, A[i, j+3], x[i])) + ) + ) + end + j += 4 + end + @inbounds while j <= N + xj = UNIT ? x[j] : x[j] / A[j, j] + UNIT || (x[j] = xj) + nxj = -xj + @simd ivdep for i = (j+1):N + x[i] = muladd(nxj, A[i, j], x[i]) + end + j += 1 + end + nothing +end +@inline function _naive_vsolve_bwd!(x, A, ::Val{UNIT}) where {UNIT} + N = length(x) + j = N + @inbounds while j > 3 + x0 = UNIT ? x[j] : x[j] / A[j, j] + x[j] = x0 + x1 = muladd(-x0, A[j-1, j], x[j-1]) + x1 = UNIT ? x1 : x1 / A[j-1, j-1] + x[j-1] = x1 + x2 = muladd(-x1, A[j-2, j-1], muladd(-x0, A[j-2, j], x[j-2])) + x2 = UNIT ? x2 : x2 / A[j-2, j-2] + x[j-2] = x2 + x3 = muladd( + -x2, + A[j-3, j-2], + muladd(-x1, A[j-3, j-1], muladd(-x0, A[j-3, j], x[j-3])) + ) + x3 = UNIT ? x3 : x3 / A[j-3, j-3] + x[j-3] = x3 + n0 = -x0 + n1 = -x1 + n2 = -x2 + n3 = -x3 + @simd ivdep for i = 1:(j-4) + x[i] = muladd( + n0, + A[i, j], + muladd( + n1, + A[i, j-1], + muladd(n2, A[i, j-2], muladd(n3, A[i, j-3], x[i])) + ) + ) + end + j -= 4 + end + @inbounds while j >= 1 + xj = UNIT ? x[j] : x[j] / A[j, j] + x[j] = xj + nxj = -xj + @simd ivdep for i = 1:(j-1) + x[i] = muladd(nxj, A[i, j], x[i]) + end + j -= 1 + end + nothing +end +# Inner-product forms for parents whose rows are the contiguous direction. +@inline function _naive_vsolve_fwd_dot!(x, A, ::Val{UNIT}) where {UNIT} + N = length(x) + @inbounds for i = 1:N + s = zero(eltype(x)) + @simd for j = 1:(i-1) + s = muladd(A[i, j], x[j], s) + end + xi = x[i] - s + x[i] = UNIT ? xi : xi / A[i, i] + end + nothing +end +@inline function _naive_vsolve_bwd_dot!(x, A, ::Val{UNIT}) where {UNIT} + N = length(x) + @inbounds for i = N:-1:1 + s = zero(eltype(x)) + @simd for j = (i+1):N + s = muladd(A[i, j], x[j], s) + end + xi = x[i] - s + x[i] = UNIT ? xi : xi / A[i, i] + end + nothing +end +@inline function _naive_vsolve!(x, A, ::Val{UNIT}, ::Val{UP}) where {UNIT,UP} + colmajor = abs(stride(A, 1)) <= abs(stride(A, 2)) + if UP + colmajor ? _naive_vsolve_bwd!(x, A, Val(UNIT)) : + _naive_vsolve_bwd_dot!(x, A, Val(UNIT)) + else + colmajor ? _naive_vsolve_fwd!(x, A, Val(UNIT)) : + _naive_vsolve_fwd_dot!(x, A, Val(UNIT)) + end +end + +for (wrap, UNIT, UP) in ( + (:LowerTriangular, false, false), + (:UnitLowerTriangular, true, false), + (:UpperTriangular, false, true), + (:UnitUpperTriangular, true, true) ) @eval begin function ldiv!( U::$wrap{T,<:StridedMatrix{T}}, b::StridedVector{T}, - ::Val{thread} = Val(true) - ) where {T<:Union{Float32,Float64},thread} - length(b) > VECTOR_RHS_CUTOFF && return LinearAlgebra.ldiv!(U, b) - nt = thread ? _nthreads() : static(1) - $dispatch( - transpose(b), - transpose(b), - transpose(parent(U)), - nt, - Val($UNIT) - ) + ::Val = Val(true) + ) where {T<:Union{Float32,Float64}} + P = parent(U) + N = length(b) + if size(P, 1) != N + throw( + DimensionMismatch( + "triangular matrix is $(size(P,1))×$(size(P,2)), right-hand side has length $N" + ) + ) + end + _naive_vsolve!(b, P, Val($UNIT), Val($UP)) return b end function ldiv!( c::StridedVector{T}, U::$wrap{T,<:StridedMatrix{T}}, b::StridedVector{T}, - ::Val{thread} = Val(true) - ) where {T<:Union{Float32,Float64},thread} - length(b) > VECTOR_RHS_CUTOFF && return LinearAlgebra.ldiv!(c, U, b) - nt = thread ? _nthreads() : static(1) - $dispatch( - transpose(c), - transpose(b), - transpose(parent(U)), - nt, - Val($UNIT) - ) + ::Val = Val(true) + ) where {T<:Union{Float32,Float64}} + P = parent(U) + N = length(b) + if size(P, 1) != N + throw( + DimensionMismatch( + "triangular matrix is $(size(P,1))×$(size(P,2)), right-hand side has length $N" + ) + ) + elseif length(c) != N + throw( + DimensionMismatch("destination has length $(length(c)), needs $N") + ) + end + c === b || copyto!(c, b) + _naive_vsolve!(c, P, Val($UNIT), Val($UP)) return c end end diff --git a/test/runtests.jl b/test/runtests.jl index 0a492ba..d87ce5f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -92,6 +92,13 @@ function test_packed_lu(::Type{T}) where {T} Y = TriangularSolve.ldiv!(UnitLowerTriangular(F), copy(B), thread) @test Y ≈ Matrix(UnitLowerTriangular(F)) \ B end + b = rand(T, n) + for thread ∈ (Val(false), Val(true)) + xv = TriangularSolve.ldiv!(UpperTriangular(F), copy(b), thread) + @test xv ≈ Matrix(UpperTriangular(F)) \ b + yv = TriangularSolve.ldiv!(UnitLowerTriangular(F), copy(b), thread) + @test yv ≈ Matrix(UnitLowerTriangular(F)) \ b + end C = rand(T, nrhs, n) @test TriangularSolve.rdiv!(copy(C), UnitLowerTriangular(F)) ≈ C / Matrix(UnitLowerTriangular(F)) @@ -134,12 +141,15 @@ end end end @testset "vector right-hand sides" begin - # sizes straddling VECTOR_RHS_CUTOFF cover both the kernel path and the - # LinearAlgebra deferral + # every size runs the naive sweeps (the BLAS deferral is gone); sizes + # cover SIMD-width remainders and the formerly-deferred n > 128 range for T ∈ (Float64, Float32), - n ∈ (1, 2, 5, 8, 16, 33, 64, 127, 128, 129, 200, 500) + n ∈ (1, 2, 5, 8, 16, 33, 64, 127, 128, 129, 200, 500, 1201) P = rand(T, n, n) + T(n) * I + # unit solves ignore the diagonal, so scale the strict triangle to + # keep solution growth bounded at large n + Pu = rand(T, n, n) ./ T(2n) + I b = rand(T, n) for wrap ∈ ( UpperTriangular, @@ -149,11 +159,14 @@ end ), thread ∈ (Val(false), Val(true)) - U = wrap(P) + unit = wrap === UnitUpperTriangular || wrap === UnitLowerTriangular + U = wrap(unit ? Pu : P) x = TriangularSolve.ldiv!(U, copy(b), thread) @test x ≈ Matrix(U) \ b rtol = sqrt(eps(T)) * n c = similar(b) @test TriangularSolve.ldiv!(c, U, copy(b), thread) ≈ x + xa = copy(b) + @test TriangularSolve.ldiv!(xa, U, xa, thread) ≈ x end end F = lu!(rand(100, 100) + 100I).factors @@ -161,22 +174,75 @@ end Uf = UpperTriangular(F) @test TriangularSolve.ldiv!(Uf, copy(bf), Val(false)) ≈ Matrix(Uf) \ bf xf = copy(bf) + cf = similar(bf) for thread ∈ (Val(false), Val(true)) TriangularSolve.ldiv!(Uf, xf, thread) xf .= bf @test iszero(@allocated TriangularSolve.ldiv!(Uf, xf, thread)) + TriangularSolve.ldiv!(cf, Uf, xf, thread) + @test iszero(@allocated TriangularSolve.ldiv!(cf, Uf, xf, thread)) + end + end + @testset "non-contiguous strided vector solves" begin + for T ∈ (Float64, Float32) + n = 37 + M = rand(T, 2n, 2n) + T(2n) * I + Pv = @view M[1:2:2n, 1:2:2n] + v = rand(T, 2n) + bv = @view v[1:2:2n] + for wrap ∈ ( + UpperTriangular, + UnitUpperTriangular, + LowerTriangular, + UnitLowerTriangular + ) + U = wrap(Pv) + xref = Matrix(U) \ Vector(bv) + @test TriangularSolve.ldiv!(U, copy(bv)) ≈ xref rtol = sqrt(eps(T)) * n + cv = @view similar(v)[1:2:2n] + @test TriangularSolve.ldiv!(cv, U, bv) ≈ xref rtol = sqrt(eps(T)) * n + end + end + end + @testset "inner-product sweep kernels" begin + # these forms are selected when the parent's rows are the contiguous + # direction; no Base strided type has that layout, so test them directly + for T ∈ (Float64, Float32), n ∈ (1, 7, 40), UNIT ∈ (false, true) + P = rand(T, n, n) + T(n) * I + b = rand(T, n) + lo = UNIT ? UnitLowerTriangular(P) : LowerTriangular(P) + up = UNIT ? UnitUpperTriangular(P) : UpperTriangular(P) + x = copy(b) + TriangularSolve._naive_vsolve_fwd_dot!(x, P, Val(UNIT)) + @test x ≈ Matrix(lo) \ b rtol = sqrt(eps(T)) * n + x = copy(b) + TriangularSolve._naive_vsolve_bwd_dot!(x, P, Val(UNIT)) + @test x ≈ Matrix(up) \ b rtol = sqrt(eps(T)) * n end end @testset "dimension mismatch throws" begin @test_throws DimensionMismatch TriangularSolve.rdiv!( - rand(4, 8), LowerTriangular(rand(6, 6) + 6I), Val(false) + rand(4, 8), + LowerTriangular(rand(6, 6) + 6I), + Val(false) ) @test_throws DimensionMismatch TriangularSolve.ldiv!( - UpperTriangular(rand(6, 6) + 6I), rand(8, 3), Val(false) + UpperTriangular(rand(6, 6) + 6I), + rand(8, 3), + Val(false) ) @test_throws DimensionMismatch TriangularSolve.ldiv!( - rand(4, 3), UpperTriangular(rand(6, 6) + 6I), rand(6, 3), Val(false) + rand(4, 3), + UpperTriangular(rand(6, 6) + 6I), + rand(6, 3), + Val(false) ) + for wrap ∈ (LowerTriangular, UnitUpperTriangular) + U = wrap(rand(6, 6) + 6I) + @test_throws DimensionMismatch TriangularSolve.ldiv!(U, rand(8)) + @test_throws DimensionMismatch TriangularSolve.ldiv!(rand(6), U, rand(8)) + @test_throws DimensionMismatch TriangularSolve.ldiv!(rand(4), U, rand(6)) + end end @testset "non-strided inputs keep the LinearAlgebra fallback" begin # a Bidiagonal parent is an AbstractMatrix{Float64} but not strided; it @@ -187,6 +253,13 @@ end B = rand(8, 4) @test TriangularSolve.ldiv!(UpperTriangular(Bd), copy(B)) ≈ Matrix(UpperTriangular(Bd)) \ B + bd = rand(8) + @test TriangularSolve.ldiv!(UpperTriangular(Bd), copy(bd)) ≈ + Matrix(UpperTriangular(Bd)) \ bd + Pb = big.(rand(8, 8)) + 8I + bb = big.(rand(8)) + @test TriangularSolve.ldiv!(LowerTriangular(Pb), copy(bb)) ≈ + Matrix(LowerTriangular(Pb)) \ bb end @testset "allocations" begin n = 200 From 4e4f8f6eba1f8e21c90dc32754ce9e619fbf7bc9 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sat, 8 Aug 2026 16:01:34 -0400 Subject: [PATCH 2/2] Update README for the native vector right-hand-side path Co-Authored-By: Chris Rackauckas --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 8872b3a..880de01 100644 --- a/README.md +++ b/README.md @@ -17,8 +17,9 @@ unit-diagonal variants of each): Three-argument out-of-place forms (`rdiv!(C, A, U)`, `ldiv!(C, U, A)`) are also provided, and each method takes a trailing `Val(true)`/`Val(false)` argument -controlling multithreading. Other element types and vector right-hand sides -fall back to `LinearAlgebra`. For example: +controlling multithreading. Vector right-hand sides run dedicated pure-Julia +substitution sweeps (never BLAS); other element types fall back to +`LinearAlgebra`. For example: ```julia julia> using TriangularSolve, LinearAlgebra, MKL;