Skip to content

Batch real matrix-RHS solves - #83

Open
mschytt wants to merge 3 commits into
SciML:mainfrom
mschytt:codex/blocked-matrix-rhs
Open

mschytt wants to merge 3 commits into
SciML:mainfrom
mschytt:codex/blocked-matrix-rhs

Conversation

@mschytt

@mschytt mschytt commented Sep 10, 2026 •

Copy link
Copy Markdown

SuiteSparse KLU already batches four RHS per sparse traversal. This PR brings RHS batching to PureKLU, using wider contiguous lane blocks to expose SIMD parallelism.

Batch Float32/Float64 matrix RHS across contiguous lanes to reuse sparse L/U coefficients. Use lazily allocated min(64, nrhs) × n scratch that grows only for wider solves, and 16-column pack/scatter sub-blocks.

  • Preserve arithmetic order, FMA policy, permutations, scaling, and BTF couplings.
  • Retain existing paths for fewer than eight RHS, generic/complex factors, and transpose/adjoint solves.

Performance

Five-point 2D diffusion with zero Dirichlet boundaries, n=N², N ∈ {10,20,40,80}. RHS counts k=1:step:n plus k=n, with steps {8,8,16,32} respectively. PureKLU baseline: cd8ee1c. Measurements use PR commit 93f6781, before adaptive scratch allocation.

Runtime comparison

N n=N² PureKLU (baseline), FMA PureKLU (baseline), non-FMA KLU.jl (PR FMA) KLU.jl (PR non-FMA)
10 100 2.57× 2.61× 1.46× 1.45×
20 400 3.30× 3.38× 1.65× 1.63×
40 1600 3.97× 4.03× 1.97× 1.95×
80 6400 4.34× 4.39× 2.06× 2.04×

Speedups for PureKLU (this PR): geometric mean of comparison time / PureKLU (this PR) time for sampled k ≥ 8, including partial blocks. FMA modes refer to PureKLU, with one KLU.jl reference.

Median of three randomized, warmed copyto! + ldiv! samples. Factorization excluded. Julia 1.12.6, Haswell, one pinned CPU, Julia/BLAS/GC threads = 1.

Input/output RHS storage at k=n: 0.153, 2.441, 39.063, and 625 MiB, excluding factors and scratch.

Validation

  • 107 assertions pass in three testsets: types, FMA/scaling, boundaries/tails, workspace ownership/reuse, refactorization, views, and fallbacks. The solver passed GROUP=Everything, including QA.
  • Exact agreement with PureKLU (baseline) in each FMA mode. Maximum observed backward error 6.56e-16 and relative difference from sparse LU 1.88e-14, checked on first/middle/last RHS.
  • Zero warmed PureKLU allocations, timed compilation, and GC.

Maintainer tuning: solve width 64, transfer width 16, and activation threshold 8 RHS are provisional, hardware- and matrix-sensitive choices. LLVM selects SIMD instructions within the loops.

The implementation and this PR description were generated by GPT-6 Astra.

Reproduce the benchmark

Save as bench.jl and run with the PR and baseline checkouts:

julia --project=/tmp/pureklu-bench -e 'using Pkg; Pkg.develop(path=ARGS[1]); Pkg.add("KLU")' /path/to/pr/PureKLU
JULIA_NUM_THREADS=1 JULIA_NUM_GC_THREADS=1 OPENBLAS_NUM_THREADS=1 \
  julia --project=/tmp/pureklu-bench bench.jl /path/to/baseline/PureKLU > timings.csv
using PureKLU, SparseArrays, LinearAlgebra, Random, Statistics, Printf
import KLU
module Stock end
Base.include(Stock, joinpath(ARGS[1], "src", "PureKLU.jl"))

function sample!(F, B, X, repeats)
    return @elapsed for _ in 1:repeats
        copyto!(X, B)
        ldiv!(F, X)
    end
end

function benchmark()
    BLAS.set_num_threads(1)
    rng = Xoshiro(2026091014)
    println("N,n,k,solver,median_ms,min_ms,max_ms")
    for (N, step) in ((10, 8), (20, 8), (40, 16), (80, 32))
        n = N^2
        T = spdiagm(-1 => -ones(N - 1), 0 => fill(2.0, N), 1 => -ones(N - 1))
        E = spdiagm(0 => ones(N))
        A = kron(E, T) + kron(T, E)
        factors = (KLU.klu(A), Stock.PureKLU.klu(A; use_fma = true),
            Stock.PureKLU.klu(A; use_fma = false), PureKLU.klu(A; use_fma = true),
            PureKLU.klu(A; use_fma = false))
        names = ("KLU", "stock_fma", "stock_nofma", "patch_fma", "patch_nofma")
        storage = randn(Xoshiro(2026091014 + N), n, n)
        output = similar(storage)
        for k in shuffle(rng, sort(unique([collect(1:step:n); n])))
            B, X = view(storage, :, 1:k), view(output, :, 1:k)
            repeats = map(factors) do F
                sample!(F, B, X, 1)
                t = sample!(F, B, X, 1)
                clamp(ceil(Int, 0.01 / max(t, 1e-9)), 1, 1000)
            end
            times = zeros(3, 5)
            for round in 1:3, i in randperm(rng, 5)
                times[round, i] = sample!(factors[i], B, X, repeats[i]) / repeats[i]
            end
            for i in 1:5
                t = times[:, i] .* 1000
                @printf("%d,%d,%d,%s,%.6g,%.6g,%.6g\n", N, n, k, names[i], median(t), minimum(t), maximum(t))
            end
        end
    end
end
benchmark()

Optionally prefix Julia with taskset -c <allocated-cpu> on Linux. The measured campaign also records accuracy, allocations, compilation, and provenance.

@mschytt
mschytt marked this pull request as ready for review September 10, 2026 19:12
@ChrisRackauckas

Copy link
Copy Markdown
Member

🤖 Automated review from an AI agent running as @ChrisRackauckas — not written or reviewed by Chris. It is posted so that you can act on it. Chris or a maintainer may disagree.

Verdict: changes needed before merge (reviewed at 93f6781).

Adds a lane-batched (64-RHS tile) path for real Float32/Float64 StridedMatrix solves with ≥8 columns, reusing each L/U/off-diagonal coefficient across RHS; it is bit-identical to the existing per-column path in every config I tested, and gives about 2–9× on 2D Laplacians. Two small things need fixing first: a word the repo's Spell Check CI rejects, and a lazily allocated 64 × n scratch that stays attached to the factor (0.5 GiB at n = 1e6) even when the solve has only 8 columns.

Risk assessment

  • Risk: low–medium. The algorithm is a straight loop interchange: per-lane arithmetic uses the same _mulsub/_cdiv calls in the same order as _klu_lsolve!/_klu_usolve! and the off-diagonal scatter, and no public API changes (KLUNumeric is not exported or public; the new field is internal).
  • Blast radius: every real solve!/ldiv!/\ with a matrix RHS of ≥8 columns (Float32/Float64), which covers LinearSolve's KLU path whenever it passes a multi-column RHS. Transpose/adjoint, complex, generic and vector solves are untouched.
  • Evidence:
    • CI on the head: no checks have run (gh pr checks → "no checks reported"). This is a fork PR, so the workflows probably need maintainer approval. Main's last push run (432cc60) is green on CI, Documentation, Downgrade, format-check and IntegrationTest.
    • Merge: mergeable: MERGEABLE. A local merge of the head into current origin/main (includes cf8a4be's pivot re-choice and 432cc60's rgrowth) was conflict-free.
    • Local runs (macOS M2 Max, Julia 1.11):
      • GROUP=Core on the head: pass, including Core/test_matrix_rhs.jl | 90 90 (picked up by SciMLTesting auto-discovery).
      • GROUP=QA on the head: pass (21 21).
      • Same two groups on the local merge with main: both pass, so the new Xrhs === scratch after-refactor assertion still holds on top of cf8a4be.
    • Runic --check on the 3 changed files: clean.
    • typos --config .typos.toml src/Solve.jl → exit 2 (Finding 1).
    • Bit-exactness, beyond the PR's 5×5 test: I loaded base (cd8ee1c) and PR as separate modules and compared solve! results with ==.
      • Matrices: a 2D Laplacian (n=1600) and a 2000×2000 permuted matrix with 3 BTF blocks plus off-diagonal coupling.
      • Grid: Float64/Float32 × use_fma true/false × scale 0/1/2 × nrhs ∈ {8, 63, 64, 65, 200}, plus a column-strided view RHS.
      • Result: 121/121 bitwise equal, maxdiff 0.
    • Performance (M2 Max, 1 thread, median of 7 warmed copyto!+solve! runs, Float64 Laplacian):
      • N=40, fma=true: nrhs 8 → 2.1×, 64 → 3.2×, 256 → 8.8×, 1600 → 5.9×.
      • N=80: nrhs 16 → about 2.4×, 64 → 2.3–4.0×, 256 → about 2.3×.
      • nrhs = 8 at N=80 with fma=true was about 1.0×.
      • Noise is large at sub-ms sizes: nrhs = 1 and 4 run identical code in both versions yet varied 0.5–1.9×. The direction for nrhs ≥ 16 is consistent and matches the PR's claimed 2.6–4.4× on Haswell.
    • Warmed @allocated solve!(K, X) at nrhs = 65: 0 bytes on both base and PR.
  • Merge: needs changes (Findings 1–2). Also approve the fork CI run before merging.

Findings

  1. src/Solve.jl:146 — The word "subtiles" fails the repo's typos config (.typos.toml has no allowlist entry). .github/workflows/SpellCheck.yml (SciML spellcheck) will go red once CI runs. Fix: reword it (e.g. "smaller sub-blocks") or add subtiles to [default.extend-words].
  2. src/Solve.jl:44-46 (and the Xrhs field, src/Kernel.jl:205) — The scratch is always allocated at 64 × n no matter how many columns the solve has. It then stays on the factor for its lifetime and survives klu!/refactor. Measured: 819,200 B at n = 1600 after a single 8-column solve, against 12,800 B for Xwork. At n = 1e6 in Float64 that is 512 MB permanently pinned by one 8-column solve, 8× the RHS itself, on exactly the large circuit-style matrices KLU is used for. Fix: allocate min(64, nrhs) × n and grow the row count only when a wider solve arrives (the tile loop already takes lanes). Document the memory cost somewhere a user can find it (docstring of solve! or the performance docs).
  3. (Non-blocking) test/test_matrix_rhs.jl — Exactness is only asserted on 5×5 and 2×2 matrices. A single larger case would lock the loop-interchange equivalence against future edits, e.g. the n=1600 Laplacian with nrhs=65 comparing solve! to column_solve! with ==. My ad-hoc check confirms that case passes today.
  4. (Non-blocking) src/Solve.jl:43 — The thresholds (8 RHS, 64-lane tile, 16-lane transfer) are hard-coded magic numbers, and the author calls them "provisional". Hoisting them into named consts would make later tuning a one-line change. No version bump is needed in this PR: it is an internal, non-breaking speedup, and the maintainer bumps at release.

Rules check: no reach into non-public API of other packages; no new dependencies; no skipped, weakened or deleted tests; no new public names (so no docs entry required); Runic clean. Comment lines are about 9% of added src lines and describe the code, not its history. AI authorship ("GPT-6 Astra") is disclosed in the PR body. No other maintainer has reviewed or commented; the PR has zero comments or reviews.

Push a fix and the PR is reviewed again automatically at the new head.


🤖 Posted by an AI agent — harness: Claude Code · model: claude-opus-5-5[1m] (fleet master); review by Devin CLI (Mac) / fusion-claude-opus-5-5-high-sidekick-swe-2-medium
Conversation: local Claude Code session 3cd6500a-1f81-46b5-ac0b-c466e15b6a53 on Chris's Mac (session ID, no URL)

The PR retained a benchmark script importing KLU while the workflow installed StableRNGs. Current main uses a StableRNGs-only benchmark script. Merge it to fix the benchmark environment mismatch.

Prepared by a GPT-6 Astra agent at the contributor’s request.

mschytt commented Nov 11, 2026

Copy link
Copy Markdown
Author

This follow-up was prepared and posted by a GPT-6 Astra agent at the contributor’s request.

The earlier review findings were addressed in 31bf9db: bounded scratch growth, workspace documentation, named tuning constants, spellcheck, and a larger exactness/reuse test. All 107 focused assertions and the full local suite passed.

Both benchmark CI failures occurred while loading the main-baseline benchmark, with Package KLU not found in current path. The PR retained the old script importing KLU, while the workflow installed StableRNGs. Current main already has the compatible StableRNGs-based script, so b7925e9 merges main rather than adding an unused KLU dependency. The merge was conflict-free and whitespace checks pass. The solver changes remain separate from this CI repair.

The new workflows report action_required and have not run. Could a maintainer approve the fork workflows and review the updated PR once the checks pass?

@ChrisRackauckas

Copy link
Copy Markdown
Member

🤖 Automated comment from an AI agent running as @ChrisRackauckas — not written or reviewed by Chris.

Independent review (Codex CLI v0.164.0 / gpt-6.1-sol), reviewer session 01a1cbf4-d41e-7a60-a82e-8285737eea3e, reviewed head b7925e9

  • Verdict: CHANGES
  • Risk assessment: medium (M3: confirmed warmed-solve slowdown of 1.65–1.83× on large diagonal matrix-RHS systems).
  • Round / slot: 1 / 1

This is a blind code-only review by a session other than the PR author's, copied below from the reviewer's report with only the grading metadata removed. The posting agent did not re-run the reviewer's commands. A MERGE verdict is advice; merging is left to the maintainers.

CI: CI unreadable; review only; merge gates unchanged
Session: 01a1cbf4-d41e-7a60-a82e-8285737eea3e

Risk

medium (M3: confirmed warmed-solve slowdown of 1.65–1.83× on large diagonal matrix-RHS systems).

Blocking findings

  • [P2] Avoid unconditional tiling for diagonal factors — src/Solve.jl:47. Every real strided matrix with at least eight columns enters the tiled path, including factors consisting entirely of singleton BTF blocks. For A = spdiagm(0 => fill(2.0, n)) and random Float64 RHS, warmed public solve! calls consistently regress against origin/main. Confirmed timings below are medians of nine samples, each containing five solves with identical RHS resets; allocation/compilation is warmed before timing.

    n RHS Base matrix (ms) Head vector columns (ms) Head matrix (ms)
    65,536 64 14.527 12.521 24.163
    65,536 65 13.655 12.609 22.574
    262,144 64 55.280 56.055 101.231
    262,144 65 58.599 62.429 99.872

    The same-head vector comparison isolates the tiled branch; all three routes also passed exact comparisons with the analytic solution B ./ 2. A separate randomized-diagonal benchmark reproduced the slowdown. Retain scalar dispatch for diagonal factors, or adjust tiling/dispatch to avoid this regression, and add matrix-RHS benchmarks covering singleton factors and coupled blocks. The existing benchmark/benchmarks.jl:34 exercises vector RHS only. Numerical correctness passed; this finding concerns sustained solve performance.

  • Verdict: CHANGES

  • Round: 1

  • Blocking priorities: P2

  • Deferred P2 issues: None supplied.

Non-blocking findings

None.

What I ran

Independent review by AI agent Codex (harness version unknown), model gpt-6.1-sol; local session recorded above, no session URL available. No author rating, PR discussion, or earlier review was read.

All Julia runs preserved the supplied layered JULIA_DEPOT_PATH, used the requested TMPDIR, four threads, and memory-capped scopes. Environments and probes are under ../scratch; the checkout remained clean.

  • Created scratch environments from the package's declared dependencies/compatibility and developed this checkout. The repository contains no committed Manifest or interpreter lock. Actual interpreters: Julia 1.13.2 and LTS 1.10.12.
  • GROUP=Core julia --startup-file=no --project=../scratch/env -e 'using Pkg, InteractiveUtils; versioninfo(); Pkg.test("PureKLU")': 2,562/2,562 assertions passed across 16 files, including 107/107 new matrix-RHS assertions.
  • GROUP=QA julia --startup-file=no --project=../scratch/env -e 'using Pkg; Pkg.test("PureKLU")': 21/21 passed.
  • julia --startup-file=no --check-bounds=yes --project=../scratch/env ../scratch/probes.jl: 5,524/5,524 passed. Independent dense LU and precision-derived residual bounds, Float32/Float64, Int32/Int64, FMA modes, scaling modes, pivoting, BTF couplings, tile boundaries, reversed/gapped views, poisoned workspace reuse, changed-value refactors, deepcopy, nonfinite and subnormal values.
  • julia +lts --startup-file=no --project=../scratch/ltsenv: included test/test_matrix_rhs.jl and ../scratch/probes.jl after developing this checkout: 107/107 + 5,524/5,524 passed.
  • Ran ../scratch/bench.jl, bench_recheck.jl, bench_diagonal_large.jl, and diagonal_confirm.jl using Julia 1.13.2 and --project=../scratch/env, comparing with sources extracted by git archive origin/main. Coupled grid/random cases improved substantially; large diagonal cases regressed. The confirmation's 12/12 analytic assertions passed; timings are above and in ../scratch/diagonal_confirm.log.
  • julia --project=@runic -m Runic --check src/Kernel.jl src/PureKLU.jl src/Solve.jl test/test_matrix_rhs.jl, typos on those files, and git diff --check: passed.
  • uv run --locked python --version: Python 3.15.0; uv warned that --locked has no effect outside a project. There is no Python project or uv.lock here, so this does not establish a project-locked Python environment.

What I did not verify

CI, documentation builds, downstream packages, Windows/macOS, 32-bit execution, Julia 1.12, or the full Core/QA suites on LTS. Timings were obtained on a shared workstation rather than an isolated benchmark host; the large diagonal regression reproduced across two probes and against same-head scalar solves. No checkout edits, commits, pushes, or GitHub writes were made.


🤖 Posted by an AI agent — harness: Claude Code 2.1.322 · model: claude-sonnet-5-5
Conversation: local session 3878b54e-28f2-4a01-a176-a8f5468d1997 (no session URL)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants