feat: allow custom max duration in DIT benchmarks

JuliaDiff · Nov 23, 2024 · 655154d · 655154d
1 parent 2743e18
commit 655154d
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 84 deletions.
diff --git a/DifferentiationInterfaceTest/src/test_differentiation.jl b/DifferentiationInterfaceTest/src/test_differentiation.jl
@@ -56,6 +56,7 @@ 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
 """
 function test_differentiation(
     backends::Vector{<:AbstractADType},
@@ -87,6 +88,7 @@ function test_differentiation(
     # benchmark options
     count_calls::Bool=true,
     benchmark_test::Bool=true,
+    benchmark_seconds::Real=1,
 )
     @assert type_stability in (:none, :prepared, :full)
     @assert allocations in (:none, :prepared, :full)
@@ -173,6 +175,7 @@ function test_differentiation(
                             subset=benchmark,
                             count_calls,
                             benchmark_test,
+                            benchmark_seconds,
                         )
                     end
                     yield()
@@ -211,6 +214,7 @@ function benchmark_differentiation(
     logging::Bool=false,
     count_calls::Bool=true,
     benchmark_test::Bool=true,
+    benchmark_seconds::Real=1,
 )
     return test_differentiation(
         backends,
@@ -223,5 +227,6 @@ function benchmark_differentiation(
         excluded,
         count_calls,
         benchmark_test,
+        benchmark_seconds,
     )
 end