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: GPU

Project.toml

@@ -3,6 +3,7 @@ uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
 version = "0.10.43"
 
 [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,10 +1,13 @@
 module TestUtils
 
+using CUDA
 using Distances
+using FillArrays
 using LinearAlgebra
 using KernelFunctions
 using Random
 using Test
+using InteractiveUtils
 
 """
     test_interface(
@@ -20,9 +23,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 +121,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 +191,102 @@ 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
-)
-    return map(n -> RowVecs(randn(rng, n, dim)), (1, 2, 3, 4))
+    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, x1::AbstractVector, x2::AbstractVector; atol=1e-6)
+    @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) == 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) == 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) == 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) == eltype(K_gpu)
+        @test isapprox(K_cpu, _to_cpu(K_gpu); atol=atol)
+    end
 end
 
+function test_gpu_against_cpu(rng::AbstractRNG, k::Kernel, data_type::Type)
+    _, x1, x2, _ = example_inputs(rng, data_type)
+    test_gpu_against_cpu(k, 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(κ.γ, 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/distances/cuda.jl

@@ -0,0 +1,123 @@
+
+const CuSubArray{T} = SubArray{T, <:Any, <:CuArray}
+const CuColVecs{T} = ColVecs{T, <:Union{CuMatrix{T}, CuSubArray{T}}}
+const CuRowVecs{T} = RowVecs{T, <:Union{CuMatrix{T}, CuSubArray{T}}}
+
+const _CuVector{T} = Union{CuVector{T}, SubArray{T, 1, <:CuArray}}
+
+# SqEuclidean
+
+# Vector{<:Real}
+
+pairwise(::SqEuclidean, x::_CuVector{<:Real}) = (x .- x').^2
+
+pairwise(::SqEuclidean, x::_CuVector{T}, y::_CuVector{T}) where {T<:Real} = (x .- y').^2
+
+colwise(::SqEuclidean, x::_CuVector{T}) where {T<:Real} = CUDA.zeros(T, length(x))
+
+colwise(::SqEuclidean, x::_CuVector{T}, y::_CuVector{T}) where {T<:Real} = (x .- y).^2
+
+# ColVecs
+
+function pairwise(::SqEuclidean, x::CuColVecs{<:Real})
+    X = CuArray(x.X) # needed to handle subarrays
+    X_ = sum(abs2, X; dims=1)
+    return X_ .+ X_' .- 2 .* (X' * X)
+end
+
+function pairwise(::SqEuclidean, x::CuColVecs{T}, y::CuColVecs{T}) where {T<:Real}
+    X = CuArray(x.X) # needed to handle subarrays
+    Y = CuArray(y.X) # needed to handle subarrays
+    X_ = sum(abs2, X; dims=1)
+    Y_ = sum(abs2, Y; dims=1)
+    return X_' .+ Y_ .- 2 .* (X' * Y)
+end
+
+function colwise(::SqEuclidean, x::CuColVecs{T}) where {T<:Real}
+    return CUDA.zeros(T, length(x))
+end
+
+function colwise(::SqEuclidean, x::CuColVecs{T}, y::CuColVecs{T}) where {T<:Real}
+    return vec(sum(abs2, x.X - y.X; dims=1))
+end
+
+# RowVecs
+
+function pairwise(::SqEuclidean, x::CuRowVecs{<:Real})
+    X = CuArray(x.X)
+    X_ = sum(abs2, X; dims=2)
+    return X_' .+ X_ .- 2 .* (X * X')
+end
+
+function pairwise(::SqEuclidean, x::CuRowVecs{T}, y::CuRowVecs{T}) where {T<:Real}
+    X = CuArray(x.X) # needed to handle subarrays
+    Y = CuArray(y.X) # needed to handle subarrays
+    X_ = sum(abs2, X; dims=2)
+    Y_ = sum(abs2, Y; dims=2)
+    return X_ .+ Y_' .- 2 .* (X * Y')
+end
+
+colwise(::SqEuclidean, x::CuRowVecs{T}) where {T<:Real} = CUDA.zeros(T, length(x))
+
+function colwise(::SqEuclidean, x::CuRowVecs{T}, y::CuRowVecs{T}) where {T<:Real}
+    return vec(sum(abs2, x.X - y.X; dims=2))
+end
+
+# Eucldiean: Derive from SqEuclidean implementations.
+
+const GPUVecType{T} = Union{_CuVector{T}, CuColVecs{T}, CuRowVecs{T}}
+
+pairwise(::Euclidean, x::GPUVecType{<:Real}) = sqrt.(pairwise(SqEuclidean(), x))
+
+function pairwise(::Euclidean, x::GPUVecType{T}, y::GPUVecType{T}) where {T<:Real}
+    return sqrt.(pairwise(SqEuclidean(), x, y))
+end
+
+colwise(::Euclidean, x::GPUVecType{T}) where {T<:Real} = CUDA.zeros(T, length(x))
+
+function colwise(::Euclidean, x::GPUVecType{T}, y::GPUVecType{T}) where {T<:Real}
+    return sqrt.(colwise(SqEuclidean(), x, y))
+end
+
+#
+# DotProduct
+#
+
+# Vector
+pairwise(::DotProduct, x::_CuVector{<:Real}) = x * x'
+
+pairwise(::DotProduct, x::_CuVector{T}, y::_CuVector{T}) where {T<:Real} = x * y'
+
+colwise(::DotProduct, x::_CuVector{<:Real}) = x.^2
+
+colwise(::DotProduct, x::_CuVector{T}, y::_CuVector{T}) where {T<:Real} = x .* y
+
+# ColVecs
+pairwise(::DotProduct, x::CuColVecs{<:Real}) = x.X' * x.X
+
+function pairwise(::DotProduct, x::CuColVecs{T}, y::CuColVecs{T}) where {T<:Real}
+    return x.X' * y.X
+end
+
+function colwise(::DotProduct, x::CuColVecs{<:Real})
+    return vec(sum(x.X.^2; dims=1))
+end
+
+function colwise(::DotProduct, x::CuColVecs{T}, y::CuColVecs{T}) where {T<:Real}
+    return vec(sum(x.X .* y.X; dims=1))
+end
+
+# RowVecs
+pairwise(::DotProduct, x::CuRowVecs{<:Real}) = x.X * x.X'
+
+function pairwise(::DotProduct, x::CuRowVecs{T}, y::CuRowVecs{T}) where {T<:Real}
+    return x.X * y.X'
+end
+
+function colwise(::DotProduct, x::CuRowVecs{<:Real})
+    return vec(sum(x.X.^2; dims=2))
+end
+
+function colwise(::DotProduct, x::CuRowVecs{T}, y::CuRowVecs{T}) where {T<:Real}
+    return vec(sum(x.X .* y.X; dims=2))
+end

src/kernels/kernelproduct.jl

@@ -46,6 +46,7 @@ Base.length(k::KernelProduct) = length(k.kernels)
 (κ::KernelProduct)(x, y) = prod(k(x, y) for k in κ.kernels)
 
 function kernelmatrix(κ::KernelProduct, x::AbstractVector)
+    # return reduce(hadamard, map(Base.Fix2(kernelmatrix, x), κ.kernels))
     return reduce(hadamard, kernelmatrix(k, x) for k in κ.kernels)
 end