feat: customizable DIT benchmarks (#636)

* feat: allow custom max duration in DIT benchmarks

* feat: customize benchmark aggregation

* Constant fix for Lux

Author: gdalle

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestLuxExt/DifferentiationInterfaceTestLuxExt.jl

@@ -181,7 +181,7 @@ function DIT.lux_scenarios(rng::AbstractRNG=default_rng())
         scen = Scenario{:gradient,:out}(
             square_loss,
             ComponentArray(ps);
-            contexts=(Constant(model), Constant(x), Constant(st)),
+            contexts=(DI.Constant(model), DI.Constant(x), DI.Constant(st)),
             res1=g,
         )
         push!(scens, scen)

DifferentiationInterfaceTest/src/test_differentiation.jl

@@ -56,6 +56,8 @@ Each setting tests/benchmarks a different subset of calls:
 
 - `count_calls=true`: whether to also count function calls during benchmarking
 - `benchmark_test=true`: whether to include tests which succeed iff benchmark doesn't error
+- `benchmark_seconds=1`: how long to run each benchmark for
+- `benchmark_aggregation=minimum`: function used to aggregate sample measurements
 """
 function test_differentiation(
     backends::Vector{<:AbstractADType},
@@ -87,6 +89,8 @@ function test_differentiation(
     # benchmark options
     count_calls::Bool=true,
     benchmark_test::Bool=true,
+    benchmark_seconds::Real=1,
+    benchmark_aggregation=minimum,
 )
     @assert type_stability in (:none, :prepared, :full)
     @assert allocations in (:none, :prepared, :full)
@@ -173,6 +177,8 @@ function test_differentiation(
                             subset=benchmark,
                             count_calls,
                             benchmark_test,
+                            benchmark_seconds,
+                            benchmark_aggregation,
                         )
                     end
                     yield()
@@ -211,6 +217,8 @@ function benchmark_differentiation(
     logging::Bool=false,
     count_calls::Bool=true,
     benchmark_test::Bool=true,
+    benchmark_seconds::Real=1,
+    benchmark_aggregation=minimum,
 )
     return test_differentiation(
         backends,
@@ -223,5 +231,7 @@ function benchmark_differentiation(
         excluded,
         count_calls,
         benchmark_test,
+        benchmark_seconds,
+        benchmark_aggregation,
     )
 end

DifferentiationInterfaceTest/src/tests/benchmark.jl

@@ -56,7 +56,7 @@ $(TYPEDFIELDS)
 
 See the documentation of [Chairmarks.jl](https://github.com/LilithHafner/Chairmarks.jl) for more details on the measurement fields.
 """
-Base.@kwdef struct DifferentiationBenchmarkDataRow
+Base.@kwdef struct DifferentiationBenchmarkDataRow{T}
     "backend used for benchmarking"
     backend::AbstractADType
     "scenario used for benchmarking"
@@ -71,16 +71,16 @@ Base.@kwdef struct DifferentiationBenchmarkDataRow
     samples::Int
     "number of evaluations used for averaging in each sample"
     evals::Int
-    "minimum runtime over all samples, in seconds"
-    time::Float64
-    "minimum number of allocations over all samples"
-    allocs::Float64
-    "minimum memory allocated over all samples, in bytes"
-    bytes::Float64
-    "minimum fraction of time spent in garbage collection over all samples, between 0.0 and 1.0"
-    gc_fraction::Float64
-    "minimum fraction of time spent compiling over all samples, between 0.0 and 1.0"
-    compile_fraction::Float64
+    "aggregated runtime over all samples, in seconds"
+    time::T
+    "aggregated number of allocations over all samples"
+    allocs::T
+    "aggregated memory allocated over all samples, in bytes"
+    bytes::T
+    "aggregated fraction of time spent in garbage collection over all samples, between 0.0 and 1.0"
+    gc_fraction::T
+    "aggregated fraction of time spent compiling over all samples, between 0.0 and 1.0"
+    compile_fraction::T
 end
 
 function record!(
@@ -91,21 +91,22 @@ function record!(
     prepared::Union{Nothing,Bool},
     bench::Benchmark,
     calls::Integer,
+    aggregation,
 )
-    bench_min = minimum(bench)
+    bench_agg = aggregation(bench)
     row = DifferentiationBenchmarkDataRow(;
         backend=backend,
         scenario=scenario,
         operator=Symbol(operator),
         prepared=prepared,
         calls=calls,
         samples=length(bench.samples),
-        evals=Int(bench_min.evals),
-        time=bench_min.time,
-        allocs=bench_min.allocs,
-        bytes=bench_min.bytes,
-        gc_fraction=bench_min.gc_fraction,
-        compile_fraction=bench_min.compile_fraction,
+        evals=Int(bench_agg.evals),
+        time=bench_agg.time,
+        allocs=bench_agg.allocs,
+        bytes=bench_agg.bytes,
+        gc_fraction=bench_agg.gc_fraction,
+        compile_fraction=bench_agg.compile_fraction,
     )
     return push!(data, row)
 end