Performance for diagonal matrix solves

[IlianPihlajamaa]

I want to repeatedly solve linear systems with a matrix that might be Diagonal. It seems that right now, LinearSolve.jl is not picking an optimal algorithm, or am I missing something? See benchmarks below. The results are similar for different sizes.

Thanks!

using LinearSolve, LinearAlgebra, Random, BenchmarkTools
function bench_diag(N)
    A1 = Diagonal(rand(N))
    A2 = copy(A1)
    b1 = rand(N)
    b2 = copy(b1)
    x = similar(b1)

    println("\nNaive:")
    @btime x = $A2\$b2
    x = A2\b2
    println(x[1])

    println("\npre-allocated:")
    @btime $x .= $A2.diag.\$b2
    println(x[1])

    println("\nLinearSolve1:")
    sol = @btime begin 
        prob = LinearProblem($A2, $b2)
        cache1 = init(prob)
        solve(cache1)
    end
    println(sol.u[1])

    prob = LinearProblem(A1, b1)
    cache1 = init(prob)
    sol = solve(cache1)
    println("\nLinearSolve2:")
    sol = @btime begin 
        cache2 = LinearSolve.set_b($cache1, $b2)
        cache2 = LinearSolve.set_A(cache2, $A2)
        solve(cache2)
    end
    println(sol.u[1])
end


bench_diag(100)

which prints

Naive:
  148.718 ns (1 allocation: 896 bytes)
0.6038470095547197

pre-allocated:
  46.613 ns (0 allocations: 0 bytes)
0.6038470095547197

LinearSolve1:
  1.710 μs (16 allocations: 4.48 KiB)
0.6038470095547197

LinearSolve2:
  393.500 ns (1 allocation: 112 bytes)
0.6038470095547197