Refactor benchmark

Author: MilesCranmer

benchmark/benchmarks.jl

@@ -10,55 +10,88 @@ function benchmark_evaluation()
     operators = OperatorEnum(;
         binary_operators=[+, -, /, *], unary_operators=[cos, exp], enable_autodiff=true
     )
-    for T in (ComplexF32, ComplexF64, Float32, Float64)
-        if !(T <: Real) && PACKAGE_VERSION < v"0.5.0" && PACKAGE_VERSION != v"0.0.0"
-            continue
-        end
-        suite[T] = BenchmarkGroup()
 
-        n = 1_000
+    config_options = [
+        [
+            (turbo=turbo, T=T, n=n, derivative=derivative) for turbo in (false, true) for
+            T in (ComplexF32, ComplexF64, Float32, Float64) for n in (100, 1_000, 10_000)
+            for derivative in (false, true)
+        ]...,
+    ]
 
-        #! format: off
-        for turbo in (false, true)
-            if turbo && !(T in (Float32, Float64))
-                continue
-            end
-            extra_key = turbo ? "_turbo" : ""
+    config_options = filter!(config_options) do config
+        !(config.T <: Real) &&
+            PACKAGE_VERSION < v"0.5.0" &&
+            PACKAGE_VERSION != v"0.0.0" &&
+            return false
+
+        config.turbo && !(config.T in (Float32, Float64)) && return false
+
+        config.T != Float32 && config.n != 1_000 && return false
+
+        config.T != Float32 && config.derivative && return false
+
+        return true
+    end
+
+    for config in config_options
+        T = config.T
+        turbo = config.turbo
+        n = config.n
+        derivative = config.derivative
+
+        derivative_s = derivative ? "derivative" : "evaluation"
+        turbo_s = turbo ? "turbo" : "standard"
+
+        haskey(suite, derivative_s) || (suite[derivative_s] = BenchmarkGroup())
+        haskey(suite[derivative_s], T) || (suite[derivative_s][T] = BenchmarkGroup())
+        haskey(suite[derivative_s][T], n) || (suite[derivative_s][T][n] = BenchmarkGroup())
+        haskey(suite[derivative_s][T][n], turbo_s) ||
+            (suite[derivative_s][T][n][turbo_s] = BenchmarkGroup())
+
+        if derivative
+            eval_grad_tree_array(
+                gen_random_tree_fixed_size(20, operators, 5, T),
+                randn(MersenneTwister(0), T, 5, n),
+                operators;
+                variable=true,
+                turbo=turbo,
+            )
+            suite[derivative_s][T][n][turbo_s] = @benchmarkable(
+                [
+                    eval_grad_tree_array(tree, X, $operators; variable=true, turbo=$turbo)
+                    for tree in trees
+                ],
+                setup = (
+                    X = randn(MersenneTwister(0), $T, 5, $n);
+                    treesize = 20;
+                    ntrees = 100;
+                    trees = [
+                        gen_random_tree_fixed_size(treesize, $operators, 5, $T) for
+                        _ in 1:ntrees
+                    ]
+                )
+            )
+        else
             eval_tree_array(
                 gen_random_tree_fixed_size(20, operators, 5, T),
                 randn(MersenneTwister(0), T, 5, n),
                 operators;
-                turbo=turbo
+                turbo=turbo,
             )
-            suite[T]["evaluation$(extra_key)"] = @benchmarkable(
+            suite[derivative_s][T][n][turbo_s] = @benchmarkable(
                 [eval_tree_array(tree, X, $operators; turbo=$turbo) for tree in trees],
-                setup=(
-                    X=randn(MersenneTwister(0), $T, 5, $n);
-                    treesize=20;
-                    ntrees=100;
-                    trees=[gen_random_tree_fixed_size(treesize, $operators, 5, $T) for _ in 1:ntrees]
+                setup = (
+                    X = randn(MersenneTwister(0), $T, 5, $n);
+                    treesize = 20;
+                    ntrees = 100;
+                    trees = [
+                        gen_random_tree_fixed_size(treesize, $operators, 5, $T) for
+                        _ in 1:ntrees
+                    ]
                 )
             )
-            if T <: Real
-                eval_grad_tree_array(
-                    gen_random_tree_fixed_size(20, operators, 5, T),
-                    randn(MersenneTwister(0), T, 5, n),
-                    operators;
-                    variable=true,
-                    turbo=turbo
-                )
-                suite[T]["derivative$(extra_key)"] = @benchmarkable(
-                    [eval_grad_tree_array(tree, X, $operators; variable=true, turbo=$turbo) for tree in trees],
-                    setup=(
-                        X=randn(MersenneTwister(0), $T, 5, $n);
-                        treesize=20;
-                        ntrees=100;
-                        trees=[gen_random_tree_fixed_size(treesize, $operators, 5, $T) for _ in 1:ntrees]
-                    )
-                )
-            end
         end
-        #! format: on
     end
     return suite
 end