Skip to content

Commit 0539ba4

Browse files
fverdugofverdugo
authored
Merge pull request #338 from gridap/autodiff_a_la_forwarddiff
Autodiff a la forwarddiff
2 parents 423a802 + 2952569 commit 0539ba4

15 files changed

+685
-0
lines changed

‎bench/ArraysBenchs/AutodiffBenchs.jl

+106
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
module AutodiffBenchs
2+
3+
using Gridap.Arrays
4+
import Gridap.Arrays: kernel_cache
5+
import Gridap.Arrays: apply_kernel!
6+
7+
function user_cell_energy(cell_u)
8+
function f(u)
9+
e = zero(eltype(u))
10+
for ui in u
11+
e += ui^2
12+
end
13+
e
14+
end
15+
apply(f,cell_u)
16+
end
17+
18+
function user_cell_residual(cell_u)
19+
apply(R(),cell_u)
20+
end
21+
22+
struct R <: Kernel end
23+
24+
function kernel_cache(k::R,u)
25+
r = copy(u)
26+
r .= zero(eltype(u))
27+
r
28+
end
29+
30+
@inline function apply_kernel!(r,k::R,u)
31+
for (i,ui) in enumerate(u)
32+
r[i] = 2*ui
33+
end
34+
r
35+
end
36+
37+
function user_cell_jacobian(cell_u)
38+
apply(J(),cell_u)
39+
end
40+
41+
struct J <: Kernel end
42+
43+
function kernel_cache(k::J,u)
44+
n = length(u)
45+
j = zeros(n,n)
46+
j
47+
end
48+
49+
@inline function apply_kernel!(j,k::J,u)
50+
n = length(u)
51+
for i in 1:n
52+
j[i,i] = 2
53+
end
54+
j
55+
end
56+
57+
L = 10^6
58+
l = 14
59+
cell_u = [ rand(l) for i in 1:L ]
60+
61+
cell_e = user_cell_energy(cell_u)
62+
cell_r = user_cell_residual(cell_u)
63+
cell_j = user_cell_jacobian(cell_u)
64+
cell_h = cell_j
65+
66+
cell_r_auto = autodiff_array_gradient(user_cell_energy,cell_u)
67+
cell_j_auto = autodiff_array_jacobian(user_cell_residual,cell_u)
68+
cell_h_auto = autodiff_array_hessian(user_cell_energy,cell_u)
69+
70+
function bench_loop(a)
71+
cache = array_cache(a)
72+
@time loop_cached(a,cache)
73+
end
74+
75+
function loop_cached(a,cache)
76+
for i in eachindex(a)
77+
ai = getindex!(cache,a,i)
78+
end
79+
end
80+
81+
println("Warm up ...")
82+
#bench_loop(cell_e)
83+
bench_loop(cell_r)
84+
bench_loop(cell_r_auto)
85+
bench_loop(cell_j)
86+
bench_loop(cell_j_auto)
87+
bench_loop(cell_h)
88+
bench_loop(cell_h_auto)
89+
90+
91+
println("Run ...")
92+
#bench_loop(cell_e)
93+
bench_loop(cell_r)
94+
bench_loop(cell_r_auto)
95+
bench_loop(cell_j)
96+
bench_loop(cell_j_auto)
97+
bench_loop(cell_h)
98+
bench_loop(cell_h_auto)
99+
100+
#using BenchmarkTools
101+
#
102+
#cache = array_cache(cell_h_auto)
103+
#i = 3
104+
#@btime a = getindex!($cache,$cell_h_auto,$i)
105+
106+
end # module

‎bench/ArraysBenchs/runbenchs.jl

+2
Original file line numberDiff line numberDiff line change
@@ -4,4 +4,6 @@ include("KernelsBenchs.jl")
44

55
include("ApplyBenchs.jl")
66

7+
include("AutodiffBenchs.jl")
8+
79
end # module

‎src/Arrays/Arrays.jl

+7
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ using Test
1818
using FillArrays
1919
using Base: @propagate_inbounds
2020
using LinearAlgebra
21+
using ForwardDiff
2122

2223
export array_cache
2324
export getindex!
@@ -89,6 +90,10 @@ export lazy_append
8990
export lazy_split
9091
export AppendedArray
9192

93+
export autodiff_array_gradient
94+
export autodiff_array_jacobian
95+
export autodiff_array_hessian
96+
9297
import Base: size
9398
import Base: getindex, setindex!
9499
import Base: similar
@@ -123,4 +128,6 @@ include("ArrayPairs.jl")
123128

124129
include("AppendedArrays.jl")
125130

131+
include("Autodiff.jl")
132+
126133
end # module

‎src/Arrays/Autodiff.jl

+114
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
2+
"""
3+
"""
4+
function autodiff_array_gradient(a,i_to_x,j_to_i=IdentityVector(length(i_to_x)))
5+
6+
i_to_xdual = apply(i_to_x) do x
7+
cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{length(x)}())
8+
xdual = cfg.duals
9+
xdual
10+
end
11+
12+
j_to_f = to_array_of_functions(a,i_to_xdual,j_to_i)
13+
j_to_x = reindex(i_to_x,j_to_i)
14+
15+
k = ForwardDiffGradientKernel()
16+
apply(k,j_to_f,j_to_x)
17+
18+
end
19+
20+
struct ForwardDiffGradientKernel <: Kernel end
21+
22+
function kernel_cache(k::ForwardDiffGradientKernel,f,x)
23+
cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{length(x)}())
24+
r = copy(x)
25+
(r, cfg)
26+
end
27+
28+
@inline function apply_kernel!(cache,k::ForwardDiffGradientKernel,f,x)
29+
r, cfg = cache
30+
@notimplementedif length(r) != length(x)
31+
ForwardDiff.gradient!(r,f,x,cfg)
32+
r
33+
end
34+
35+
"""
36+
"""
37+
function autodiff_array_jacobian(a,i_to_x,j_to_i=IdentityVector(length(i_to_x)))
38+
39+
i_to_xdual = apply(i_to_x) do x
40+
cfg = ForwardDiff.JacobianConfig(nothing, x, ForwardDiff.Chunk{length(x)}())
41+
xdual = cfg.duals
42+
xdual
43+
end
44+
45+
j_to_f = to_array_of_functions(a,i_to_xdual,j_to_i)
46+
j_to_x = reindex(i_to_x,j_to_i)
47+
48+
k = ForwardDiffJacobianKernel()
49+
apply(k,j_to_f,j_to_x)
50+
51+
end
52+
53+
struct ForwardDiffJacobianKernel <: Kernel end
54+
55+
function kernel_cache(k::ForwardDiffJacobianKernel,f,x)
56+
cfg = ForwardDiff.JacobianConfig(nothing, x, ForwardDiff.Chunk{length(x)}())
57+
n = length(x)
58+
j = zeros(eltype(x),n,n)
59+
(j, cfg)
60+
end
61+
62+
@inline function apply_kernel!(cache,k::ForwardDiffJacobianKernel,f,x)
63+
j, cfg = cache
64+
@notimplementedif size(j,1) != length(x)
65+
@notimplementedif size(j,2) != length(x)
66+
ForwardDiff.jacobian!(j,f,x,cfg)
67+
j
68+
end
69+
70+
"""
71+
"""
72+
function autodiff_array_hessian(a,i_to_x,j_to_i=IdentityVector(length(i_to_x)))
73+
agrad = i_to_y -> autodiff_array_gradient(a,i_to_y,j_to_i)
74+
autodiff_array_jacobian(agrad,i_to_x,j_to_i)
75+
end
76+
77+
function to_array_of_functions(a,x,ids=IdentityVector(length(x)))
78+
k = ArrayOfFunctionsKernel(a,x)
79+
j = IdentityVector(length(ids))
80+
apply(k,j)
81+
end
82+
83+
struct ArrayOfFunctionsKernel{A,X} <: Kernel
84+
a::A
85+
x::X
86+
end
87+
88+
function kernel_cache(k::ArrayOfFunctionsKernel,j)
89+
xi = testitem(k.x)
90+
l = length(k.x)
91+
x = MutableFill(xi,l)
92+
ax = k.a(x)
93+
axc = array_cache(ax)
94+
(ax, x, axc)
95+
end
96+
97+
@inline function apply_kernel!(cache,k::ArrayOfFunctionsKernel,j)
98+
ax, x, axc = cache
99+
@inline function f(xj)
100+
x.value = xj
101+
axj = getindex!(axc,ax,j)
102+
end
103+
f
104+
end
105+
106+
mutable struct MutableFill{T} <: AbstractVector{T}
107+
value::T
108+
length::Int
109+
end
110+
111+
Base.size(a::MutableFill) = (a.length,)
112+
113+
@inline Base.getindex(a::MutableFill,i::Integer) = a.value
114+

‎src/Exports.jl

+1
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ using Gridap.TensorValues: ⊗; export ⊗
118118
@publish FESpaces TrialFESpace
119119
@publish FESpaces TestFESpace
120120
@publish FESpaces FETerm
121+
@publish FESpaces FEEnergy
121122
@publish FESpaces AffineFETerm
122123
@publish FESpaces LinearFETerm
123124
@publish FESpaces FESource

‎src/FESpaces/FEAutodiff.jl

+52
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
2+
"""
3+
"""
4+
function autodiff_cell_residual_from_energy(uh_to_cell_energy::Function,uh,cell_ids=_default_cell_ids(uh))
5+
@assert is_a_fe_function(uh)
6+
U = get_fe_space(uh)
7+
cell_u_to_cell_energy = _change_argument_to_cell_u(uh_to_cell_energy,U)
8+
cell_u = get_cell_values(uh)
9+
cell_r = autodiff_array_gradient(cell_u_to_cell_energy,cell_u,cell_ids)
10+
cell_r
11+
end
12+
13+
"""
14+
"""
15+
function autodiff_cell_jacobian_from_energy(uh_to_cell_energy::Function,uh,cell_ids=_default_cell_ids(uh))
16+
@assert is_a_fe_function(uh)
17+
U = get_fe_space(uh)
18+
cell_u_to_cell_energy = _change_argument_to_cell_u(uh_to_cell_energy,U)
19+
cell_u = get_cell_values(uh)
20+
cell_j = autodiff_array_hessian(cell_u_to_cell_energy,cell_u,cell_ids)
21+
cell_j
22+
end
23+
24+
"""
25+
"""
26+
function autodiff_cell_jacobian_from_residual(uh_to_cell_residual::Function,uh,cell_ids=_default_cell_ids(uh))
27+
@assert is_a_fe_function(uh)
28+
U = get_fe_space(uh)
29+
cell_u_to_cell_residual = _change_argument_to_cell_u(uh_to_cell_residual,U)
30+
cell_u = get_cell_values(uh)
31+
cell_j = autodiff_array_jacobian(cell_u_to_cell_residual,cell_u,cell_ids)
32+
cell_j
33+
end
34+
35+
function _default_cell_ids(uh)
36+
cell_u = get_cell_values(uh)
37+
cell_ids = IdentityVector(length(cell_u))
38+
end
39+
40+
function _change_argument_to_cell_u(uh_to_cell_energy,U)
41+
function f(cell_u)
42+
cell_shapefuns = get_cell_shapefuns(U)
43+
cell_map = get_cell_map(U)
44+
ref_style = RefStyle(get_cell_dof_basis(U))
45+
array = lincomb(cell_shapefuns,cell_u)
46+
uh = GenericCellField(array,cell_map,ref_style)
47+
cell_e = uh_to_cell_energy(uh)
48+
cell_e
49+
end
50+
f
51+
end
52+

‎src/FESpaces/FESpaces.jl

+11
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ export interpolate
162162
export interpolate_everywhere
163163
export interpolate_dirichlet
164164
export get_cell_dof_basis
165+
export get_cell_shapefuns
165166

166167
export SingleFieldFEFunction
167168

@@ -245,6 +246,12 @@ export apply_statelaw
245246
export CellField
246247
export update_state_variables!
247248

249+
export autodiff_cell_residual_from_energy
250+
export autodiff_cell_jacobian_from_energy
251+
export autodiff_cell_jacobian_from_residual
252+
253+
export FEEnergy
254+
248255
include("CellBases.jl")
249256

250257
include("CellDofBases.jl")
@@ -299,4 +306,8 @@ include("FESpaceFactories.jl")
299306

300307
include("StateLaws.jl")
301308

309+
include("FEAutodiff.jl")
310+
311+
include("FETermsWithAutodiff.jl")
312+
302313
end # module

0 commit comments

Comments
 (0)