AD Performance (#467)

* Tuples rather than vectors

* Testing

* Bump patch

* Formatting

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Docstring line lengths

* Data generation

* Add StableRNGs to test Project

* Add ad tests to test utils

* Test constant kernels ad perf

* Improve ad perf testing

* Fix rational performance issues

* Improve FBM AD performance

* Fix FBM AD performance

* Test matern kernels peformance

* Optimise NeuralNetworkKernel

* Enable all nn tests

* linear and polynomial performance

* Uncomment nn test

* Add performance tests to basekernels

* Uncomment tests

* Finish up merge conflict resolution

* Test non-base-kernels

* Check transform performance

* Update runtests

* 1.6-compatible test_broken

* Failing tests on LTS

* Formatting

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Fix formatting

* Formatting

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Improve docstrings

* Apply suggestions from code review

Co-authored-by: Théo Galy-Fajou <[email protected]>

* Remove redundant only calls

* Tidy up docstring

* __example_inputs -> example_inputs

* Eltype in ones

* Sort ard promotion

* Bump patch

* Fix typo

* Factor out _to_colvecs

* Use struct rather than closue

* Test all rationalquadratickernel methods

* Formatting

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Uncomment tests

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Théo Galy-Fajou <[email protected]>

Author: 3 people

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.42"
+version = "0.10.43"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/TestUtils.jl

@@ -149,11 +149,13 @@ end
 """
     test_interface([rng::AbstractRNG], k::Kernel, ::Type{T}; kwargs...) where {T<:Real}
 
-Run the [`test_interface`](@ref) tests for randomly generated inputs of types `Vector{T}`, `Vector{Vector{T}}`, `ColVecs{T}`, and `RowVecs{T}`.
+Run the [`test_interface`](@ref) tests for randomly generated inputs of types `Vector{T}`,
+`Vector{Vector{T}}`, `ColVecs{T}`, and `RowVecs{T}`.
 
 For other input types, please provide the data manually.
 
-The keyword arguments are forwarded to the invocations of [`test_interface`](@ref) with the randomly generated inputs.
+The keyword arguments are forwarded to the invocations of [`test_interface`](@ref) with the
+randomly generated inputs.
 """
 function test_interface(rng::AbstractRNG, k::Kernel, ::Type{T}; kwargs...) where {T<:Real}
     @testset "Vector{$T}" begin
@@ -174,4 +176,27 @@ function test_interface(k::Kernel, T::Type{<:Real}=Float64; kwargs...)
     return test_interface(Random.GLOBAL_RNG, k, T; kwargs...)
 end
 
+"""
+    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))
+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))
+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))
+end
+
 end # module

src/basekernels/fbm.jl

@@ -50,8 +50,9 @@ _mod(x::RowVecs) = vec(sum(abs2, x.X; dims=2))
 
 function kernelmatrix(κ::FBMKernel, x::AbstractVector)
     modx = _mod(x)
+    modx_wide = modx * ones(eltype(modx), 1, length(modx)) # ad perf hack -- is unit tested
     modxx = pairwise(SqEuclidean(), x)
-    return _fbm.(modx, modx', modxx, κ.h)
+    return _fbm.(modx_wide, modx_wide', modxx, only(κ.h))
 end
 
 function kernelmatrix!(K::AbstractMatrix, κ::FBMKernel, x::AbstractVector)
@@ -63,7 +64,9 @@ end
 
 function kernelmatrix(κ::FBMKernel, x::AbstractVector, y::AbstractVector)
     modxy = pairwise(SqEuclidean(), x, y)
-    return _fbm.(_mod(x), _mod(y)', modxy, κ.h)
+    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
+    return _fbm.(modx_wide, mody_wide', modxy, only(κ.h))
 end
 
 function kernelmatrix!(
@@ -77,10 +80,10 @@ end
 function kernelmatrix_diag(κ::FBMKernel, x::AbstractVector)
     modx = _mod(x)
     modxx = colwise(SqEuclidean(), x)
-    return _fbm.(modx, modx, modxx, κ.h)
+    return _fbm.(modx, modx, modxx, only(κ.h))
 end
 
 function kernelmatrix_diag(κ::FBMKernel, x::AbstractVector, y::AbstractVector)
     modxy = colwise(SqEuclidean(), x, y)
-    return _fbm.(_mod(x), _mod(y), modxy, κ.h)
+    return _fbm.(_mod(x), _mod(y), modxy, only(κ.h))
 end

src/basekernels/nn.jl

@@ -37,6 +37,26 @@ function (κ::NeuralNetworkKernel)(x, y)
     return asin(dot(x, y) / sqrt((1 + sum(abs2, x)) * (1 + sum(abs2, y))))
 end
 
+function kernelmatrix(
+    k::NeuralNetworkKernel, x::AbstractVector{<:Real}, y::AbstractVector{<:Real}
+)
+    return kernelmatrix(k, _to_colvecs(x), _to_colvecs(y))
+end
+
+function kernelmatrix(k::NeuralNetworkKernel, x::AbstractVector{<:Real})
+    return kernelmatrix(k, _to_colvecs(x))
+end
+
+function kernelmatrix_diag(
+    k::NeuralNetworkKernel, x::AbstractVector{<:Real}, y::AbstractVector{<:Real}
+)
+    return kernelmatrix_diag(k, _to_colvecs(x), _to_colvecs(y))
+end
+
+function kernelmatrix_diag(k::NeuralNetworkKernel, x::AbstractVector{<:Real})
+    return kernelmatrix_diag(k, _to_colvecs(x))
+end
+
 function kernelmatrix(::NeuralNetworkKernel, x::ColVecs, y::ColVecs)
     validate_inputs(x, y)
     X_2 = sum(x.X .* x.X; dims=1)

src/basekernels/polynomial.jl

@@ -26,10 +26,28 @@ LinearKernel(; c::Real=0.0) = LinearKernel(c)
 
 @functor LinearKernel
 
-kappa(κ::LinearKernel, xᵀy::Real) = xᵀy + only(κ.c)
+__linear_kappa(c::Real, xᵀy::Real) = xᵀy + c
+
+kappa(κ::LinearKernel, xᵀy::Real) = __linear_kappa(only(κ.c), xᵀy)
 
 metric(::LinearKernel) = DotProduct()
 
+function kernelmatrix(k::LinearKernel, x::AbstractVector, y::AbstractVector)
+    return __linear_kappa.(only(k.c), pairwise(metric(k), x, y))
+end
+
+function kernelmatrix(k::LinearKernel, x::AbstractVector)
+    return __linear_kappa.(only(k.c), pairwise(metric(k), x))
+end
+
+function kernelmatrix_diag(k::LinearKernel, x::AbstractVector, y::AbstractVector)
+    return __linear_kappa.(only(k.c), colwise(metric(k), x, y))
+end
+
+function kernelmatrix_diag(k::LinearKernel, x::AbstractVector)
+    return __linear_kappa.(only(k.c), colwise(metric(k), x))
+end
+
 Base.show(io::IO, κ::LinearKernel) = print(io, "Linear Kernel (c = ", only(κ.c), ")")
 
 """
@@ -68,10 +86,32 @@ function Functors.functor(::Type{<:PolynomialKernel}, x)
     return (c=x.c,), reconstruct_polynomialkernel
 end
 
-kappa(κ::PolynomialKernel, xᵀy::Real) = (xᵀy + only(κ.c))^κ.degree
+struct _PolynomialKappa
+    degree::Int
+end
+
+(κ::_PolynomialKappa)(c::Real, xᵀy::Real) = (xᵀy + c)^κ.degree
+
+kappa(κ::PolynomialKernel, xᵀy::Real) = _PolynomialKappa(κ.degree)(only(κ.c), xᵀy)
 
 metric(::PolynomialKernel) = DotProduct()
 
+function kernelmatrix(k::PolynomialKernel, x::AbstractVector, y::AbstractVector)
+    return _PolynomialKappa(k.degree).(only(k.c), pairwise(metric(k), x, y))
+end
+
+function kernelmatrix(k::PolynomialKernel, x::AbstractVector)
+    return _PolynomialKappa(k.degree).(only(k.c), pairwise(metric(k), x))
+end
+
+function kernelmatrix_diag(k::PolynomialKernel, x::AbstractVector, y::AbstractVector)
+    return _PolynomialKappa(k.degree).(only(k.c), colwise(metric(k), x, y))
+end
+
+function kernelmatrix_diag(k::PolynomialKernel, x::AbstractVector)
+    return _PolynomialKappa(k.degree).(only(k.c), colwise(metric(k), x))
+end
+
 function Base.show(io::IO, κ::PolynomialKernel)
     return print(io, "Polynomial Kernel (c = ", only(κ.c), ", degree = ", κ.degree, ")")
 end

src/basekernels/rational.jl

@@ -31,12 +31,32 @@ end
 
 @functor RationalKernel
 
-function kappa(κ::RationalKernel, d::Real)
-    return (one(d) + d / only(κ.α))^(-only(κ.α))
-end
+__rational_kappa(α::Real, d::Real) = (one(d) + d / α)^(-α)
+
+kappa(κ::RationalKernel, d::Real) = __rational_kappa(only(κ.α), d)
 
 metric(k::RationalKernel) = k.metric
 
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix(k::RationalKernel, x::AbstractVector, y::AbstractVector)
+    return __rational_kappa.(only(k.α), pairwise(metric(k), x, y))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix(k::RationalKernel, x::AbstractVector)
+    return __rational_kappa.(only(k.α), pairwise(metric(k), x))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix_diag(k::RationalKernel, x::AbstractVector, y::AbstractVector)
+    return __rational_kappa.(only(k.α), colwise(metric(k), x, y))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix_diag(k::RationalKernel, x::AbstractVector)
+    return __rational_kappa.(only(k.α), colwise(metric(k), x))
+end
+
 function Base.show(io::IO, κ::RationalKernel)
     return print(io, "Rational Kernel (α = ", only(κ.α), ", metric = ", κ.metric, ")")
 end
@@ -69,18 +89,59 @@ struct RationalQuadraticKernel{Tα<:Real,M} <: SimpleKernel
     end
 end
 
+const _RQ_Euclidean = RationalQuadraticKernel{<:Real,<:Euclidean}
+
 @functor RationalQuadraticKernel
 
-function kappa(κ::RationalQuadraticKernel, d::Real)
-    return (one(d) + d^2 / (2 * only(κ.α)))^(-only(κ.α))
-end
-function kappa(κ::RationalQuadraticKernel{<:Real,<:Euclidean}, d²::Real)
-    return (one(d²) + d² / (2 * only(κ.α)))^(-only(κ.α))
-end
+__rq_kappa(α::Real, d::Real) = (one(d) + d^2 / (2 * α))^(-α)
+__rq_kappa_euclidean(α::Real, d²::Real) = (one(d²) + d² / (2 * α))^(-α)
+
+kappa(κ::RationalQuadraticKernel, d::Real) = __rq_kappa(only(κ.α), d)
+kappa(κ::_RQ_Euclidean, d²::Real) = __rq_kappa_euclidean(only(κ.α), d²)
 
 metric(k::RationalQuadraticKernel) = k.metric
 metric(::RationalQuadraticKernel{<:Real,<:Euclidean}) = SqEuclidean()
 
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix(k::RationalQuadraticKernel, x::AbstractVector, y::AbstractVector)
+    return __rq_kappa.(only(k.α), pairwise(metric(k), x, y))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix(k::RationalQuadraticKernel, x::AbstractVector)
+    return __rq_kappa.(only(k.α), pairwise(metric(k), x))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix_diag(k::RationalQuadraticKernel, x::AbstractVector, y::AbstractVector)
+    return __rq_kappa.(only(k.α), colwise(metric(k), x, y))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix_diag(k::RationalQuadraticKernel, x::AbstractVector)
+    return __rq_kappa.(only(k.α), colwise(metric(k), x))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix(k::_RQ_Euclidean, x::AbstractVector, y::AbstractVector)
+    return __rq_kappa_euclidean.(only(k.α), pairwise(SqEuclidean(), x, y))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix(k::_RQ_Euclidean, x::AbstractVector)
+    return __rq_kappa_euclidean.(only(k.α), pairwise(SqEuclidean(), x))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix_diag(k::_RQ_Euclidean, x::AbstractVector, y::AbstractVector)
+    return __rq_kappa_euclidean.(only(k.α), colwise(SqEuclidean(), x, y))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix_diag(k::_RQ_Euclidean, x::AbstractVector)
+    return __rq_kappa_euclidean.(only(k.α), colwise(SqEuclidean(), x))
+end
+
 function Base.show(io::IO, κ::RationalQuadraticKernel)
     return print(
         io, "Rational Quadratic Kernel (α = ", only(κ.α), ", metric = ", κ.metric, ")"
@@ -121,12 +182,32 @@ end
 
 @functor GammaRationalKernel
 
-function kappa(κ::GammaRationalKernel, d::Real)
-    return (one(d) + d^only(κ.γ) / only(κ.α))^(-only(κ.α))
-end
+__grk_kappa(α::Real, γ::Real, d::Real) = (one(d) + d^γ / α)^(-α)
+
+kappa(κ::GammaRationalKernel, d::Real) = __grk_kappa(only(κ.α), only(κ.γ), d)
 
 metric(k::GammaRationalKernel) = k.metric
 
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix(k::GammaRationalKernel, x::AbstractVector, y::AbstractVector)
+    return __grk_kappa.(only(k.α), only(k.γ), pairwise(metric(k), x, y))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix(k::GammaRationalKernel, x::AbstractVector)
+    return __grk_kappa.(only(k.α), only(k.γ), pairwise(metric(k), x))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix_diag(k::GammaRationalKernel, x::AbstractVector, y::AbstractVector)
+    return __grk_kappa.(only(k.α), only(k.γ), colwise(metric(k), x, y))
+end
+
+# AD-performance optimisation. Is unit tested.
+function kernelmatrix_diag(k::GammaRationalKernel, x::AbstractVector)
+    return __grk_kappa.(only(k.α), only(k.γ), colwise(metric(k), x))
+end
+
 function Base.show(io::IO, κ::GammaRationalKernel)
     return print(
         io,

src/kernels/normalizedkernel.jl

@@ -20,15 +20,17 @@ end
 (κ::NormalizedKernel)(x, y) = κ.kernel(x, y) / sqrt(κ.kernel(x, x) * κ.kernel(y, y))
 
 function kernelmatrix(κ::NormalizedKernel, x::AbstractVector, y::AbstractVector)
-    return kernelmatrix(κ.kernel, x, y) ./
-           sqrt.(
-        kernelmatrix_diag(κ.kernel, x) .* permutedims(kernelmatrix_diag(κ.kernel, y))
-    )
+    x_diag = kernelmatrix_diag(κ.kernel, x)
+    x_diag_wide = x_diag * ones(eltype(x_diag), 1, length(y)) # ad perf hack. Is unit tested
+    y_diag = kernelmatrix_diag(κ.kernel, y)
+    y_diag_wide = y_diag * ones(eltype(y_diag), 1, length(x)) # ad perf hack. Is unit tested
+    return kernelmatrix(κ.kernel, x, y) ./ sqrt.(x_diag_wide .* y_diag_wide')
 end
 
 function kernelmatrix(κ::NormalizedKernel, x::AbstractVector)
     x_diag = kernelmatrix_diag(κ.kernel, x)
-    return kernelmatrix(κ.kernel, x) ./ sqrt.(x_diag .* permutedims(x_diag))
+    x_diag_wide = x_diag * ones(eltype(x_diag), 1, length(x)) # ad perf hack. Is unit tested
+    return kernelmatrix(κ.kernel, x) ./ sqrt.(x_diag_wide .* x_diag_wide')
 end
 
 function kernelmatrix_diag(κ::NormalizedKernel, x::AbstractVector)

src/kernels/scaledkernel.jl

@@ -26,19 +26,19 @@ end
 (k::ScaledKernel)(x, y) = only(k.σ²) * k.kernel(x, y)
 
 function kernelmatrix(κ::ScaledKernel, x::AbstractVector, y::AbstractVector)
-    return κ.σ² .* kernelmatrix(κ.kernel, x, y)
+    return only(κ.σ²) * kernelmatrix(κ.kernel, x, y)
 end
 
 function kernelmatrix(κ::ScaledKernel, x::AbstractVector)
-    return κ.σ² .* kernelmatrix(κ.kernel, x)
+    return only(κ.σ²) * kernelmatrix(κ.kernel, x)
 end
 
 function kernelmatrix_diag(κ::ScaledKernel, x::AbstractVector)
-    return κ.σ² .* kernelmatrix_diag(κ.kernel, x)
+    return only(κ.σ²) * kernelmatrix_diag(κ.kernel, x)
 end
 
 function kernelmatrix_diag(κ::ScaledKernel, x::AbstractVector, y::AbstractVector)
-    return κ.σ² .* kernelmatrix_diag(κ.kernel, x, y)
+    return only(κ.σ²) * kernelmatrix_diag(κ.kernel, x, y)
 end
 
 function kernelmatrix!(

src/transform/ardtransform.jl

@@ -35,9 +35,14 @@ dim(t::ARDTransform) = length(t.v)
 (t::ARDTransform)(x::Real) = only(t.v) * x
 (t::ARDTransform)(x) = t.v .* x
 
-_map(t::ARDTransform, x::AbstractVector{<:Real}) = t.v' .* x
-_map(t::ARDTransform, x::ColVecs) = ColVecs(t.v .* x.X)
-_map(t::ARDTransform, x::RowVecs) = RowVecs(t.v' .* x.X)
+# Quite specific implementations required to pass correctness and performance tests.
+_map(t::ARDTransform, x::AbstractVector{<:Real}) = x * only(t.v)
+function _map(t::ARDTransform, x::ColVecs)
+    return ColVecs((t.v * ones(eltype(t.v), 1, size(x.X, 2))) .* x.X)
+end
+function _map(t::ARDTransform, x::RowVecs)
+    return RowVecs(x.X .* (ones(eltype(t.v), size(x.X, 1)) * collect(t.v')))
+end
 
 Base.isequal(t::ARDTransform, t2::ARDTransform) = isequal(t.v, t2.v)
 

src/utils.jl

@@ -97,6 +97,8 @@ Base.zero(x::ColVecs) = ColVecs(zero(x.X))
 
 dim(x::ColVecs) = size(x.X, 1)
 
+_to_colvecs(x::AbstractVector{<:Real}) = ColVecs(reshape(x, 1, :))
+
 pairwise(d::PreMetric, x::ColVecs) = Distances_pairwise(d, x.X; dims=2)
 pairwise(d::PreMetric, x::ColVecs, y::ColVecs) = Distances_pairwise(d, x.X, y.X; dims=2)
 function pairwise(d::PreMetric, x::AbstractVector, y::ColVecs)

test/Project.toml

@@ -13,6 +13,7 @@ PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"