IterativeSolvers.jl and DistributedArrays.jl

I would like to use the iterative solvers from IterativeSolvers.jl with DistributedArrays.jl, with no luck for now.

Here is a MWE, where I try to invert a diagonal positive definite matrix using the CG algorithm :

using LinearAlgebra, IterativeSolvers, DistributedArrays
n = 100
A = Diagonal(drand(n) .^ 2 .+ √eps())
b = drand(n)
cg(A, b)

The error message mentions some ambiguities in the methods definition, and proposes to implement the following method

copyto!(::DistributedArrays.DArray, ::Base.Broadcast.Broadcasted{T, Axes, F, Args} where {T<:Base.Broadcast.AbstractArrayStyle{0}, Axes, F, Args<:Tuple})

If anyone encountered this issue previously, any tips on how to do this would be greatly appreciated. Thanks!

1 Like

Note this doesn’t work with KrylovKit either:

using LinearAlgebra, KrylovKit, DistributedArrays
n = 100
b = drand(n)
linsolve(x -> b .* x, b, b)

complains about strides(DArray). I’d open an issue in the relevant packages.

@antoine-levitt Sure, I did just that here, thanks.

In fact the problem is not too bad:

  [1] stride(A::DArray{Float64, 1, Vector{Float64}}, k::Int64)
    @ Base ./abstractarray.jl:501
  [2] axpby!(alpha::Float64, x::DArray{Float64, 1, Vector{Float64}}, beta::Float64, y::DArray{Float64, 1, Vector{Float64}})
    @ LinearAlgebra.BLAS /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/LinearAlgebra/src/blas.jl:649
  [3] _rmul!(b::KrylovKit.OrthonormalBasis{DArray{Float64, 1, Vector{Float64}}}, G::LinearAlgebra.Givens{Float64})
    @ KrylovKit ~/.julia/packages/KrylovKit/OLgKs/src/dense/givens.jl:33

so the problem is that DArrays does not support axpby!. We can fix this:

import LinearAlgebra
LinearAlgebra.axpby!(alpha::Float64, x::DArray{Float64, 1, Vector{Float64}}, beta::Float64, y::DArray{Float64, 1, Vector{Float64}}) = axpby!(alpha, x.localpart, beta, y.localpart)

which works! This should go into a PR on DistributedArrays, but it looks like that package is not actively maintained…