GPU Support

[rossviljoen]

The issues I know of that currently prevent GPU support are:

• KernelFunctions uses hardcoded vectors for many kernels (e.g. ScaleTransform) Remove hard-coded Vector fields and add vectors constructors KernelFunctions.jl#299
• Distances.jl doesn't work with CuArrays Quite different performance of pairwise on CPU vs GPU JuliaStats/Distances.jl#143
• mean(fx) and cov(fx) return standard arrays (not CuArrays)

[willtebbutt]

Distances.jl doesn't work with CuArrays Quite different performance of pairwise on CPU vs GPU JuliaStats/Distances.jl#143

Since we define our own version of pairwise etc in KernelFunctions.jl, we could (for now) implement our own version of this functionality.

mean(fx) and cov(fx) return standard arrays (not CuArrays)

What are your thoughts on a fix for this? I would have imagined that it would be straightforward to modify these to just return a similar kind of array to whatever is passed in as an input, but perhaps I'm missing something?

[rossviljoen]

mean(fx) is a very easy fix - just need to replace fill with FillArrays.Fill etc. for mean functions.

The problem with cov(fx) is that

broadcast(+, ::CuMatrix, ::Diagonal{<:Real, <:Fill}) -> Matrix

so, I suppose a simple fix would be to replace the current implementation

Statistics.cov(f::FiniteGP) = cov(f.f, f.x) + f.Σy

with

function Statistics.cov(f::FiniteGP)
    C = cov(f.f, f.x)
    return broadcast!(+, similar(C), C, f.Σy)
end

[willtebbutt]

Hmm that's unfortunate -- it would be a shame to move to a mutating implementation. Another option would be to create a function my_add that defines + in the way that we want it defined (so as to avoid type-piracy), and maybe open an issue on FillArrays.jl about this? Possibly there's a simple way to make these types play nicely together.

[rossviljoen]

It looks like the reason it doesn't work is that GPUArrays can only broadcast over wrapped array types if the wrapped array is also an AbstractGPUArray - i.e. Diagonal{<:Real, <:AbstractGPUArray}

https://github.com/JuliaGPU/GPUArrays.jl/blob/bb9ca6d1f11e82a1d495cb9cf39cee9c215491e0/src/host/broadcast.jl#L22

I think you're right - we'd need a custom _add function doing something like:

function _add(A::AbstractGPUArray, B::Diagonal{T, <:FillArrays.AbstractFill}) where T
    Base.materialize(Base.broadcasted(Base.BroadcastStyle(typeof(A)), +, A, B))
end

although that would introduce a dependency on GPUArrays.jl in AbstractGPs.

---

I don't see an easy way to fix it in general though since wrapped arrays have to be dealt with explicity (see this and related issues). So, I think you'd have to somehow deal with wrapped fill arrays explicitly as well?

A more general version of the above using Adapt.jl's WrappedArray type:

const WrappedFillArray{T,N} = Adapt.WrappedArray{T, N, FillArrays.AbstractFill,FillArrays.AbstractFill{T,N}}

function _add(A::AbstractGPUArray, B::WrappedFillArray) where T
    Base.materialize(Base.broadcasted(Base.BroadcastStyle(typeof(A)), +, A, B))
end