Benchmark results for Mooncake

[yebai]

Thanks for the nice work!

Is it possible to include the benchmarking results for Mooncake too? It is not yet performance-optimal, but @willtebbutt is working on it.

[nsiccha]

@yebai and @willtebbutt, should be live now at https://nsiccha.github.io/StanBlocks.jl/performance.html#visualization
Could either of you confirm that this is how I should benchmark Mooncake?

Setting up function at the very beginning:

StanBlocks.jl/docs/julia/common.jl

Lines 174 to 181 in f0840d5

	compute_property(e, ::Val{:mooncake!}) = begin
	(;julia_lpdf) = e
	rule = Mooncake.build_rrule(julia_lpdf, randn(e.dimension))
	mooncake_lpdf = Mooncake.CoDual(julia_lpdf, zero_tangent(julia_lpdf))
	((g,x),)->(Mooncake.__value_and_gradient!!(
	rule, mooncake_lpdf, Mooncake.CoDual(x, g)
	)[1], g)
	end

Benchmarking:

StanBlocks.jl/docs/julia/common.jl

Line 191 in f0840d5

compute_property(e, ::Val{:mooncake_benchmark}) = (@be (randn(e.dimension),zeros(e.dimension)) e.mooncake!)

[nsiccha]

Now that I see it like that, maybe I have to interpolate e.mooncake!? I'll check out whether that makes a difference tomorrow, but the results are roughly in line with what I had observed before, and the non-interpolated version seems to work well for all the other benchmarks as far as I can tell.

[willtebbutt]

Okay. I think you're benchmarking Mooncake in a reasonable manner. I've done some performance profiling locally and, at least on the first example, there's some really weird performance issues (the kind of thing that I wouldn't expect to see). I'm not entirely sure what's going on yet, but I suspect it's going to involve a fix on Mooncake's end. Will keep you apprised.

[nsiccha]

Great! The fact that it's possible to get 0 allocations both for plain Julia, enzyme and mooncake using that approach also indicates to me that it's fine the way I do it. Will close this issue then, feel free to open another issue if I should change or update+rerun things.