Add utilities that don't return the value

.gitignore

@@ -12,3 +12,5 @@
 /docs/Manifest.toml
 /test/Manifest.toml
 /benchmark/Manifest.toml
+
+*.csv

benchmark/Project.toml

@@ -1,5 +1,7 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
@@ -8,4 +10,5 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

benchmark/benchmarks.jl

@@ -66,145 +66,89 @@ reverse_fallback_backends = [
     ReverseDiffBackend(; custom=false),
 ]
 
-all_custom_backends = vcat(forward_custom_backends, reverse_custom_backends)
-all_fallback_backends = vcat(forward_fallback_backends, reverse_fallback_backends)
-all_backends = vcat(all_custom_backends, all_fallback_backends)
+all_backends = vcat(
+    forward_custom_backends,
+    forward_fallback_backends,
+    reverse_custom_backends,
+    reverse_fallback_backends,
+)
 
 ## Suite
 
-SUITE = BenchmarkGroup()
+function make_suite()
+    SUITE = BenchmarkGroup()
 
-### Scalar to scalar
-
-scalar_to_scalar = Layer(randn(), randn(), tanh)
-
-for backend in all_backends
-    handles_types(backend, Number, Number) || continue
-    SUITE["value_and_derivative"][(1, 1)][string(backend)] = @benchmarkable begin
-        value_and_derivative($backend, $scalar_to_scalar, x)
-    end setup = (x = randn())
-end
-
-for backend in all_fallback_backends
-    handles_types(backend, Number, Number) || continue
-    if autodiff_mode(backend) == :forward
-        SUITE["value_and_pushforward"][(1, 1)][string(backend)] = @benchmarkable begin
-            value_and_pushforward($backend, $scalar_to_scalar, x, dx)
-        end setup = (x = randn(); dx = randn())
-    else
-        SUITE["value_and_pullback"][(1, 1)][string(backend)] = @benchmarkable begin
-            value_and_pullback($backend, $scalar_to_scalar, x, dy)
-        end setup = (x = randn(); dy = randn())
-    end
-end
-
-### Scalar to vector
-
-for m in [10]
-    scalar_to_vector = Layer(randn(m), randn(m), tanh)
+    ### Scalar to scalar
+    scalar_to_scalar = Layer(randn(), randn(), tanh)
 
     for backend in all_backends
-        handles_types(backend, Number, Vector) || continue
-        SUITE["value_and_multiderivative"][(1, m)][string(backend)] = @benchmarkable begin
-            value_and_multiderivative($backend, $scalar_to_vector, x)
-        end setup = (x = randn())
-        SUITE["value_and_multiderivative!"][(1, m)][string(backend)] = @benchmarkable begin
-            value_and_multiderivative!(multider, $backend, $scalar_to_vector, x)
-        end setup = (x = randn(); multider = zeros($m))
+        add_derivative_benchmarks!(SUITE, backend, scalar_to_scalar, 1, 1)
     end
-
-    for backend in all_fallback_backends
-        handles_types(backend, Number, Vector) || continue
-        if autodiff_mode(backend) == :forward
-            SUITE["value_and_pushforward"][(1, m)][string(backend)] = @benchmarkable begin
-                value_and_pushforward($backend, $scalar_to_vector, x, dx)
-            end setup = (x = randn(); dx = randn())
-            SUITE["value_and_pushforward!"][(1, m)][string(backend)] = @benchmarkable begin
-                value_and_pushforward!(dy, $backend, $scalar_to_vector, x, dx)
-            end setup = (x = randn(); dx = randn(); dy = zeros($m))
-        else
-            SUITE["value_and_pullback"][(1, m)][string(backend)] = @benchmarkable begin
-                value_and_pullback($backend, $scalar_to_vector, x, dy)
-            end setup = (x = randn(); dy = ones($m))
-            SUITE["value_and_pullback!"][(1, m)][string(backend)] = @benchmarkable begin
-                value_and_pullback!(dx, $backend, $scalar_to_vector, x, dy)
-            end setup = (x = randn(); dy = ones($m); dx = 0.0)
-        end
+    for backend in forward_fallback_backends
+        add_pushforward_benchmarks!(SUITE, backend, scalar_to_scalar, 1, 1)
     end
-end
-
-### Vector to scalar
-
-for n in [10]
-    vector_to_scalar = Layer(randn(n), randn(), tanh)
-
-    for backend in all_backends
-        handles_types(backend, Vector, Number) || continue
-        SUITE["value_and_gradient"][(n, 1)][string(backend)] = @benchmarkable begin
-            value_and_gradient($backend, $vector_to_scalar, x)
-        end setup = (x = randn($n))
-        SUITE["value_and_gradient!"][(n, 1)][string(backend)] = @benchmarkable begin
-            value_and_gradient!(grad, $backend, $vector_to_scalar, x)
-        end setup = (x = randn($n); grad = zeros($n))
+    for backend in reverse_fallback_backends
+        add_pullback_benchmarks!(SUITE, backend, scalar_to_scalar, 1, 1)
     end
 
-    for backend in all_fallback_backends
-        handles_types(backend, Vector, Number) || continue
-        if autodiff_mode(backend) == :forward
-            SUITE["value_and_pushforward"][(n, 1)][string(backend)] = @benchmarkable begin
-                value_and_pushforward($backend, $vector_to_scalar, x, dx)
-            end setup = (x = randn($n); dx = randn($n))
-            SUITE["value_and_pushforward!"][(n, 1)][string(backend)] = @benchmarkable begin
-                value_and_pushforward!(dy, $backend, $vector_to_scalar, x, dx)
-            end setup = (x = randn($n); dx = randn($n); dy = 0.0)
-        else
-            SUITE["value_and_pullback"][(n, 1)][string(backend)] = @benchmarkable begin
-                value_and_pullback($backend, $vector_to_scalar, x, dy)
-            end setup = (x = randn($n); dy = randn())
-            SUITE["value_and_pullback!"][(n, 1)][string(backend)] = @benchmarkable begin
-                value_and_pullback!(dx, $backend, $vector_to_scalar, x, dy)
-            end setup = (x = randn($n); dy = randn(); dx = zeros($n))
+    ### Scalar to vector
+    for m in [10]
+        scalar_to_vector = Layer(randn(m), randn(m), tanh)
+
+        for backend in all_backends
+            add_multiderivative_benchmarks!(SUITE, backend, scalar_to_vector, 1, m)
+        end
+        for backend in forward_fallback_backends
+            add_pushforward_benchmarks!(SUITE, backend, scalar_to_vector, 1, m)
+        end
+        for backend in reverse_fallback_backends
+            add_pullback_benchmarks!(SUITE, backend, scalar_to_vector, 1, m)
         end
     end
-end
-
-### Vector to vector
 
-for (n, m) in [(10, 10)]
-    vector_to_vector = Layer(randn(m, n), randn(m), tanh)
+    ### Vector to scalar
+    for n in [10]
+        vector_to_scalar = Layer(randn(n), randn(), tanh)
 
-    for backend in all_backends
-        handles_types(backend, Vector, Vector) || continue
-        SUITE["value_and_jacobian"][(n, m)][string(backend)] = @benchmarkable begin
-            value_and_jacobian($backend, $vector_to_vector, x)
-        end setup = (x = randn($n))
-        SUITE["value_and_jacobian!"][(n, m)][string(backend)] = @benchmarkable begin
-            value_and_jacobian!(jac, $backend, $vector_to_vector, x)
-        end setup = (x = randn($n); jac = zeros($m, $n))
+        for backend in all_backends
+            add_gradient_benchmarks!(SUITE, backend, vector_to_scalar, n, 1)
+        end
+        for backend in forward_fallback_backends
+            add_pushforward_benchmarks!(SUITE, backend, vector_to_scalar, n, 1)
+        end
+        for backend in reverse_fallback_backends
+            add_pullback_benchmarks!(SUITE, backend, vector_to_scalar, n, 1)
+        end
     end
 
-    for backend in all_fallback_backends
-        handles_types(backend, Vector, Vector) || continue
-        if autodiff_mode(backend) == :forward
-            SUITE["value_and_pushforward"][(n, m)][string(backend)] = @benchmarkable begin
-                value_and_pushforward($backend, $vector_to_vector, x, dx)
-            end setup = (x = randn($n); dx = randn($n))
-            SUITE["value_and_pushforward!"][(n, m)][string(backend)] = @benchmarkable begin
-                value_and_pushforward!(dy, $backend, $vector_to_vector, x, dx)
-            end setup = (x = randn($n); dx = randn($n); dy = zeros($m))
-        else
-            SUITE["value_and_pullback"][(n, m)][string(backend)] = @benchmarkable begin
-                value_and_pullback($backend, $vector_to_vector, x, dy)
-            end setup = (x = randn($n); dy = randn($m))
-            SUITE["value_and_pullback!"][(n, m)][string(backend)] = @benchmarkable begin
-                value_and_pullback!(dx, $backend, $vector_to_vector, x, dy)
-            end setup = (x = randn($n); dy = randn($m); dx = zeros($n))
+    ### Vector to vector
+    for (n, m) in [(10, 10)]
+        vector_to_vector = Layer(randn(m, n), randn(m), tanh)
+
+        for backend in all_backends
+            add_jacobian_benchmarks!(SUITE, backend, vector_to_vector, n, m)
+        end
+        for backend in forward_fallback_backends
+            add_pushforward_benchmarks!(SUITE, backend, vector_to_vector, n, m)
+        end
+        for backend in reverse_fallback_backends
+            add_pullback_benchmarks!(SUITE, backend, vector_to_vector, n, m)
         end
     end
+
+    return SUITE
 end
 
+include("utils.jl")
+
+SUITE = make_suite()
+
 # Run benchmarks locally
 # results = BenchmarkTools.run(SUITE; verbose=true)
 
 # Compare commits locally
 # using BenchmarkCI; BenchmarkCI.judge(baseline="origin/main"); BenchmarkCI.displayjudgement()
+
+# Parse into dataframe
+# include("dataframe.jl")
+# data = parse_benchmark_results(results; path=joinpath(@__DIR__, "results.csv"))

benchmark/dataframe.jl

@@ -0,0 +1,36 @@
+using BenchmarkTools
+using CSV
+using DataFrames
+using Statistics
+
+function parse_benchmark_results_aux(result::BenchmarkTools.Trial, level=nothing)
+    data = DataFrame(
+        :samples => [length(result.times)],
+        :time_median => [median(result.times)],
+        :memory_median => [median(result.memory)],
+        :allocs_median => [median(result.allocs)],
+    )
+    return data
+end
+
+function parse_benchmark_results_aux(results::BenchmarkGroup, level=1)
+    data = DataFrame()
+    level_symbol = Symbol(string("level_$level"))
+    for (key, val) in pairs(results)
+        subdata = parse_benchmark_results_aux(val, level + 1)
+        subdata[!, level_symbol] = fill(key, size(subdata, 1))
+        append!(data, subdata)
+    end
+    select!(data, level_symbol, Not(level_symbol))
+    return data
+end
+
+function parse_benchmark_results(results::BenchmarkGroup; path=nothing)
+    data = parse_benchmark_results_aux(results)
+    if !isnothing(path)
+        open(path, "w") do file
+            CSV.write(file, data)
+        end
+    end
+    return data
+end