SimpleKernel performance with Zygote

It looks like we've got a performance bug in the above. Specifically, the primal looks to be fine:

julia> @benchmark kernelmatrix(SEKernel(), $(randn(100)))
BenchmarkTools.Trial:
  memory estimate:  158.25 KiB
  allocs estimate:  8
  --------------
  minimum time:     92.684 μs (0.00% GC)
  median time:      155.135 μs (0.00% GC)
  mean time:        181.207 μs (12.32% GC)
  maximum time:     15.995 ms (98.92% GC)
  --------------
  samples:          10000
  evals/sample:     1

while Zygotes forwards-pass is about 100x worse:

julia> @benchmark Zygote.pullback(kernelmatrix, SEKernel(), $(randn(100)))
BenchmarkTools.Trial:
  memory estimate:  7.25 MiB
  allocs estimate:  290068
  --------------
  minimum time:     34.106 ms (0.00% GC)
  median time:      35.180 ms (0.00% GC)
  mean time:        36.456 ms (2.37% GC)
  maximum time:     60.521 ms (25.00% GC)
  --------------
  samples:          138
  evals/sample:     1

Given that the bar for a reasonably performant forwards-pass is 1-2x the primal, this is definitely sub-optimal.

I'm reasonably sure that the culprit is the use of Base.Fix1 here: https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/blob/13985cc6cb8903004f33e6fe9d2540571594a3c8/src/matrix/kernelmatrix.jl#L95

maping interesting types in is generally a bad idea when used in conjunction with Zygote unfortunately.

@theogf @devmotion I'm not sure what our best options are here. I think there are two questions here:

1. could we tweak the current implementation a bit into something that Zygote likes, and get sane performance?
2. what escape hatches do we have / could we have to make it straightforward to hand-improve performance where necessary?

For example, the use of Base.Fix1 here might actually be a real win, because we could hand-optimise the implementation of map(::Base.Fix1{typeof(kappa), SEKernel}, ::Array{<:Real}) quite straightforwardly. Unfortunately, this would essentially mean that any kernel that we care about having good performance with would require us to hand-implement stuff. If we did this using ForwardDiff, it might be fine though 🤷

n.b. it's definitely not the pairwise computations:

julia> @benchmark Zygote.pullback(KernelFunctions.pairwise, KernelFunctions.SqEuclidean(), randn(100))
BenchmarkTools.Trial:
  memory estimate:  82.25 KiB
  allocs estimate:  45
  --------------
  minimum time:     28.473 μs (0.00% GC)
  median time:      59.303 μs (0.00% GC)
  mean time:        74.604 μs (15.96% GC)
  maximum time:     15.654 ms (99.13% GC)
  --------------
  samples:          10000
  evals/sample:     1

they've been hand-optimised for ages, so it would have been really surprising if there were a problem there.

Start in src/matrix/kernelmatrix.jl at the linked line using Base.Fix1, and reproduce the Julia benchmarks for kernelmatrix and Zygote.pullback. Compare the forward-pass behavior with the pairwise benchmark to isolate the regression. Done means documenting or implementing an approach that brings the Zygote forward pass substantially closer to the primal without regressing the existing kernel path.