diff --git a/src/TriangularSolve.jl b/src/TriangularSolve.jl index 153f940..ba8e3e9 100644 --- a/src/TriangularSolve.jl +++ b/src/TriangularSolve.jl @@ -932,6 +932,58 @@ 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 + +for (wrap, dispatch, UNIT) in ( + (:LowerTriangular, :div_dispatch!, false), + (:UnitLowerTriangular, :div_dispatch!, true), + (:UpperTriangular, :div_dispatch_L!, false), + (:UnitUpperTriangular, :div_dispatch_L!, 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) + ) + 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) + ) + return c + end + end +end + ldiv!(A, B, ::Val = Val(true)) = LinearAlgebra.ldiv!(A, B) ldiv!(Y, A, B, ::Val = Val(true)) = LinearAlgebra.ldiv!(Y, A, B) rdiv!(A, B, ::Val = Val(true)) = LinearAlgebra.rdiv!(A, B) diff --git a/test/runtests.jl b/test/runtests.jl index 0897499..0a492ba 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -133,6 +133,40 @@ end } end end + @testset "vector right-hand sides" begin + # sizes straddling VECTOR_RHS_CUTOFF cover both the kernel path and the + # LinearAlgebra deferral + for T ∈ (Float64, Float32), + n ∈ (1, 2, 5, 8, 16, 33, 64, 127, 128, 129, 200, 500) + + P = rand(T, n, n) + T(n) * I + b = rand(T, n) + for wrap ∈ ( + UpperTriangular, + UnitUpperTriangular, + LowerTriangular, + UnitLowerTriangular + ), + thread ∈ (Val(false), Val(true)) + + U = wrap(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 + end + end + F = lu!(rand(100, 100) + 100I).factors + bf = rand(100) + Uf = UpperTriangular(F) + @test TriangularSolve.ldiv!(Uf, copy(bf), Val(false)) ≈ Matrix(Uf) \ bf + xf = copy(bf) + for thread ∈ (Val(false), Val(true)) + TriangularSolve.ldiv!(Uf, xf, thread) + xf .= bf + @test iszero(@allocated TriangularSolve.ldiv!(Uf, xf, thread)) + end + end @testset "dimension mismatch throws" begin @test_throws DimensionMismatch TriangularSolve.rdiv!( rand(4, 8), LowerTriangular(rand(6, 6) + 6I), Val(false)