Add benchmarks (#58)

* Remove reference folder

* Add benchmarks

.github/workflows/Benchmark.yml

@@ -0,0 +1,30 @@
+name: Run benchmarks
+on:
+  pull_request:
+    types: [labeled, opened, synchronize, reopened]
+# Only trigger the benchmark job when you add `run benchmark` label to the PR
+jobs:
+  Benchmark:
+    runs-on: ubuntu-latest
+    if: contains(github.event.pull_request.labels.*.name, 'run benchmark')
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@latest
+      - name: Cache artifacts
+        uses: actions/cache@v3
+        env:
+          cache-name: cache-artifacts
+        with:
+          path: ~/.julia/artifacts 
+          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
+          restore-keys: |
+            ${{ runner.os }}-test-${{ env.cache-name }}-
+            ${{ runner.os }}-test-
+            ${{ runner.os }}-
+      - name: Install dependencies
+        run: julia -e 'using Pkg; pkg"add JSON PkgBenchmark [email protected]"'
+      - name: Run benchmarks
+        run: julia benchmark/run_benchmarks.jl
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+

.gitignore

@@ -6,4 +6,5 @@
 /docs/Manifest.toml
 /docs/build/
 /docs/src/index.md
-/benchmark/Manifest.toml
+/benchmark/Manifest.toml
+/references/*

benchmark/Project.toml

@@ -1,5 +1,7 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+SimpleDiffEq = "05bca326-078c-5bf0-a5bf-ce7c7982d7fd"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"

benchmark/benchmark.jl

@@ -1,49 +1,74 @@
-using ADTypes: AbstractSparsityDetector, jacobian_sparsity
 using BenchmarkTools
+
+using ADTypes: AbstractSparsityDetector, jacobian_sparsity
 using SparseConnectivityTracer
 using SparseConnectivityTracer: DuplicateVector, RecursiveSet, SortedVector
-using NNlib: conv
 
-include("../test/brusselator_definition.jl")
+using SparseArrays: sprand
+using SimpleDiffEq: ODEProblem, solve, SimpleEuler
+using Flux: Conv
 
-const METHODS = (
-    TracerSparsityDetector(BitSet),
-    TracerSparsityDetector(Set{UInt64}),
-    TracerSparsityDetector(DuplicateVector{UInt64}),
-    TracerSparsityDetector(RecursiveSet{UInt64}),
-    TracerSparsityDetector(SortedVector{UInt64}),
-)
+const SET_TYPES = (BitSet, Set{UInt64}, DuplicateVector{UInt64}, SortedVector{UInt64})
 
-function benchmark_brusselator(N::Integer, method::AbstractSparsityDetector)
-    f! = Brusselator!(N)
-    x = rand(N, N, 2)
-    y = similar(x)
+## ODEs
+include("../test/brusselator_definition.jl")
 
-    return @benchmark jacobian_sparsity($f!, $y, $x, $method)
+## Iterated multiplication by random sparse matrices
+function iterated_sparse_matmul(x; sparsity=0.1, iterations=5)
+    n = length(x)
+    y = copy(x)
+    for _ in 1:iterations
+        A = sprand(n, n, sparsity)
+        y = A * y
+    end
+    return y
 end
 
-function benchmark_conv(N, method::AbstractSparsityDetector)
-    x = rand(N, N, 3, 1) # WHCN image 
-    w = rand(5, 5, 3, 2)  # corresponds to Conv((5, 5), 3 => 2)
-    f(x) = conv(x, w)
+## Deep Learning
+const LAYER = Conv((5, 5), 3 => 2)
 
-    return @benchmark jacobian_sparsity($f, $x, $method)
-end
+## Define Benchmark suite
+suite = BenchmarkGroup()
 
-## Run Brusselator benchmarks
-for N in (6, 24, 100)
-    for method in METHODS
-        @info "Benchmarking Brusselator of size $N with $method..."
-        b = benchmark_brusselator(N, method)
-        display(b)
+for S in SET_TYPES
+    setname = string(S)
+    method = TracerSparsityDetector(S)
+
+    J_global = suite["Jacobian"]["Global"]
+
+    ## Sparse matrix multiplication
+    for sparsity in (0.01, 0.05, 0.1, 0.25)
+        x = rand(50)
+        f(x) = iterated_sparse_matmul(x; sparsity=sparsity, iterations=5)
+        J_global["sparse_matmul"]["sparsity=$sparsity"][setname] = @benchmarkable jacobian_sparsity(
+            $f, $x, $method
+        )
+    end
+
+    ## ODEs
+    for N in (6, 24, 100)
+        f! = Brusselator!(N)
+        x = rand(N, N, 2)
+        y = similar(x)
+        J_global["brusselator"]["N=$N"][setname] = @benchmarkable jacobian_sparsity(
+            $f!, $y, $x, $method
+        )
+
+        solver = SimpleEuler()
+        prob = ODEProblem(brusselator_2d_loop!, x, (0.0, 1.0), f!.params)
+        function brusselator_ode_solve(x)
+            return solve(ODEProblem(brusselator_2d_loop!, x, (0.0, 1.0), f!.params), solver; dt=0.5).u[end]
+        end
+        J_global["brusselator_ode_solve"]["N=$N"][setname] = @benchmarkable jacobian_sparsity(
+            $brusselator_ode_solve, $x, $method
+        )
     end
-end
 
-## Run conv benchmarks
-for N in (28, 128)
-    for method in METHODS # Symbolics fails on this example
-        @info "Benchmarking NNlib.conv on image of size ($N, $N, 3) with with $method..."
-        b = benchmark_conv(N, method)
-        display(b)
+    ## Deep Learning
+    for N in (28, 128)
+        J_global["conv"]["size=$(N)x$(N)x3"][setname] = @benchmarkable jacobian_sparsity(
+            $LAYER, $(rand(N, N, 3, 1)), $method
+        )
     end
+    # TODO: benchmark local sparsity tracers on LeNet-5 CNN
 end

benchmark/run_benchmarks.jl

@@ -0,0 +1,10 @@
+# To run benchmarks locally, BenchmarkCI should be added to root project. Then call:
+# ```bash
+# julia benchmark/run_benchmarks.jl
+# ```
+
+using BenchmarkCI
+on_CI = haskey(ENV, "GITHUB_ACTIONS")
+
+BenchmarkCI.judge()
+on_CI ? BenchmarkCI.postjudge() : BenchmarkCI.displayjudgement()

test/brusselator_definition.jl

@@ -5,7 +5,7 @@
 #! format: off
 brusselator_f(x, y, t) = ifelse((((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) && (t >= 1.1), 5., 0.)
 limit(a, N) = a == N+1 ? 1 : a == 0 ? N : a
-function brusselator_2d_loop(du, u, p, t)
+function brusselator_2d_loop!(du, u, p, t)
     A, B, alpha, xyd, dx, N = p; alpha = alpha/dx^2
     @inbounds for I in CartesianIndices((N, N))
         i, j = Tuple(I)
@@ -37,4 +37,4 @@ function Brusselator!(N::Integer)
     return Brusselator!(N, params)
 end
 
-(b!::Brusselator!)(y, x) = brusselator_2d_loop(y, x, b!.params, nothing)
+(b!::Brusselator!)(y, x) = brusselator_2d_loop!(y, x, b!.params, nothing)