GPU Attempt 2

.buildkite/pipeline.yml

@@ -0,0 +1,14 @@
+steps:
+  - label: "Julia v1"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+      - JuliaCI/julia-test#v1: ~
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    if: build.message !~ /\[skip tests\]/
+    timeout_in_minutes: 60
+
+env:
+  GROUP: CUDA

Project.toml

@@ -1,8 +1,9 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.43"
+version = "0.10.44"
 
 [deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 CompositionsBase = "a33af91c-f02d-484b-be07-31d278c5ca2b"
@@ -25,6 +26,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 ChainRulesCore = "1"
 Compat = "3.7, 4"
 CompositionsBase = "0.1"
+CUDA = "3"
 Distances = "0.10"
 FillArrays = "0.10, 0.11, 0.12, 0.13"
 Functors = "0.1, 0.2"

src/KernelFunctions.jl

@@ -47,6 +47,7 @@ export IndependentMOKernel,
 export tensor, ⊗, compose
 
 using Compat
+using CUDA
 using ChainRulesCore: ChainRulesCore, Tangent, ZeroTangent, NoTangent
 using ChainRulesCore: @thunk, InplaceableThunk
 using CompositionsBase
@@ -75,6 +76,7 @@ include("distances/pairwise.jl")
 include("distances/dotproduct.jl")
 include("distances/delta.jl")
 include("distances/sinus.jl")
+include("distances/cuda.jl")
 
 include("transform/transform.jl")
 include("transform/scaletransform.jl")

src/TestUtils.jl

@@ -1,6 +1,8 @@
 module TestUtils
 
+using CUDA
 using Distances
+using FillArrays
 using LinearAlgebra
 using KernelFunctions
 using Random
@@ -20,9 +22,9 @@ Run various consistency checks on `k` at the inputs `x0`, `x1`, and `x2`.
 `x0` and `x1` should be of the same length with different values, while `x0` and `x2` should
 be of different lengths.
 
-These tests are intended to pick up on really substantial issues with a kernel implementation
-(e.g. substantial asymmetry in the kernel matrix, large negative eigenvalues), rather than to
-test the numerics in detail, which can be kernel-specific.
+These tests are intended to pick up on really substantial issues with a kernel
+implementation (e.g. substantial asymmetry in the kernel matrix, large negative
+eigenvalues), rather than to test the numerics in detail, which can be kernel-specific.
 """
 function test_interface(
     k::Kernel,
@@ -118,6 +120,18 @@ function test_interface(
     )
 end
 
+function test_interface(
+    rng::AbstractRNG, k::Kernel, ::Type{<:ColVecs{T, CuMatrix{T}}}; dim_in=2, kwargs...
+) where {T<:Real}
+    return test_interface(
+        k,
+        _to_cuda_gpu(ColVecs(randn(rng, T, dim_in, 11))),
+        _to_cuda_gpu(ColVecs(randn(rng, T, dim_in, 11))),
+        _to_cuda_gpu(ColVecs(randn(rng, T, dim_in, 13)));
+        kwargs...,
+    )
+end
+
 function test_interface(
     rng::AbstractRNG, k::Kernel, ::Type{<:RowVecs{T}}; dim_in=2, kwargs...
 ) where {T<:Real}
@@ -176,27 +190,108 @@ function test_interface(k::Kernel, T::Type{<:Real}=Float64; kwargs...)
     return test_interface(Random.GLOBAL_RNG, k, T; kwargs...)
 end
 
+const FloatType = Union{Float32, Float64}
+
 """
     example_inputs(rng::AbstractRNG, type)
 
 Return a tuple of 4 inputs of type `type`. See `methods(example_inputs)` for information
 around supported types. It is recommended that you utilise `StableRNGs.jl` for `rng` here
 to ensure consistency across Julia versions.
 """
-function example_inputs(rng::AbstractRNG, ::Type{Vector{Float64}})
-    return map(n -> randn(rng, Float64, n), (1, 2, 3, 4))
+function example_inputs(rng::AbstractRNG, ::Type{Vector{T}}) where {T<:FloatType}
+    return map(n -> randn(rng, T, n), (1, 2, 3, 4))
 end
 
 function example_inputs(
-    rng::AbstractRNG, ::Type{ColVecs{Float64,Matrix{Float64}}}; dim::Int=2
-)
-    return map(n -> ColVecs(randn(rng, dim, n)), (1, 2, 3, 4))
+    rng::AbstractRNG, ::Type{ColVecs{T,Matrix{T}}}; dim::Int=2
+) where {T<:FloatType}
+    return map(n -> ColVecs(randn(rng, T, dim, n)), (1, 2, 3, 4))
 end
 
 function example_inputs(
-    rng::AbstractRNG, ::Type{RowVecs{Float64,Matrix{Float64}}}; dim::Int=2
+    rng::AbstractRNG, ::Type{RowVecs{T,Matrix{T}}}; dim::Int=2
+) where {T<:FloatType}
+    return map(n -> RowVecs(randn(rng, T, n, dim)), (1, 2, 3, 4))
+end
+
+function example_inputs(rng::AbstractRNG, ::Type{CuVector{T}}) where {T<:FloatType}
+    return map(_to_cuda_gpu, example_inputs(rng, Vector{T}))
+end
+
+function example_inputs(
+    rng::AbstractRNG, ::Type{ColVecs{T, CuMatrix{T}}}; dim::Int=2
+) where {T<:FloatType}
+    return map(_to_cuda_gpu, example_inputs(rng, ColVecs{T, Matrix{T}}; dim=dim))
+end
+
+function example_inputs(
+    rng::AbstractRNG, ::Type{RowVecs{T, CuMatrix{T}}}; dim::Int=2
+) where {T<:FloatType}
+    return map(_to_cuda_gpu, example_inputs(rng, RowVecs{T, Matrix{T}}; dim=dim))
+end
+
+function test_gpu_against_cpu(
+    k::Kernel, T::Type, x1::AbstractVector, x2::AbstractVector; atol=1e-3
 )
-    return map(n -> RowVecs(randn(rng, n, dim)), (1, 2, 3, 4))
+    @assert length(x1) != length(x2)
+
+    k_cpu = _to_cpu(k)
+    x1_cpu = _to_cpu(x1)
+    x2_cpu = _to_cpu(x2)
+    let
+        K_cpu = kernelmatrix(k_cpu, x1_cpu)
+        K_gpu = kernelmatrix(k, x1)
+        @test size(K_cpu) == size(K_gpu)
+        @test eltype(K_cpu) == T
+        @test eltype(K_cpu) == eltype(K_gpu)
+        @test isapprox(K_cpu, _to_cpu(K_gpu); atol=atol)
+    end
+    let
+        K_cpu = kernelmatrix(k_cpu, x1_cpu, x2_cpu)
+        K_gpu = kernelmatrix(k, x1, x2)
+        @test size(K_cpu) == size(K_gpu)
+        @test eltype(K_cpu) == T
+        @test eltype(K_cpu) == eltype(K_gpu)
+        @test isapprox(K_cpu, _to_cpu(K_gpu); atol=atol)
+    end
+    let
+        K_cpu = kernelmatrix_diag(k_cpu, x1_cpu)
+        K_gpu = kernelmatrix_diag(k, x1)
+        @test size(K_cpu) == size(K_gpu)
+        @test eltype(K_cpu) == T
+        @test eltype(K_cpu) == eltype(K_gpu)
+        @test isapprox(K_cpu, _to_cpu(K_gpu); atol=atol)
+    end
+    let
+        K_cpu = kernelmatrix_diag(k_cpu, x1_cpu, x1_cpu)
+        K_gpu = kernelmatrix_diag(k, x1, x1)
+        @test size(K_cpu) == size(K_gpu)
+        @test eltype(K_cpu) == T
+        @test eltype(K_cpu) == eltype(K_gpu)
+        @test isapprox(K_cpu, _to_cpu(K_gpu); atol=atol)
+    end
 end
 
+function test_gpu_against_cpu(rng::AbstractRNG, T::Type, k::Kernel, data_type::Type)
+    _, x1, x2, _ = example_inputs(rng, data_type)
+    test_gpu_against_cpu(k, T, x1, x2)
+end
+
+_to_cpu(x::CuArray{<:Real}) = Array(x)
+_to_cpu(x::ColVecs{<:Real, <:CuMatrix}) = ColVecs(_to_cpu(x.X))
+_to_cpu(x::RowVecs{<:Real, <:CuMatrix}) = RowVecs(_to_cpu(x.X))
+_to_cpu(x::FillArrays.AbstractFill{<:Real}) = x
+
+_to_cpu(k::Kernel) = k
+_to_cpu(k::TransformedKernel) = TransformedKernel(_to_cpu(k.kernel), _to_cpu(k.transform))
+_to_cpu(t::Transform) = t
+_to_cpu(t::LinearTransform) = LinearTransform(_to_cpu(t.A))
+_to_cpu(t::ChainTransform) = ChainTransform(map(_to_cpu, t.transforms))
+
+_to_cuda_gpu(x::Array{<:Real}) = CuArray(x)
+_to_cuda_gpu(x::ColVecs{T, Matrix{T}} where {T<:Real}) = ColVecs(_to_cuda_gpu(x.X))
+_to_cuda_gpu(x::RowVecs{T, Matrix{T}} where {T<:Real}) = RowVecs(_to_cuda_gpu(x.X))
+_to_cuda_gpu(x::FillArrays.AbstractFill{<:Real}) = x
+
 end # module

src/basekernels/exponential.jl

@@ -137,10 +137,28 @@ end
 
 @functor GammaExponentialKernel
 
-kappa(κ::GammaExponentialKernel, d::Real) = exp(-d^only(κ.γ))
+_gamma_kappa(γ::Real, d::Real) = exp(-d^γ)
+
+kappa(κ::GammaExponentialKernel, d::Real) = _gamma_kappa(only(κ.γ), d)
 
 metric(k::GammaExponentialKernel) = k.metric
 
+function kernelmatrix(k::GammaExponentialKernel, x::AbstractVector)
+    return _gamma_kappa.(only(k.γ), pairwise(metric(k), x))
+end
+
+function kernelmatrix(k::GammaExponentialKernel, x::AbstractVector, y::AbstractVector)
+    return _gamma_kappa.(only(k.γ), pairwise(metric(k), x, y))
+end
+
+function kernelmatrix_diag(k::GammaExponentialKernel, x::AbstractVector)
+    return _gamma_kappa.(only(k.γ), colwise(metric(k), x))
+end
+
+function kernelmatrix_diag(k::GammaExponentialKernel, x::AbstractVector, y::AbstractVector)
+    return _gamma_kappa.(only(k.γ), colwise(metric(k), x, y))
+end
+
 iskroncompatible(::GammaExponentialKernel) = true
 
 function Base.show(io::IO, κ::GammaExponentialKernel)

src/basekernels/fbm.jl

@@ -48,9 +48,19 @@ _mod(x::AbstractVector{<:AbstractVector{<:Real}}) = sum.(abs2, x)
 _mod(x::ColVecs) = vec(sum(abs2, x.X; dims=1))
 _mod(x::RowVecs) = vec(sum(abs2, x.X; dims=2))
 
+_ones(el_type, x::AbstractVector, sz...) = ones(el_type, sz...)
+function _ones(
+    el_type,
+    x::Union{CuVector{<:Real}, ColVecs{<:Real, <:CuMatrix}, RowVecs{<:Real, <:CuMatrix}},
+    sz...,
+)
+    return CUDA.ones(el_type, sz...)
+end
+
 function kernelmatrix(κ::FBMKernel, x::AbstractVector)
     modx = _mod(x)
-    modx_wide = modx * ones(eltype(modx), 1, length(modx)) # ad perf hack -- is unit tested
+    _v = _ones(eltype(modx), x, 1, length(modx))
+    modx_wide = modx * _v # ad perf hack -- is unit tested
     modxx = pairwise(SqEuclidean(), x)
     return _fbm.(modx_wide, modx_wide', modxx, only(κ.h))
 end
@@ -64,8 +74,10 @@ end
 
 function kernelmatrix(κ::FBMKernel, x::AbstractVector, y::AbstractVector)
     modxy = pairwise(SqEuclidean(), x, y)
-    modx_wide = _mod(x) * ones(eltype(modxy), 1, length(y)) # ad perf hack -- is unit tested
-    mody_wide = _mod(y) * ones(eltype(modxy), 1, length(x)) # ad perf hack -- is unit tested
+    _vx = _ones(eltype(modxy), x, 1, length(y))
+    _vy = _ones(eltype(modxy), y, 1, length(x))
+    modx_wide = _mod(x) * _vx # ad perf hack -- is unit tested
+    mody_wide = _mod(y) * _vy # ad perf hack -- is unit tested
     return _fbm.(modx_wide, mody_wide', modxy, only(κ.h))
 end
 

src/basekernels/matern.jl

@@ -77,7 +77,7 @@ end
 
 Matern32Kernel(; metric=Euclidean()) = Matern32Kernel(metric)
 
-kappa(::Matern32Kernel, d::Real) = (1 + sqrt(3) * d) * exp(-sqrt(3) * d)
+kappa(::Matern32Kernel, d::T) where {T<:Real} = (1 + T(sqrt(3)) * d) * exp(-T(sqrt(3)) * d)
 
 metric(k::Matern32Kernel) = k.metric
 
@@ -108,7 +108,7 @@ end
 
 Matern52Kernel(; metric=Euclidean()) = Matern52Kernel(metric)
 
-kappa(::Matern52Kernel, d::Real) = (1 + sqrt(5) * d + 5 * d^2 / 3) * exp(-sqrt(5) * d)
+kappa(::Matern52Kernel, d::T) where {T<:Real} = (1 + T(sqrt(5)) * d + 5 * d^2 / 3) * exp(-T(sqrt(5)) * d)
 
 metric(k::Matern52Kernel) = k.metric
 

src/basekernels/wiener.jl

@@ -62,22 +62,22 @@ function (::WienerKernel{1})(x, y)
     X = sqrt(sum(abs2, x))
     Y = sqrt(sum(abs2, y))
     minXY = min(X, Y)
-    return 1 / 3 * minXY^3 + 1 / 2 * minXY^2 * euclidean(x, y)
+    return minXY^3 / 3 + minXY^2 * euclidean(x, y) / 2
 end
 
 function (::WienerKernel{2})(x, y)
     X = sqrt(sum(abs2, x))
     Y = sqrt(sum(abs2, y))
     minXY = min(X, Y)
-    return 1 / 20 * minXY^5 + 1 / 12 * minXY^3 * euclidean(x, y) * (X + Y - 1 / 2 * minXY)
+    return minXY^5 / 20 + minXY^3 * euclidean(x, y) * (X + Y - minXY / 2) / 12
 end
 
 function (::WienerKernel{3})(x, y)
     X = sqrt(sum(abs2, x))
     Y = sqrt(sum(abs2, y))
     minXY = min(X, Y)
-    return 1 / 252 * minXY^7 +
-           1 / 720 * minXY^4 * euclidean(x, y) * (5 * max(X, Y)^2 + 2 * X * Y + 3 * minXY^2)
+    return minXY^7 / 252 +
+           minXY^4 * euclidean(x, y) * (5 * max(X, Y)^2 + 2 * X * Y + 3 * minXY^2) / 720
 end
 
 function Base.show(io::IO, ::WienerKernel{I}) where {I}