Fix type instability in Recipe structs and eliminate hot-loop allocations by dmaldona · Pull Request #185 · numlinalg/RandLinearAlgebra.jl

dmaldona · 2026-08-05T03:31:58Z

Description

Fixes type instability across all Recipe structs, eliminates heap allocations in compressor and solver hot paths, rewrites QR/LQ sub-solvers to use pre-allocated LAPACK workspace, and fixes 3 correctness bugs. Based on profiling results from the July 2026 benchmark audit (REVIEW_JULY7.md).

37 files changed across src/ and test/, touching all three subsystems (Compressors, Solvers, Approximators).

Motivation and Context

Benchmark profiling revealed that the library's Recipe pattern — designed for zero-allocation iteration — was not following through on that promise. Solvers allocated 73 KiB–2.1 MiB per iteration due to:

Abstract field types on Recipe structs forcing dynamic dispatch
qr!() allocating a new QRCompactWY factorization object every call
SRHT/FJLT creating temporary matrices via fancy indexing and sprandn
ColumnProjection using typeof(A)(undef, ...) which breaks on views and allocates

Changes

Type stability (12 Recipe structs parameterized):
GaussianRecipe, CountSketchRecipe, SparseSignRecipe, SamplingRecipe, IdentityRecipe, RandSVDRecipe, RangeFinderRecipe, LUPPRecipe, QRCPRecipe, AgmonRecipe, BasicLoggerRecipe — all abstract field types (Number, AbstractMatrix, Cardinality, CompressorRecipe, Union types) replaced with concrete type parameters.

Sub-solver rewrite (biggest allocation win):

QRSolverRecipe and LQSolverRecipe rewritten to use LAPACK.geqrt3!/gemqrt!/trtrs! with pre-allocated T (block reflector) and work buffers
update_sub_solver! changed from field reassignment (solver.A = A, which triggers implicit conversion copies for SubArray inputs) to copyto!(solver.A, A)
Shared LAPACK helpers moved to SubSolvers.jl

Compressor allocation fixes:

SRHT mul!: pre-allocated extraction buffer + sign vector (was ~1 MiB/call)
FJLT update_compressor!: in-place randn!(nonzeros(S.op)) (was 32 KiB/call)
CountSketch: non-allocating sign generation

Correctness fixes:

LUPP complete_selector: buffer sized with output dims, not compressor dims (crashed on non-square matrices)
RandSVD rapproximate!: Matrix{eltype(A)} instead of hardcoded Matrix{Float64}
FullResidualRecipe/LSGradientRecipe: zeros(eltype(b), ...) to match type parameter
ColumnProjection: similar(A, ...) instead of typeof(A)(undef, ...)

Test improvements:

~25 bare == comparisons missing @test prefix now properly assert
New tests: non-square LUPP, Float32 RandSVD, view ColumnProjection, zero-allocation FJLT/SRHT

How has this been tested

Full test suite: julia --project -e 'using Pkg; Pkg.test()' — all 3299 tests pass.

Allocation benchmarks (scripts/drilldown.jl, scripts/profile_bench_fast.jl):

Metric	Before	After	Improvement
Kaczmarz alloc/iter (1000 iters)	73.2 KiB	1.0 KiB	73×
Kaczmarz total (1000 iters)	71.4 MiB	1.0 MiB	71×
ColumnProjection alloc/iter	2.1 MiB	21.7 KiB	99×
FJLT update_compressor!	32.2 KiB	0 B	eliminated
SRHT mul!	1.0 MiB	1.3 KiB	~800×
ldiv!(LQSolver)	8.8 KiB	608 B	15×
ldiv!(QRSolver)	22.6 KiB	4.4 KiB	5×
Kaczmarz+Sampling time/iter	655 us	200 us	3.3× faster

Types of changes

Fix
Performance
Test

Checklists:

Code and Comments

My code follows the code style of this project.
I have updated all package dependencies (if any).
I have included all relevant files to realize the functionality of the PR.
I have exported relevant functionality (if any).

Testing

I have added unit tests to cover my changes.
All new and existing tests passed.
I have achieved sufficient code coverage.

🤖 Generated with Claude Code

@inbounds

…ions Parameterize all Recipe structs to eliminate dynamic dispatch from abstract field types. Fix heap allocations in compressor and solver hot paths. Rewrite QR/LQ sub-solvers to use pre-allocated LAPACK workspace. Fix correctness bugs in LUPP, RandSVD, and error methods. Type stability: - Parameterize GaussianRecipe, CountSketchRecipe, SparseSignRecipe, SamplingRecipe, IdentityRecipe, RandSVDRecipe, RangeFinderRecipe, LUPPRecipe, QRCPRecipe, AgmonRecipe, BasicLoggerRecipe with concrete type parameters replacing abstract field types Allocation fixes — compressors: - SRHT mul!: pre-allocate extraction buffer and sign vector, eliminating ~1 MiB/call from fancy indexing and broadcast temporaries - FJLT update_compressor!: resample values in-place with randn!(nonzeros()) instead of allocating a new sparse matrix every call (32 KiB/call) - CountSketch update_compressor!: replace allocating [-1.0, 1.0] literal with non-allocating ifelse(rand(Bool), 1.0, -1.0) Allocation fixes — sub-solvers: - Rewrite QRSolverRecipe and LQSolverRecipe to use LAPACK.geqrt3!/gemqrt!/ trtrs! with pre-allocated T (block reflector) and work buffers, replacing qr!() which allocated a new QRCompactWY object every call - Use copyto! in update_sub_solver! instead of field reassignment, avoiding implicit conversion copies when the source is a SubArray (2 MiB/call) - Move shared LAPACK import and _qt_char helper to SubSolvers.jl Allocation fixes — solvers: - ColumnProjection complete_solver: replace fragile typeof(A)(undef, ...) with similar(), fixing breakage on views/sparse/adjoint inputs Correctness fixes: - LUPP complete_selector: fix buffer sized with compressor dims instead of output dims, causing DimensionMismatch on non-square matrices - RandSVD rapproximate!: replace hardcoded Matrix{Float64} with Matrix{eltype(A)} to preserve input element types - FullResidualRecipe/LSGradientRecipe: use zeros(eltype(b), ...) instead of zeros(...) to match declared type parameter Performance: - Add @inbounds to tight sparse loops in CountSketch, Sampling, SparseSign Test improvements: - Fix ~25 bare == comparisons missing @test prefix - Add tests for non-square LUPP, Float32 RandSVD, view ColumnProjection, zero-allocation FJLT/SRHT Results (Dense 4096x256, Kaczmarz-SparseSign, 1000 iters): - Per-iteration allocation: 73 KiB -> 1.0 KiB (73x reduction) - Total GC pressure: 71.4 MiB -> 1.0 MiB - ColumnProjection allocation: 2.1 MiB/iter -> 21.7 KiB/iter (99x) - All 3299 tests pass Co-Authored-By: Claude <noreply@anthropic.com>

codecov · 2026-08-05T03:37:38Z

Codecov Report

❌ Patch coverage is 98.11321% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/Solvers/SubSolvers.jl	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

vp314

There are many changes made here. Some of them are straightforward and understandable. Some of them are baffling.

vp314 · 2026-08-11T15:25:06Z

    # Allocate the information in the buffer using the types of A and b
-    compressed_mat = typeof(A)(undef, rows_a, sample_size) #Stores A*compressor
-    residual_vec = typeof(b)(undef, rows_a) #Stores b - Ax 
+    compressed_mat = similar(A, rows_a, sample_size)


In the PR write-up, the claim is that using typeof(A)(undef,...) and calling view produces allocations, but this does not seem to be true.

julia> m,n = 100,100; D = typeof(A)(undef, m,n); @allocated view(D, :, 1:5) 128 julia> m,n = 100,100; A = randn(m, n); B = similar(A); @allocated view(B, :, 1:5) 128

vp314 · 2026-08-11T15:31:55Z

    end

-    QA = Matrix{Float64}(undef, size(Q, 2), size(A, 2))
+    QA = Matrix{eltype(A)}(undef, size(Q, 2), size(A, 2))


Why not use similar here?

vp314 · 2026-08-11T15:33:00Z

    compressor = complete_compressor(ingredients.compressor, A)
-    n_rows, n_cols = size(compressor)
-    SA = Matrix{eltype(A)}(undef, n_rows, n_cols)
+    SA = Matrix{eltype(A)}(undef, size(compressor, 1), size(A, 2))


Why not use similar here?

vp314 · 2026-08-11T15:36:56Z

+- `A::MA`: reference to the coefficient matrix.
+- `b::VB`: reference to the constant vector.
+- `x::VX`: reference to the current solution iterate (updated each iteration).
+- `r::Union{Nothing, VB}`: the residual vector ``Ax - b``. Starts as


There may be more wisdom in preallocating r and including a boolean flag indicating whether it has been calculated from the data (i.e., A, x, and b) or if it was just allocated to values that carry not contextual meaning. This may be a larger code change that is not thematically connected here, so we may want to make it an issue.

vp314 · 2026-08-11T15:46:30Z

-    S.op = sprandn(type, n_rows, n_cols, S.sparsity)
-    # Resample the non-zero values 
+    randn!(nonzeros(S.op))
    rand!(S.signs)


Please put comments back in.

It might make sense to resample the rows to which nonzeros are assigned at minimum for this modified FJLT method (see next comment).

This modifies the algorithm such that it is no longer doing "textbook" FJLT. It is doing something else, which probably works just fine. We should modfiy the docstring and perhaps the name of this recipe to indicate that what is being done is not the same.

vp314 · 2026-08-11T15:50:12Z

+                row_P = P_rows[k_P]
+                val_P = P_nz[k_P]

                rng_B = nzrange(B, row_P)
                for k_B in rng_B
-                    col_C = B_rows[k_B] # Row index in B -> Column index in C
+                    col_C = B_rows[k_B]


Please put the comments back in.

vp314 · 2026-08-11T15:51:12Z

+    extraction::Matrix{Float64}
+    sign_vec::Vector{Float64}


These are missing from the docstring, so it is unclear what their value/purpose is.

vp314 · 2026-08-11T15:52:13Z

+    sign_vec = Vector{Float64}(undef, padded_size)
+    @inbounds for i in eachindex(signs)
+        sign_vec[i] = ifelse(signs[i], 1.0, -1.0)
+    end


A blank line after the for loop would better match the style.

vp314 · 2026-08-11T15:53:57Z

+    @inbounds for i in eachindex(signs)
+        sign_vec[i] = ifelse(signs[i], 1.0, -1.0)
+    end


What is the distinction between the signs and the sign_vec? It seems like this is a redundancy. Does using sign_vec over signs offer any benefit? Is it the @inbounds macro?

vp314 · 2026-08-11T15:56:00Z

-        # using fwht with signs. Scale the rows of the 
-        S.padding .*= ifelse.(S.signs, 1, -1) 
+        # using fwht with signs. Scale the rows of the
+        S.padding .*= S.sign_vec


Does this and line 385 confer such a savings that it makes sense to keep signs and sign_vec?

dmaldona requested review from nathanielpritchard and vp314 August 5, 2026 03:31

vp314 requested changes Aug 11, 2026

View reviewed changes

gisslab mentioned this pull request Aug 21, 2026

Add LeverageScore distribution #186

Open

24 tasks

Navigation Menu

Uh oh!

Fix type instability in Recipe structs and eliminate hot-loop allocations - #185

Fix type instability in Recipe structs and eliminate hot-loop allocations#185
dmaldona wants to merge 1 commit into
mainfrom
perf/recipe-type-stability-and-allocation-fixes

dmaldona commented Aug 5, 2026

Uh oh!

codecov Bot commented Aug 5, 2026

Uh oh!

vp314 left a comment

Uh oh!

vp314 Aug 11, 2026

Uh oh!

vp314 Aug 11, 2026

Uh oh!

vp314 Aug 11, 2026

Uh oh!

vp314 Aug 11, 2026

Uh oh!

vp314 Aug 11, 2026

Uh oh!

vp314 Aug 11, 2026

Uh oh!

vp314 Aug 11, 2026

Uh oh!

vp314 Aug 11, 2026

Uh oh!

vp314 Aug 11, 2026

Uh oh!

vp314 Aug 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

dmaldona commented Aug 5, 2026

Description

Motivation and Context

Changes

How has this been tested

Types of changes

Checklists:

Uh oh!

codecov Bot commented Aug 5, 2026

Codecov Report

Uh oh!

vp314 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants