PERF: Use a plain loop in _pivoting! for small tableaus

benchmark/README.md

@@ -11,6 +11,8 @@ subgroup of `SUITE`. Currently covered:
 
 - [`ddp.jl`](ddp.jl): `DiscreteDP` (`src/markov/ddp.jl`), under
   `SUITE["ddp"]`;
+- [`lcp_lemke.jl`](lcp_lemke.jl): `lcp_lemke` (`src/lcp_lemke.jl`, with the
+  pivoting kernels of `src/pivoting.jl`), under `SUITE["lcp_lemke"]`;
 - [`mc_tools.jl`](mc_tools.jl): `MarkovChain` (`src/markov/mc_tools.jl`),
   under `SUITE["mc_tools"]`.
 
@@ -57,6 +59,17 @@ convergence behavior.
 Note that the suite holds the dense `Q` arrays of the two large cases in
 memory (roughly 0.5 GB in total).
 
+### `lcp_lemke` ([`lcp_lemke.jl`](lcp_lemke.jl))
+
+Lemke's algorithm on random LCPs with positive definite `M` (so that a
+solution exists and the algorithm terminates with it), each generated with
+its own fixed-seed RNG:
+
+| Key | Description |
+|:----|:------------|
+| `dense_n{10,50,200}` | `lcp_lemke` end to end; n = 10 and 50 exercise the loop kernel of `_pivoting!`, n = 200 the BLAS kernel |
+| `dense_n10_prealloc` | `lcp_lemke!` with caller-owned arrays (repeated-solve regime) |
+
 ### `mc_tools` ([`mc_tools.jl`](mc_tools.jl))
 
 The suite is organized by function, over random models (dense stochastic

benchmark/benchmarks.jl

@@ -17,6 +17,7 @@ using BenchmarkTools
 const SUITE = BenchmarkGroup()
 
 SUITE["ddp"] = include("ddp.jl")
+SUITE["lcp_lemke"] = include("lcp_lemke.jl")
 SUITE["mc_tools"] = include("mc_tools.jl")
 
 #= Standalone execution =#

benchmark/lcp_lemke.jl

@@ -0,0 +1,52 @@
+#=
+Benchmarks for lcp_lemke.jl (Lemke's algorithm, with the pivoting kernels
+of pivoting.jl)
+
+The problem sizes are chosen around the loop/BLAS dispatch of `_pivoting!`
+(`PIVOTING_BLAS_CUTOFF`): n = 10 and 50 exercise the loop kernel, n = 200
+the BLAS kernel. The `_prealloc` case times the repeated-solve regime
+through `lcp_lemke!` with caller-owned arrays.
+=#
+using QuantEcon
+using BenchmarkTools
+using Random
+using LinearAlgebra
+
+#= Model generator =#
+
+# Random LCP with positive definite M, so that a solution exists and
+# Lemke's algorithm terminates with it
+function lcp_random_pd(rng, n)
+    A = randn(rng, n, n)
+    M = A'A + I
+    q = randn(rng, n)
+    return M, q
+end
+
+#= Suite =#
+
+suite = BenchmarkGroup()
+
+# A fresh generator per case, so that adding or reordering cases does not
+# alter the model data of the other cases
+new_lcp_rng() = MersenneTwister(1234)
+
+for n in (10, 50, 200)
+    M, q = lcp_random_pd(new_lcp_rng(), n)
+    # the fixture must exercise pivoting, not the trivial all(q .>= 0) path
+    @assert lcp_lemke(M, q).num_iter > 0
+    suite["dense_n$n"] = @benchmarkable lcp_lemke($M, $q)
+end
+
+# Repeated-solve regime: caller-owned output and workspace arrays
+let n = 10
+    M, q = lcp_random_pd(new_lcp_rng(), n)
+    @assert lcp_lemke(M, q).num_iter > 0
+    z = Vector{Float64}(undef, n)
+    tableau = Matrix{Float64}(undef, n, 2n+2)
+    basis = Vector{Int}(undef, n)
+    suite["dense_n10_prealloc"] =
+        @benchmarkable lcp_lemke!($z, $tableau, $basis, $M, $q)
+end
+
+suite

src/pivoting.jl

@@ -8,12 +8,29 @@ using LinearAlgebra: BLAS, BlasFloat
 const TOL_PIV = 1e-10
 const TOL_RATIO_DIFF = 1e-15
 
+# For tableaus with at most this many elements, `_pivoting!` uses a plain
+# SIMD loop instead of BLAS. A rank-1 update has no data reuse for BLAS
+# kernels to exploit, so the loop matches single-threaded BLAS throughput,
+# while below this size BLAS multithreading cannot amortize its
+# synchronization latency (and the loop also avoids the call overhead).
+# The total tableau size is what the crossover tracks (the memory traffic
+# of the update), across square and rectangular shapes alike; it is
+# machine dependent, but only through a square root of the hardware
+# constants, and sits at 16k-19k elements on the machines measured (square
+# and rectangular shapes on laptop and server hardware). 2^14 errs on the
+# safe side, as using the loop slightly below the true crossover costs a
+# few percent where using BLAS above it costs nothing.
+const PIVOTING_BLAS_CUTOFF = 2^14
 
 """
     _pivoting!(tableau, pivot_col, pivot_row, col_buf)
 
 Perform a pivoting step. Modify `tableau` in place.
 
+For a strided tableau of BLAS-compatible eltype with more than
+`PIVOTING_BLAS_CUTOFF` elements the update is delegated to BLAS, and is
+performed by a plain loop otherwise; the two paths agree up to roundoff.
+
 # Arguments
 
 - `tableau::AbstractMatrix`: Array containing the tableau.
@@ -30,15 +47,27 @@ Perform a pivoting step. Modify `tableau` in place.
 function _pivoting!(tableau::AbstractMatrix{T},
                     pivot_col::Integer, pivot_row::Integer,
                     col_buf::Vector{T}) where {T<:AbstractFloat}
+    if tableau isa StridedMatrix{<:BlasFloat} &&
+            length(tableau) > PIVOTING_BLAS_CUTOFF
+        _pivoting_blas!(tableau, pivot_col, pivot_row, col_buf)
+    else
+        _pivoting_loop!(tableau, pivot_col, pivot_row, col_buf)
+    end
+    return tableau
+end
+
+function _pivoting_blas!(tableau::StridedMatrix{T},
+                         pivot_col::Integer, pivot_row::Integer,
+                         col_buf::Vector{T}) where {T<:BlasFloat}
     @inbounds @views begin
         pivot_elt = tableau[pivot_row, pivot_col]
-        _div!(tableau[pivot_row, :], pivot_elt)
+        BLAS.scal!(inv(pivot_elt), tableau[pivot_row, :])
 
         x = col_buf
         copyto!(x, tableau[:, pivot_col])
         x[pivot_row] = zero(T)
         y = tableau[pivot_row, :]
-        _rank1_update!(x, y, tableau)  # tableau -= x * y'
+        BLAS.ger!(-one(T), x, y, tableau)  # tableau -= x * y'
 
         # Eliminate possible floating-point residue
         for i in 1:size(tableau, 1)
@@ -50,16 +79,48 @@ function _pivoting!(tableau::AbstractMatrix{T},
     return tableau
 end
 
-_div!(x::StridedVector{T}, c::T) where {T<:BlasFloat} = BLAS.scal!(inv(c), x)
+# Single fused pass over the tableau in column-major order; the pivot-row
+# entry of each column is normalized and written last, and x[pivot_row] = 0
+# keeps the update from touching the pivot row
+function _pivoting_loop!(tableau::AbstractMatrix{T},
+                         pivot_col::Integer, pivot_row::Integer,
+                         col_buf::Vector{T}) where {T<:AbstractFloat}
+    nrows, ncols = size(tableau)
+    @inbounds begin
+        # For BLAS eltypes, multiply by the reciprocal, matching the BLAS
+        # kernel (safe: any pivot passing the tolerance checks is far above
+        # the inv overflow threshold); for narrower eltypes (e.g. Float16)
+        # inv can overflow where direct division stays finite, so divide.
+        # `use_inv` is a compile-time constant, so the branch specializes
+        # away.
+        use_inv = T <: BlasFloat
+        pivot_elt = tableau[pivot_row, pivot_col]
+        inv_pivot_elt = use_inv ? inv(pivot_elt) : pivot_elt
+
+        x = col_buf
+        for i in 1:nrows
+            x[i] = tableau[i, pivot_col]
+        end
+        x[pivot_row] = zero(T)
 
-# Fallback
-_div!(x, c) = rdiv!(x, c)
+        for j in 1:ncols
+            y_j = use_inv ? tableau[pivot_row, j] * inv_pivot_elt :
+                            tableau[pivot_row, j] / pivot_elt
+            @simd for i in 1:nrows
+                tableau[i, j] -= x[i] * y_j
+            end
+            tableau[pivot_row, j] = y_j
+        end
 
-_rank1_update!(x, y, A::StridedMatrix{T}) where {T<:BlasFloat} =
-    BLAS.ger!(-one(T), x, y, A)
+        # Eliminate possible floating-point residue
+        for i in 1:nrows
+            tableau[i, pivot_col] = zero(T)
+        end
+        tableau[pivot_row, pivot_col] = one(T)
+    end
 
-# Fallback
-_rank1_update!(x, y, A) = A .-= x .* y'
+    return tableau
+end
 
 
 """

test/test_pivoting.jl

@@ -39,4 +39,42 @@ using QuantEcon: _pivoting!, _lex_min_ratio_test!
             @test isapprox(tableau, tableau_opt)
         end
     end
+
+    @testset "Loop and BLAS kernels agree" begin
+        rng = MersenneTwister(0)
+        pivcol, pivrow = 3, 2
+        cutoff = QuantEcon.PIVOTING_BLAS_CUTOFF
+        # small, and rectangular shapes exactly at and just above the
+        # dispatch boundary in total tableau size
+        for (nrows, ncols) in ((10, 22), (64, cutoff ÷ 64),
+                               (64, cutoff ÷ 64 + 1))
+            tableau_0 = rand(rng, nrows, ncols) .+ 0.5
+            col_buf = Vector{Float64}(undef, nrows)
+
+            tableau_loop = copy(tableau_0)
+            QuantEcon._pivoting_loop!(tableau_loop, pivcol, pivrow, col_buf)
+            tableau_blas = copy(tableau_0)
+            QuantEcon._pivoting_blas!(tableau_blas, pivcol, pivrow, col_buf)
+            @test tableau_loop ≈ tableau_blas rtol=1e-13
+
+            tableau = copy(tableau_0)
+            @inferred _pivoting!(tableau, pivcol, pivrow, col_buf)
+            @test tableau ==
+                (nrows * ncols > cutoff ? tableau_blas : tableau_loop)
+            # the pivot column is reduced to the unit vector exactly
+            @test tableau[:, pivcol] ==
+                [i == pivrow ? 1. : 0. for i in 1:nrows]
+        end
+    end
+
+    @testset "Non-BLAS eltype normalization stays finite" begin
+        # inv(p) overflows Float16, so the loop kernel must divide
+        # directly for non-BLAS eltypes
+        p = Float16(1e-5)
+        tableau = Float16[p p 0; p 2p p]
+        col_buf = Vector{Float16}(undef, 2)
+        _pivoting!(tableau, 1, 1, col_buf)
+        @test all(isfinite, tableau)
+        @test tableau == Float16[1 1 0; 0 p p]
+    end
 end