diff --git a/benchmark/Project.toml b/benchmark/Project.toml index cf9c7b9c1..aaa7b4a0c 100644 --- a/benchmark/Project.toml +++ b/benchmark/Project.toml @@ -1,3 +1,11 @@ [deps] BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392" +LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[compat] +BenchmarkTools = "1" +ForwardDiff = "0.10" +Zygote = "0.6" diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 0cb80e5ce..cb9d59f22 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -1,5 +1,8 @@ using BenchmarkTools using KernelFunctions +using LogExpFunctions: logistic +using Zygote +using ForwardDiff N1 = 10 N2 = 20 @@ -16,31 +19,83 @@ Yv = collect.(eachcol(Y)) # Create the general suite of benchmarks SUITE = BenchmarkGroup() +# Create a list of kernel and their constructors kernels = Dict( - "SqExponential" => SqExponentialKernel(), "Exponential" => ExponentialKernel() + # Constant Kernels + "Constant" => ((2.0,), x -> ConstantKernel(; c=x)), + "White" => ((), () -> WhiteKernel()), + # Cosine Kernel + "Cosine" => ((), () -> CosineKernel()), + # Exponential Kernels + "Exponential" => ((), () -> ExponentialKernel()), + # "Gibbs" => ((), () -> GibbsKernel(; lengthscale=x -> sin.(x))), + "SqExponential" => ((), () -> SqExponentialKernel()), + "GammaExponential" => ((1.0,), x -> GammaExponentialKernel(; γ=2 * logistic(x))), + # Exponentiated Kernel + "Exponentiated" => ((), () -> ExponentiatedKernel()), ) inputtypes = Dict("ColVecs" => (Xc, Yc), "RowVecs" => (Xr, Yr), "Vecs" => (Xv, Yv)) functions = Dict( - "kernelmatrixX" => (kernel, X, Y) -> kernelmatrix(kernel, X), - "kernelmatrixXY" => (kernel, X, Y) -> kernelmatrix(kernel, X, Y), - "kernelmatrix_diagX" => (kernel, X, Y) -> kernelmatrix_diag(kernel, X), - "kernelmatrix_diagXY" => (kernel, X, Y) -> kernelmatrix_diag(kernel, X, Y), + "kernelmatrixX" => (fk, args, X, Y) -> kernelmatrix(fk(args...), X), + "kernelmatrixXY" => (fk, args, X, Y) -> kernelmatrix(fk(args...), X, Y), + "kernelmatrix_diagX" => (fk, args, X, Y) -> kernelmatrix_diag(fk(args...), X), + "kernelmatrix_diagXY" => (fk, args, X, Y) -> kernelmatrix_diag(fk(args...), X, Y), ) -for (kname, kernel) in kernels - SUITE[kname] = sk = BenchmarkGroup() +# Test the allocated functions +SUITE["Allocated Functions"] = suite_alloc = BenchmarkGroup() +for (kname, (kargs, kf)) in kernels + suite_alloc[kname] = suite_kernel = BenchmarkGroup() for (inputname, (X, Y)) in inputtypes - sk[inputname] = si = BenchmarkGroup() + suite_kernel[inputname] = suite_input = BenchmarkGroup() for (fname, f) in functions - si[fname] = @benchmarkable $f($kernel, $X, $Y) + suite_input[fname] = @benchmarkable $f($kf, $kargs, $X, $Y) end end end +# Test the AD frameworks +## Zygote +SUITE["Zygote"] = suite_zygote = BenchmarkGroup() +for (kname, (kargs, kf)) in kernels + suite_zygote[kname] = suite_kernel = BenchmarkGroup() + for (inputname, (X, Y)) in inputtypes + suite_kernel[inputname] = suite_input = BenchmarkGroup() + for (fname, f) in functions + # Forward-pass + suite_input[fname * "_forward"] = @benchmarkable Zygote.pullback( + $kargs, $X, $Y + ) do args, x, y + $f($kf, args, x, y) + end + # Reverse pass + out, pb = Zygote.pullback(kargs, X, Y) do args, x, y + f(kf, args, x, y) + end + suite_input[fname * "_reverse"] = @benchmarkable $pb($out) + end + end +end + +## ForwardDiff +# Right now there is no canonical way to turn (kargs, X, Y) into an array. +# SUITE["ForwardDiff"] = suite_forwarddiff = BenchmarkGroup() +# for (kname, (kargs, kf)) in kernels +# suite_forwarddiff[kname] = suite_kernel = BenchmarkGroup() +# for (inputname, (X, Y)) in inputtypes +# suite_kernel[inputname] = suite_input = BenchmarkGroup() +# for (fname, f) in functions +# suite_input[fname] = @benchmarkable ForwardDiff.gradient($kargs, $X, $Y) do args, x, y +# $f($kf, args, x, y) +# end +# end +# end +# end + # Uncomment the following to run benchmark locally # tune!(SUITE) -# results = run(SUITE, verbose=true) +# results = run(SUITE; verbose=true)