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()
end
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?

versioninfo()
Julia Version 1.9.3
Commit bed2cd540a1 (2023-08-24 14:43 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 64 × AMD Ryzen Threadripper PRO 3975WX 32-Cores
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, znver2)
  Threads: 1 on 64 virtual cores
Environment:
  JULIA_PKG_USE_CLI_GIT = true
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 
  JULIA_IMAGE_THREADS = 1
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.

3 Likes

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