Reduce allocations LinearSolve

Hi - I am trying to solve a linear problem repeatedly using LinearSolve.jl without allocating any new memory by mutating the input and solution arrays. Here is a minimal example

using LinearAlgebra, LinearSolve, IterativeSolvers, BenchmarkTools

const n = 1000
lhs = rand(n-1) # pre-allocating lhs, rhs
rhs = rand(n-1,n-1)
b_hat = Vector{Float64}(undef, n-1) # pre-allocating results
lin_prob = LinearProblem(rhs, lhs) # setting up the linear problem
lin_solve = init(lin_prob)

function update_b_hat!(b_hat::Vector{Float64}, rhs::Matrix{Float64}, lhs::Vector{Float64},
    lin_solve::LinearSolve.LinearCache{Matrix{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, Nothing, LU{Float64, Matrix{Float64}}, IterativeSolvers.Identity, IterativeSolvers.Identity, Float64, Nothing})
   lin_solve.A .= rhs
   lin_solve.b .= lhs
   b_hat .= solve(lin_solve).u # 6 allocations in solve()
lhs_new = rand(n-1)
rhs_new = rand(n-1,n-1)
update_b_hat!(b_hat, rhs_new, lhs_new, lin_solve)

Running the function update_b_hat! produces 8 allocations. The allocations are a problem as run this function in an inner loop. Is there a way to update both sides of the linear problem and solve it without making any new allocations?

I get a MethodError when trying to run the code (with the latest versions of LinearSolve), because one of the type parameters in the signature doesn’t match. Removing the type annotation of lin_solve::LinearSolve.LinearCache{...} and solves this issue and makes the code a lot more readable.

Also, I don’t get any allocations from update_b_hat! :thinking:

julia> @btime update_b_hat!($b_hat, $rhs_new, $lhs_new, $lin_solve)
  775.754 μs (0 allocations: 0 bytes)

How are you benchmarking the function? And what are the versions of the packages/Julia?

Julia Version 1.9.3
Commit bed2cd540a1 (2023-08-24 14:43 UTC)
Build Info:
  Official release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 64 × AMD Ryzen Threadripper PRO 3975WX 32-Cores
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, znver2)
  Threads: 1 on 64 virtual cores
pkg> status
Status `/tmp/jl_fVsy9u/Project.toml`
  [6e4b80f9] BenchmarkTools v1.3.2
  [42fd0dbc] IterativeSolvers v0.9.2
  [7ed4a6bd] LinearSolve v2.6.0
  [37e2e46d] LinearAlgebra
1 Like

Weird! I use the @btime from BenchmarkTools with dollar signs in front of all the arguments. I am using LinearSolve 1.29.1 and Julia 1.7.2

1 Like

Julia v1.9 has some of the missing ldiv! dispatches required for some cases to further reduce allocations.


Thank you for the quick replies! The issue was the versions of LinearSolve and Julia. When I upgraded past LinearSolve v2.0 and Julia v1.9, I now get no allocations

1 Like