Add utilities that don't return the value (#27)

* Add special cases

* More complex backends

* Fix benchmarks

* More benchmarks

* Reactivate JET test

* Remove weird build folder

* Don't import dangerous stuff before JET test

* Re-skip JET

* Wrong JET test

* Better docstrings, more exhaustive tests, benchmark with NN layer

* Add utilities that don't return the value

* Structured benchmarks with dataframe export

.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