-
-
Notifications
You must be signed in to change notification settings - Fork 83
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #989 from SciML/gd/sparsead
Add sparse AD comparison
- Loading branch information
Showing
3 changed files
with
1,397 additions
and
0 deletions.
There are no files selected for viewing
169 changes: 169 additions & 0 deletions
169
benchmarks/AutomaticDifferentiationSparse/BrusselatorSparseAD.jmd
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,169 @@ | ||
--- | ||
title: Sparse AD benchmarks | ||
author: Guillaume Dalle | ||
--- | ||
|
||
```julia | ||
using ADTypes | ||
using LinearAlgebra, SparseArrays | ||
using BenchmarkTools, DataFrames | ||
import DifferentiationInterface as DI | ||
import SparseDiffTools as SDT | ||
using SparseConnectivityTracer: TracerSparsityDetector | ||
using SparseMatrixColorings: GreedyColoringAlgorithm | ||
using Symbolics: SymbolicsSparsityDetector | ||
using Test | ||
``` | ||
|
||
## Definitions | ||
|
||
```julia | ||
const N = 32 | ||
const xyd_brusselator = range(0; stop=1, length=N) | ||
const p = (3.4, 1.0, 10.0, step(xyd_brusselator)) | ||
|
||
brusselator_f(x, y, t) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.0 | ||
limit(a, N) = | ||
if a == N + 1 | ||
1 | ||
elseif a == 0 | ||
N | ||
else | ||
a | ||
end; | ||
``` | ||
|
||
```julia | ||
function brusselator_2d_loop(du, u, p, t) | ||
A, B, alpha, dx = p | ||
alpha = alpha / dx^2 | ||
@inbounds for I in CartesianIndices((N, N)) | ||
i, j = Tuple(I) | ||
x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]] | ||
ip1, im1, jp1, jm1 = limit(i + 1, N), | ||
limit(i - 1, N), limit(j + 1, N), | ||
limit(j - 1, N) | ||
du[i, j, 1] = | ||
alpha * | ||
(u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] - 4u[i, j, 1]) + | ||
B + | ||
u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] + brusselator_f(x, y, t) | ||
du[i, j, 2] = | ||
alpha * | ||
(u[im1, j, 2] + u[ip1, j, 2] + u[i, jp1, 2] + u[i, jm1, 2] - 4u[i, j, 2]) + | ||
A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2] | ||
end | ||
end; | ||
``` | ||
|
||
```julia | ||
function init_brusselator_2d(xyd) | ||
N = length(xyd) | ||
u = zeros(N, N, 2) | ||
for I in CartesianIndices((N, N)) | ||
x = xyd[I[1]] | ||
y = xyd[I[2]] | ||
u[I, 1] = 22 * (y * (1 - y))^(3 / 2) | ||
u[I, 2] = 27 * (x * (1 - x))^(3 / 2) | ||
end | ||
return u | ||
end; | ||
``` | ||
|
||
```julia | ||
x0 = init_brusselator_2d(xyd_brusselator); | ||
y0 = similar(x0); | ||
|
||
f!(y, x) = brusselator_2d_loop(y, x, p, 0.0); | ||
``` | ||
|
||
## Sparsity detection | ||
|
||
```julia | ||
S1 = ADTypes.jacobian_sparsity(f!, y0, x0, TracerSparsityDetector()) | ||
S2 = ADTypes.jacobian_sparsity(f!, y0, x0, SymbolicsSparsityDetector()) | ||
@test S1 == S2 | ||
``` | ||
|
||
```julia | ||
td1 = @belapsed ADTypes.jacobian_sparsity($f!, $y0, $x0, TracerSparsityDetector()) | ||
println("Sparsity detection with SparseConnectivityTracer: $td1 s") | ||
``` | ||
|
||
```julia | ||
td2 = @belapsed ADTypes.jacobian_sparsity($f!, $y0, $x0, SymbolicsSparsityDetector()) | ||
println("Sparsity detection with Symbolics: $td2 s") | ||
``` | ||
|
||
```julia | ||
println("Speedup from new sparsity detection method (>1 is better): $(td2 / td1)") | ||
``` | ||
|
||
## Coloring | ||
|
||
```julia | ||
S = S1 | ||
c1 = ADTypes.column_coloring(S, GreedyColoringAlgorithm()) | ||
c2 = SDT.matrix_colors(S) | ||
@test c1 == c2 | ||
``` | ||
|
||
```julia | ||
tc1 = @belapsed ADTypes.column_coloring($S, GreedyColoringAlgorithm()) | ||
println("Coloring with SparseMatrixColorings: $tc1 s") | ||
``` | ||
|
||
```julia | ||
tc2 = @belapsed SDT.matrix_colors($S) | ||
println("Coloring with SDT: $tc2 s") | ||
``` | ||
|
||
```julia | ||
println("Speedup from new coloring method (>1 is better): $(tc2 / tc1)") | ||
``` | ||
|
||
## Compressed differentiation | ||
|
||
```julia | ||
backend = AutoSparse( | ||
AutoForwardDiff(); | ||
sparsity_detector=TracerSparsityDetector(), | ||
coloring_algorithm=GreedyColoringAlgorithm(), | ||
); | ||
``` | ||
|
||
```julia | ||
extras = DI.prepare_jacobian(f!, similar(y0), backend, x0); | ||
J1 = DI.jacobian!(f!, similar(y0), similar(S, eltype(x0)), backend, x0, extras) | ||
|
||
cache = SDT.sparse_jacobian_cache( | ||
backend, SDT.JacPrototypeSparsityDetection(; jac_prototype=S), f!, similar(y0), x0 | ||
); | ||
J2 = SDT.sparse_jacobian!(similar(S, eltype(x0)), backend, cache, f!, similar(y0), x0) | ||
|
||
@test J1 == J2 | ||
``` | ||
|
||
```julia | ||
tj1 = @belapsed DI.jacobian!($f!, _y, _J, $backend, $x0, _extras) evals = 1 samples = 100 setup = ( | ||
_y = similar(y0); | ||
_J = similar(S, eltype(x0)); | ||
_extras = DI.prepare_jacobian(f!, similar(y0), backend, x0) | ||
) | ||
println("Jacobian with DifferentiationInterface: $tj1 s") | ||
``` | ||
|
||
```julia | ||
tj2 = @belapsed SDT.sparse_jacobian!(_J, $backend, _cache, $f!, _y, x0) evals = 1 samples = 100 setup = ( | ||
_y = similar(y0); | ||
_J = similar(S, eltype(x0)); | ||
_cache = SDT.sparse_jacobian_cache( | ||
backend, SDT.JacPrototypeSparsityDetection(; jac_prototype=S), f!, similar(y0), x0 | ||
) | ||
) | ||
println("Jacobian with SparseDiffTools: $tj2 s") | ||
``` | ||
|
||
```julia | ||
println("Speedup from new differentiation method (>1 is better): $(tj2 / tj1)") | ||
``` |
Oops, something went wrong.