Support custom vector types

[amontoison]

In some specific applications, users may want to use their own vector types. For example, in Oceananigans.jl, each vector is a slice of a 3D matrix and is of the type Field. (cc @glwagner)
--> CliMA/Oceananigans.jl#3812

The best solution is to document how we can create a modular wrapper that fits the user's needs, with a very simple and limited API to implement.

I quickly experimented with a custom type that represents a vector of dual numbers, and it works.
Off-topic: @michel2323, this could be a way to perform low-level differentiation directly in Krylov.jl.

using ForwardDiff, Krylov
import ForwardDiff.Dual
import Krylov.FloatOrComplex

struct ADVector{T, V, N} <: AbstractVector{T}
    v::Vector{Dual{V, T, N}}
end

function ADVector{T, V, N}(::UndefInitializer, n::Int) where {T, V, N}
    ad_vector = [ForwardDiff.Dual{V, T, N}(0.0) for _ in 1:n]
    ADVector{T, V, N}(ad_vector)
end

function ADVector{T}(::UndefInitializer, n::Int) where T
    ad_vector = [ForwardDiff.Dual{Nothing, T, 0}(0.0) for _ in 1:n]
    ADVector{T, Nothing, 0}(ad_vector)
end

function ADVector{T}(v::Vector{T}) where T
    n = length(v)
    ad_vector = [ForwardDiff.Dual{Nothing, T, 0}(v[i]) for i in 1:n]
    return ADVector{T, Nothing, 0}(ad_vector)
end

function Base.length(v::ADVector)
    return length(v.v)
end

function Base.size(v::ADVector)
    return (length(v.v),)
end

function Base.getindex(v::ADVector, i::Int)
    return v.v[i]
end

function Base.setindex!(v::ADVector{T, V, N}, value::ForwardDiff.Dual{V, T, N}, i::Int) where {T, V, N}
    v.v[i] = value
    return v
end

mul!(y::ADVector{T}, A::Matrix{T}, x::Vector{T}) where T <: FloatOrComplex = mul!(y.v, A, x.v)

Krylov.kdot(n::Integer, x::ADVector{T}, dx::Integer, y::ADVector{T}, dy::Integer) where T <: FloatOrComplex = dot(x.v, y.v)
Krylov.knrm2(n::Integer, x::ADVector{T}, dx::Integer) where T <: FloatOrComplex = norm(x.v)

Krylov.kscal!(n::Integer, s, x::ADVector{T}, dx::Integer) where T <: FloatOrComplex = (x.v .*= s)

Krylov.kaxpy!(n::Integer, s, x::ADVector{T}, dx::Integer, y::ADVector{T}, dy::Integer) where T <: FloatOrComplex = axpy!(s, x.v, y.v)

Krylov.kaxpby!(n::Integer, s, x::ADVector{T}, dx::Integer, t, y::ADVector{T}, dy::Integer) where T <: FloatOrComplex = axpby!(s, x.v, t, y.v)

Krylov.kcopy!(n::Integer, x::ADVector{T}, dx::Integer, y::ADVector{T}, dy::Integer) where T <: FloatOrComplex = copyto!(y.v, x.v)

Krylov.kfill!(x::ADVector{T}, val::T) where T <: FloatOrComplex = fill!(x.v, val)

I should clean up the low-level API (functions starting with k), as these are not publicly available or documented for the user.

@michel2323 @glwagner
I believe this is also the best way to introduce MPI usage in the future.
It will be independent of Krylov.jl and can be applied to any package.

[glwagner]

In some specific applications, users may want to use their own vector types. For example, in Oceananigans.jl, each vector is a slice of a 3D matrix and is of the type Field. (cc @glwagner)

To clarify, what we have is a situation in which our data should be considered as a 1D vector for linear algebra, but represents a 3D field in a physical context.

Consider a 2D array that represents, for example, an image:

julia> c = rand(3, 3)
3×3 Matrix{Float64}:
 0.515044  0.981678  0.0473222
 0.437767  0.142471  0.756243
 0.284507  0.647186  0.445389

A "diffusion" operation, which can be interpreted as a kind of filter, is represented by the following operator:

diffuse(i, j, c) = - 4 * c[i, j] + c[i+1, j] + c[i, j+1] + c[i-1, j] + c[i, j-1]

The operator diffuse can also be written as a 9 x 9 matrix acting on the "vectorization" of c:

julia> c_vector = c[:]
9-element Vector{Float64}:
 0.5150436482179653
 0.43776719843960443
 0.28450746235635493
 0.9816775509672155
 0.14247143653959915
 0.6471861061623826
 0.0473222143036538
 0.7562428994779301
 0.4453890152898158

Said again, we can apply the diffuse operator to c either by forming a product A_diffuse * c_vector where A_diffuse is a (sparse) matrix that represents the operator diffuse --- or by applying the function diffuse(i, j, c) in a loop over i = 1:size(c, 1) and j = size(c, 2).

In Oceananigans we are concerned with 3D fluid dynamics, so you can take the above analogy but apply to a 3D matrix. If the matrix has Nx, Ny, Nz = size(c), then the "operator" has size (Nx * Ny * Nz)^2.

While we could in principle represent our operators with sparse matrices and then evaluate matrix-vector products with sparse linear algebra, there's no intrinsic benefit to spending memory on the sparse matrix; we are interested in iterative methods that allow us to avoid allocating operator memory entirely.

I think there are probably many such contexts, where the "native" data representation is an M-dimensional matrix with dimension sizes, N = [N1, N2, ..., NM], but which can also be represented as a vector of length prod(N), which is acted on by a matrix of size prod(N)^2.

[amontoison]

Thanks for the explanation!
I started to clean the low-level API in #907 to provide what you need.
We probably just need to create a wrapper in Oceanigans.jl to represent the 3D field as an abstract vector and define these functions now:

• mul!(y, A, x)
• length(x)
• Krylov.kdot
• Krylov.scal!
• Krylov.kaxpy!
• Krylov.kaxpby!
• Krylov.knorm
• Krylov.kfill!
• Krylov.kcopy!
• Krylov.ref!

I will document that after INFORMS.

[amontoison]

@glwagner
I just added an example to the documentation explaining how a 3D field in a physical context, represented as a 3D array, can also represent a 1D vector for linear algebra:
https://jso.dev/Krylov.jl/dev/matrix_free/#Example-with-discretized-PDE.

I’m wondering if we can do something similar for Field so that we can still target BLAS, CUBLAS, rocBLAS, ....

If it’s not possible, we will define our own wrapper as we discussed last week (not documented yet).
The low-level API has now been cleaned up, and I released version 0.9.8 today, so we can already try it in Oceanigans.jl.

[glwagner]

@glwagner I just added an example to the documentation explaining how a 3D field in a physical context, represented as a 3D array, can also represent a 1D vector for linear algebra: https://jso.dev/Krylov.jl/dev/matrix_free/#Example-with-discretized-PDE.

I’m wondering if we can do something similar for Field so that we can still target BLAS, CUBLAS, rocBLAS, ....

If it’s not possible, we will define our own wrapper as we discussed last week (not documented yet). The low-level API has now been cleaned up, and I released version 0.9.8 today, so we can try it in Oceanigans.jl.

Nice!

I just want to note for clarity -- the key tricky part of the implementation is skipped over in that example, which is how boundary conditions are treated. For performance we cannot use short-circuiting operations which check whether we are on the boundary; instead (at least for a finite volume or difference code) the typical approach is to use "halo regions". In other words, the underlying data will be a 3D continuous array of size (Nx+2, Ny+2, Nz+2) where each index ranges from 0, to N+1. This allows the operators to be applied on the boundaries with no out-of-bounds issue. The boundary conditions are usually satisfied by "filling" the halo regions with appropriate values.

This means that we cannot use neat tricks like reshape to move a 3D matrix to a vector because the data is no longer contiguous. Instead, we require a wrapper in order to subtype AbstractVector.

[amontoison]

Thanks for the explanation, @glwagner!
I thought that the boundary conditions only lead to modifications in the right-hand side b, but it means that your operator is more complex and difficult to implement if you do that.

I will document how we can define our own vector type with an example of discretized PDE with halo regions.

We will have all the code that we need for Oceanigans.jl after that.

[glwagner]

Thanks for the explanation, @glwagner! I thought that the boundary conditions only lead to modifications in the right-hand side b, but it means that your operator is more complex and difficult to implement if you do that.

Let me clarify. Even homogeneous boundary conditions require some special attention to the boundary. The reason is not the condition itself but rather because the operator being applied is not generally correct or even applicable in-bounds for points that touch the boundary. In the code you've given in the example, this is accomplished with if-statements. However, such short-circuiting if statements generally lead to a loss of performance within such hot inner loops and are thus often avoided in performance intensive code.

There is an additional point that inhomogeneous boundary conditions can be accounted for in the Poisson equation by changing the right hand side of the equation. But these changes would be implemented in addition to any special attention paid to evaluating a stencil on the boundary as described above.

[glwagner]

I will document how we can define our own vector type with an example of discretized PDE with halo regions.

We will have all the code that we need for Oceanigans.jl after that.

Agreed. For a minimal / example implementation what is needed is something like

struct ThreeDimensionalizedVector{FT, D} <: AbstractVector{FT}
    data :: D
    function ThreeDimensionalizedVector(data)
        FT = eltype(data)
        D = typeof(data)
        return new{FT, D}(data)
    end
end

Then if data is an OffsetArray we can implement an "if-less" stencil.