Add more kernels to the benchmark and AD benchmarks

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"

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)