Better tools for benchmarking

[yebai]

Can have a utility that integrates better with Julia benchmarking tools ?

using Chairmarks, DynamicPPL

model_expr = to_benchmark_expr(posterior, AutoForwardDiff())

@be $model_expr

[penelopeysm]

I think it's a bit tricky because there's quite a lot of setup needed here, I think it would be doable if we limit ourselves to Chairmarks, but I think there's no way to retain compatibility with BenchmarkTools as well.

[penelopeysm]

I like the idea though, we could also make to_benchmark_expr(model) benchmark evaluation and to_benchmark_expr(model, adtype) benchmark gradients

[penelopeysm]

Actually I'm afraid this is a no-go because $ in @be interpolates values rather than expressions 😢

julia> expr = :(sum(randn(10)))
:(sum(randn(10)))

julia> @be sum(randn(10))
Benchmark: 3091 samples with 295 evaluations
 min    88.417 ns (2 allocs: 144 bytes)
 median 91.244 ns (2 allocs: 144 bytes)
 mean   103.663 ns (2 allocs: 144 bytes, 0.39% gc time)
 max    11.975 μs (2 allocs: 144 bytes, 98.32% gc time)

julia> @be $expr
ERROR: MethodError: objects of type Expr are not callable
The object of type `Expr` exists, but no method is defined for this combination of argument types when trying to treat it as a callable object.
Stacktrace:
  [1] _benchmark_3(::Expr, ::Int64, ::Bool)
    @ Chairmarks ~/.julia/packages/Chairmarks/yJzRJ/src/benchmarking.jl:145
  [2] (::Chairmarks.var"#29#31"{Nothing, Bool, Int64, Bool, Tuple{…}, Tuple{}, Base.RefValue{…}, Tuple{…}})(i::Int64)
    @ Chairmarks ~/.julia/packages/Chairmarks/yJzRJ/src/benchmarking.jl:124
  [3] ntuple
    @ ./ntuple.jl:19 [inlined]
  [4] _benchmark_2(args1::Any, setup::Any, teardown::Any, gc::Bool, evals::Int64, warmup::Bool, fs::Any)
    @ Chairmarks ~/.julia/packages/Chairmarks/yJzRJ/src/benchmarking.jl:121
  [5] _benchmark_1(init::Any, setup::Any, teardown::Any, evals::Union{…}, samples::Union{…}, seconds::Union{…}, gc::Bool, fs::Any)
    @ Chairmarks ~/.julia/packages/Chairmarks/yJzRJ/src/benchmarking.jl:40
  [6] benchmark(init::Nothing, setup::Nothing, fs::Tuple{…}, teardown::Nothing; evals::Nothing, samples::Nothing, seconds::Float64, gc::Bool)
    @ Chairmarks ~/.julia/packages/Chairmarks/yJzRJ/src/benchmarking.jl:17
  [7] benchmark(init::Nothing, setup::Nothing, fs::Tuple{Expr}, teardown::Nothing)
    @ Chairmarks ~/.julia/packages/Chairmarks/yJzRJ/src/benchmarking.jl:12
  [8] benchmark(init::Nothing, setup::Nothing, f::Expr, teardown::Nothing; kw::@Kwargs{})
    @ Chairmarks ~/.julia/packages/Chairmarks/yJzRJ/src/benchmarking.jl:5
  [9] benchmark(setup::Nothing, f::Expr, teardown::Nothing; kw::@Kwargs{})
    @ Chairmarks ~/.julia/packages/Chairmarks/yJzRJ/src/benchmarking.jl:4
 [10] benchmark(f::Expr; kw::@Kwargs{})
    @ Chairmarks ~/.julia/packages/Chairmarks/yJzRJ/src/benchmarking.jl:3
 [11] top-level scope
    @ REPL[13]:1
Some type information was truncated. Use `show(err)` to see complete types.

[penelopeysm]

we could write our own macro, that would work: so something like

@model_be model
@model_be model adtype

would be quite easy to write; but I'm not sure it would qualify as 'integrating with existing packages'

[yebai]

That's unfortunate. One possible alternative is

using Chairmarks, DynamicPPL

f = @to_benchmark(posterior, AutoForwardDiff()) # f evaluates log density and gradient

@be f()

In general, I think it is easier to allow one to directly use benchmarking tools intead of wrapping around them.

[yebai]

(re-opened so we can keep discussing this)

[penelopeysm]

You can't do that, though, because the LogDensityFunction needs to be set up before the benchmark is run

@be (ldfsetup, paramsetup) LogDensityProblems.logdensity(_[1], _[2])

otherwise the benchmarking of f also includes the setup time.

[yebai]

I see, thanks. Maybe this works?

(ldfsetup, paramsetup) = @to_benchmark(posterior, AutoForwardDiff(); parameters = optionally_something) 

@be (ldfsetup, paramsetup) LogDensityProblems.logdensity(_[1], _[2])

[penelopeysm]

I think that runs into the same issue that you can't interpolate Expr. What you're suggesting isn't very different from this though:

ldf, params = make_ldf_and_params(model, adtype)
@be LogDensityProblems.logdensity($ldf, $params)

[penelopeysm]

the necessary definition of make_ldf_and_params is mostly the same as in run_ad -- so we could extract that into a separate function quite easily

DynamicPPL.jl/src/test_utils/ad.jl

Lines 260 to 274 in 06ab810

	# Convert Boolean `test` to an AbstractADCorrectnessTestSetting
	if test isa Bool
	test = test ? WithBackend() : NoTest()
	end
	# Extract parameters
	if isnothing(params)
	params = varinfo[:]
	end
	params = [p for p in params] # Concretise
	# Calculate log-density and gradient with the backend of interest
	verbose && @info "Running AD on $(model.f) with $(adtype)\n"
	verbose && println(" params : $(params)")
	ldf = LogDensityFunction(model, getlogdensity, varinfo; adtype=adtype)

[yebai]

ldf, params = make_ldf_and_params(model, adtype)
@be LogDensityProblems.logdensity($ldf, $params)

I like the idea -- yes, a utility to setup ldf and params would be very helpful.

[penelopeysm]

Could probably also be used in Mooncake test suite!