Replace the vector-RHS BLAS deferral with a rank-4 pure-Julia sweep (never-BLAS vectors)

Project.toml

@@ -1,7 +1,7 @@
 name = "TriangularSolve"
 uuid = "d5829a12-d9aa-46ab-831f-fb7c9ab06edf"
 authors = ["chriselrod <elrodc@gmail.com> and contributors"]
-version = "0.2.4"
+version = "0.2.5"
 
 [deps]
 CloseOpenIntervals = "fb6a15b2-703c-40df-9091-08a04967cfa9"

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;
 

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

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,34 +159,90 @@ 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
     bf = rand(100)
     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