Skip to content

Add specialized StaticArrays-based methods #215

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Apr 27, 2017
Merged

Add specialized StaticArrays-based methods #215

merged 4 commits into from
Apr 27, 2017

Conversation

jrevels
Copy link
Member

@jrevels jrevels commented Apr 13, 2017

Adds StaticArrays as dependency, and uses it to support a low-dimensional, stack-allocated differentiation API.

julia> using BenchmarkTools, ForwardDiff, StaticArrays

julia> x = rand(3, 3, 3);

julia> sx = SArray{Tuple{3,3,3}}(x);

julia> out = similar(x);

julia> cfg = ForwardDiff.GradientConfig(sum, x);

# naive old method
julia> @btime ForwardDiff.gradient(sum, $x);
  6.585 μs (3 allocations: 3.23 KiB)

# non-allocating old method
julia> @btime ForwardDiff.gradient!($out, sum, $x, $cfg);
  666.856 ns (0 allocations: 0 bytes)

# stack-allocated new method
julia> @btime ForwardDiff.gradient(sum, $sx);
  9.394 ns (0 allocations: 0 bytes)

Currently, I've only added support for gradients and Jacobians. EDIT: Added special cases for Hessians as well, see below. I still need to add back in the tests I originally used in #213. Done.

@aaowens
Copy link
Contributor

aaowens commented Apr 17, 2017

Is this expected? The speedup for gradient! is minor.

julia> @btime ForwardDiff.gradient(sum, $sx);
  10.415 ns (0 allocations: 0 bytes)

julia> out2 = similar(sx);

julia> cfg2 = ForwardDiff.GradientConfig(sum, sx);

julia> @btime ForwardDiff.gradient!($out, sum, $x, $cfg);
  705.035 ns (0 allocations: 0 bytes)

julia> @btime ForwardDiff.gradient!($out2, sum, $sx, $cfg2);
  657.787 ns (0 allocations: 0 bytes)

This PR sounds really useful for tools like Optim on low-dimensional functions, both because of the potential differentiation speedup, and because returning the gradient as an SArray should enable the fast matrix solve path.

@jrevels
Copy link
Member Author

jrevels commented Apr 17, 2017

Is this expected? The speedup for gradient! is minor.

Yes, this is expected. The main benefit of gradient!(out, f, ::SArray) is that an SArray is passed to the target function, which might speed up evaluation quite a bit depending on the target function. However, for simple target functions like sum, gradient!(out, f, ::SArray) will spend the bulk of its time extracting derivatives to out.

Note this extraction process is not unrolled like it is for gradient(f, ::SArray). I'm not sure how much benefit we'd see from unrolling it (since out is heap-allocated anyway), but I'll play around with it.

@jrevels
Copy link
Member Author

jrevels commented Apr 18, 2017

Just added some special cases for Hessians as well:

julia> x = rand(3, 3);

julia> sx = SArray{Tuple{3, 3}}(x);

julia> out = similar(x, 9, 9);

julia> cfg = ForwardDiff.HessianConfig(prod, x);

# naive old method
julia> @btime ForwardDiff.hessian(prod, $x);
  37.159 μs (17 allocations: 17.11 KiB)

# almost non-allocating old method
julia> @btime ForwardDiff.hessian!($out, prod, $x, $cfg);
  2.115 μs (2 allocations: 832 bytes)

# stack-allocated new method
julia> @btime ForwardDiff.hessian(prod, $sx);
  1.220 μs (0 allocations: 0 bytes)

@jrevels jrevels changed the title [WIP] add specialized StaticArrays-based methods Add specialized StaticArrays-based methods Apr 18, 2017
@jrevels
Copy link
Member Author

jrevels commented Apr 18, 2017

Tests pass locally, and everything here is ready. However, I developed this against StaticArrays master, so this PR is now waiting on JuliaArrays/StaticArrays.jl#142.

@jrevels
Copy link
Member Author

jrevels commented Apr 26, 2017

Argh, 32 bit Windows...AppVeyor was passing before this PR, and it seems like this PR doesn't touch the code that would cause that test to fail, so I'm not sure what's going on here.

@@ -1,6 +1,5 @@
environment:
matrix:
- JULIAVERSION: "julianightlies/bin/winnt/x86/julia-latest-win32.exe"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cutting the Gordian knot.

@jrevels jrevels merged commit edad35c into master Apr 27, 2017
@jrevels jrevels deleted the jr/static branch May 19, 2017 12:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants