Skip to content

Commit

Permalink
More benchmarks
Browse files Browse the repository at this point in the history
  • Loading branch information
@gdalle
gdalle committed Mar 7, 2024
1 parent 3deaef1 commit c5e2fa3
Show file tree
Hide file tree
Showing 12 changed files with 284 additions and 115 deletions.
180 changes: 140 additions & 40 deletions benchmark/benchmarks.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,78 +12,178 @@ using PolyesterForwardDiff: PolyesterForwardDiff
using ReverseDiff: ReverseDiff
using Zygote: Zygote

## Settings

BenchmarkTools.DEFAULT_PARAMETERS.evals = 1
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 1

## Functions

struct ScalarToVector
n::Int
end

(f::ScalarToVector)(x::Real) = collect((1:(f.n)) .* x)
scalar_to_scalar(x::Real)::Real = 2x
(f::ScalarToVector)(x::Real)::AbstractVector = fill(2x, f.n)
vector_to_scalar(x::AbstractVector{<:Real})::Real = dot(eachindex(x), x)
vector_to_vector(x::AbstractVector{<:Real})::AbstractVector = eachindex(x) .* x

scalar_to_scalar(x::Real) = x
vector_to_scalar(x::AbstractVector{<:Real}) = dot(eachindex(x), x)
vector_to_vector(x::AbstractVector{<:Real}) = eachindex(x) .* x
## Backends

forward_backends = [
forward_custom_backends = [
EnzymeForwardBackend(; custom=true),
EnzymeForwardBackend(; custom=false),
FiniteDiffBackend(; custom=true),
FiniteDiffBackend(; custom=false),
ForwardDiffBackend(; custom=true),
PolyesterForwardDiffBackend(4; custom=true),
]

forward_fallback_backends = [
EnzymeForwardBackend(; custom=false),
FiniteDiffBackend(; custom=false),
ForwardDiffBackend(; custom=false),
]

reverse_backends = [
reverse_custom_backends = [
ZygoteBackend(; custom=true),
ZygoteBackend(; custom=false),
EnzymeReverseBackend(; custom=true),
EnzymeReverseBackend(; custom=false),
ReverseDiffBackend(; custom=true),
]

reverse_fallback_backends = [
ZygoteBackend(; custom=false),
EnzymeReverseBackend(; custom=false),
ReverseDiffBackend(; custom=false),
]

backends = vcat(forward_backends, reverse_backends)
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)

input_dims = [10]
output_dims = [10]
## Suite

SUITE = BenchmarkGroup()

for backend in backends
### Scalar to scalar

## Scalar to scalar
if !(backend isa ReverseDiffBackend)
SUITE["derivative"][(1, 1)][string(backend)] = @benchmarkable begin
value_and_derivative($backend, $scalar_to_scalar, x)
end setup = (x = 1.0) evals = 1 seconds = 1
for backend in all_backends
handles_types(backend, Number, Number) || continue
SUITE["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 ad_mode(backend) == :forward
SUITE["pushforward"][(1, 1)][string(backend)] = @benchmarkable begin
value_and_pushforward($backend, $scalar_to_scalar, x, dx)
end setup = (x = randn(); dx = randn())
else
SUITE["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 = ScalarToVector(m)

for backend in all_backends
handles_types(backend, Number, Vector) || continue
SUITE["multiderivative"][(1, m)][string(backend)] = @benchmarkable begin
value_and_multiderivative($backend, $scalar_to_vector, x)
end setup = (x = randn())
SUITE["multiderivative!"][(1, m)][string(backend)] = @benchmarkable begin
value_and_multiderivative!(multider, $backend, $scalar_to_vector, x)
end setup = (x = randn(); multider = zeros($m))
end

## Scalar to vector
if !(
backend isa ReverseDiffBackend ||
backend isa EnzymeReverseBackend ||
backend isa EnzymeForwardBackend
)
for m in output_dims
scalar_to_vector = ScalarToVector(m)
SUITE["multiderivative"][(1, m)][string(backend)] = @benchmarkable begin
value_and_multiderivative($backend, $scalar_to_vector, x)
end setup = (x = 1.0) evals = 1 seconds = 1
for backend in all_fallback_backends
handles_types(backend, Number, Vector) || continue
if ad_mode(backend) == :forward
SUITE["pushforward"][(1, m)][string(backend)] = @benchmarkable begin
value_and_pushforward($backend, $scalar_to_vector, x, dx)
end setup = (x = randn(); dx = randn())
SUITE["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["pullback"][(1, m)][string(backend)] = @benchmarkable begin
value_and_pullback($backend, $scalar_to_vector, x, dy)
end setup = (x = randn(); dy = ones($m))
SUITE["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
end
end

## Vector to scalar
for n in input_dims
### Vector to scalar

for n in [10]
for backend in all_backends
handles_types(backend, Vector, Number) || continue
SUITE["gradient"][(n, 1)][string(backend)] = @benchmarkable begin
value_and_gradient($backend, $vector_to_scalar, x)
end setup = (x = randn($n)) evals = 1 seconds = 1
end setup = (x = randn($n))
SUITE["gradient!"][(n, 1)][string(backend)] = @benchmarkable begin
value_and_gradient!(grad, $backend, $vector_to_scalar, x)
end setup = (x = randn($n); grad = zeros($n))
end

for backend in all_fallback_backends
handles_types(backend, Vector, Number) || continue
if ad_mode(backend) == :forward
SUITE["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["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["pullback"][(n, 1)][string(backend)] = @benchmarkable begin
value_and_pullback($backend, $vector_to_scalar, x, dy)
end setup = (x = randn($n); dy = randn())
SUITE["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))
end
end
end

### Vector to vector

for n in [10]
m = n
for backend in all_backends
handles_types(backend, Vector, Vector) || continue
SUITE["jacobian"][(n, m)][string(backend)] = @benchmarkable begin
value_and_jacobian($backend, $vector_to_vector, x)
end setup = (x = randn($n))
SUITE["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))
end

## Vector to vector
if !(backend isa EnzymeReverseBackend)
for n in input_dims, m in output_dims
backend isa EnzymeReverseBackend && continue
SUITE["jacobian"][(n, m)][string(backend)] = @benchmarkable begin
value_and_jacobian($backend, $vector_to_vector, x)
end setup = (x = randn($n)) evals = 1 seconds = 1
for backend in all_fallback_backends
handles_types(backend, Vector, Vector) || continue
if ad_mode(backend) == :forward
SUITE["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["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["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["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))
end
end
end
Expand Down
58 changes: 17 additions & 41 deletions ext/DifferentiationInterfaceFiniteDiffExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,17 +10,9 @@ using FiniteDiff:
finite_difference_jacobian
using LinearAlgebra: dot, mul!

const DEFAULT_INPLACE = Val{false} # see https://docs.sciml.ai/FiniteDiff/stable/#f-Definitions

## Stepping

struct Step{F,X,D}
f::F
x::X
d::D
end

(step::Step)(t::Number) = step.f(step.x + t * step.d)
# see https://docs.sciml.ai/FiniteDiff/stable/#f-Definitions
const FUNCTION_INPLACE = Val{true}
const FUNCTION_NOT_INPLACE = Val{false}

## Backend construction

Expand All @@ -36,38 +28,22 @@ end
## Primitives

function DI.value_and_pushforward!(
dy::Y, ::FiniteDiffBackend{custom,fdtype}, f, x::X, dx
) where {X<:Number,Y<:Number,custom,fdtype}
y = f(x)
der = finite_difference_derivative(f, x, fdtype, eltype(dy), y)
new_dy = der * dx
return y, new_dy
end

function DI.value_and_pushforward!(
dy::Y, ::FiniteDiffBackend{custom,fdtype}, f, x::X, dx
) where {X<:Number,Y<:AbstractArray,custom,fdtype}
y = f(x)
finite_difference_gradient!(dy, f, x, fdtype, eltype(dy), DEFAULT_INPLACE, y)
dy .*= dx
return y, dy
end

function DI.value_and_pushforward!(
dy::Y, ::FiniteDiffBackend{custom,fdtype}, f, x::X, dx
) where {X<:AbstractArray,Y<:Number,custom,fdtype}
dy::Y, ::FiniteDiffBackend{custom,fdtype}, f, x, dx
) where {Y<:Number,custom,fdtype}
y = f(x)
g = finite_difference_gradient(f, x, fdtype, eltype(dy), DEFAULT_INPLACE, y)
new_dy = dot(g, dx)
step(t::Number)::Number = f(x .+ t .* dx)
new_dy = finite_difference_derivative(step, zero(eltype(dx)), fdtype, eltype(y), y)
return y, new_dy
end

function DI.value_and_pushforward!(
dy::Y, ::FiniteDiffBackend{custom,fdtype}, f, x::X, dx
) where {X<:AbstractArray,Y<:AbstractArray,custom,fdtype}
dy::Y, ::FiniteDiffBackend{custom,fdtype}, f, x, dx
) where {Y<:AbstractArray,custom,fdtype}
y = f(x)
J = finite_difference_jacobian(f, x, fdtype, eltype(dy))
mul!(vec(dy), J, vec(dx))
step(t::Number)::AbstractArray = f(x .+ t .* dx)
finite_difference_gradient!(
dy, step, zero(eltype(dx)), fdtype, eltype(y), FUNCTION_NOT_INPLACE, y
)
return y, dy
end

Expand All @@ -85,31 +61,31 @@ function DI.value_and_multiderivative!(
multider::AbstractArray, ::FiniteDiffBackend{true,fdtype}, f, x::Number
) where {fdtype}
y = f(x)
finite_difference_gradient!(multider, f, x, fdtype, eltype(y), DEFAULT_INPLACE, y)
finite_difference_gradient!(multider, f, x, fdtype, eltype(y), FUNCTION_NOT_INPLACE, y)
return y, multider
end

function DI.value_and_multiderivative(
::FiniteDiffBackend{true,fdtype}, f, x::Number
) where {fdtype}
y = f(x)
multider = finite_difference_gradient(f, x, fdtype, eltype(y), DEFAULT_INPLACE, y)
multider = finite_difference_gradient(f, x, fdtype, eltype(y), FUNCTION_NOT_INPLACE, y)
return y, multider
end

function DI.value_and_gradient!(
grad::AbstractArray, ::FiniteDiffBackend{true,fdtype}, f, x::AbstractArray
) where {fdtype}
y = f(x)
finite_difference_gradient!(grad, f, x, fdtype, eltype(y), DEFAULT_INPLACE, y)
finite_difference_gradient!(grad, f, x, fdtype, eltype(y), FUNCTION_NOT_INPLACE, y)
return y, grad
end

function DI.value_and_gradient(
::FiniteDiffBackend{true,fdtype}, f, x::AbstractArray
) where {fdtype}
y = f(x)
grad = finite_difference_gradient(f, x, fdtype, eltype(y), DEFAULT_INPLACE, y)
grad = finite_difference_gradient(f, x, fdtype, eltype(y), FUNCTION_NOT_INPLACE, y)
return y, grad
end

Expand Down
2 changes: 1 addition & 1 deletion ext/DifferentiationInterfaceZygoteExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ DI.ZygoteBackend(; custom::Bool=true) = ChainRulesReverseBackend(ZygoteRuleConfi
const ZygoteBackendType{custom} = ChainRulesReverseBackend{custom,<:ZygoteRuleConfig}

function Base.show(io::IO, backend::ZygoteBackendType{custom}) where {custom}
return print(io, "ZygoteBackend{$custom}()")
return print(io, "ZygoteBackend{$(custom ? "custom" : "fallback")}()")
end

## Utilities
Expand Down
10 changes: 7 additions & 3 deletions src/DifferentiationInterface.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,14 @@ include("scalar_array.jl")
include("array_scalar.jl")
include("array_array.jl")

export ChainRulesReverseBackend,
ChainRulesForwardBackend,
EnzymeReverseBackend,
export AbstractBackend, AbstractForwardBackend, AbstractReverseBackend
export ad_mode
export handles_input_type, handles_output_type, handles_types

export ChainRulesForwardBackend,
ChainRulesReverseBackend,
EnzymeForwardBackend,
EnzymeReverseBackend,
FiniteDiffBackend,
ForwardDiffBackend,
PolyesterForwardDiffBackend,
Expand Down
Loading

0 comments on commit c5e2fa3

Please sign in to comment.