Autodiff a la forwarddiff

bench/ArraysBenchs/AutodiffBenchs.jl

@@ -0,0 +1,106 @@
+module AutodiffBenchs
+
+using Gridap.Arrays
+import Gridap.Arrays: kernel_cache
+import Gridap.Arrays: apply_kernel!
+
+function user_cell_energy(cell_u)
+  function f(u)
+    e = zero(eltype(u))
+    for ui in u
+      e += ui^2
+    end
+    e
+  end
+  apply(f,cell_u)
+end
+
+function user_cell_residual(cell_u)
+  apply(R(),cell_u)
+end
+
+struct R <: Kernel end
+
+function kernel_cache(k::R,u)
+  r = copy(u)
+  r .= zero(eltype(u))
+  r
+end
+
+@inline function apply_kernel!(r,k::R,u)
+  for (i,ui) in enumerate(u)
+    r[i] = 2*ui
+  end
+  r
+end
+
+function user_cell_jacobian(cell_u)
+  apply(J(),cell_u)
+end
+
+struct J <: Kernel end
+
+function kernel_cache(k::J,u)
+  n = length(u)
+  j = zeros(n,n)
+  j
+end
+
+@inline function apply_kernel!(j,k::J,u)
+  n = length(u)
+  for i in 1:n
+    j[i,i] = 2
+  end
+  j
+end
+
+L = 10^6
+l = 14
+cell_u = [ rand(l) for i in 1:L ]
+
+cell_e = user_cell_energy(cell_u)
+cell_r = user_cell_residual(cell_u)
+cell_j = user_cell_jacobian(cell_u)
+cell_h = cell_j
+
+cell_r_auto = autodiff_array_gradient(user_cell_energy,cell_u)
+cell_j_auto = autodiff_array_jacobian(user_cell_residual,cell_u)
+cell_h_auto = autodiff_array_hessian(user_cell_energy,cell_u)
+
+function bench_loop(a)
+  cache = array_cache(a)
+  @time loop_cached(a,cache)
+end
+
+function loop_cached(a,cache)
+  for i in eachindex(a)
+    ai = getindex!(cache,a,i)
+  end
+end
+
+println("Warm up ...")
+#bench_loop(cell_e)
+bench_loop(cell_r)
+bench_loop(cell_r_auto)
+bench_loop(cell_j)
+bench_loop(cell_j_auto)
+bench_loop(cell_h)
+bench_loop(cell_h_auto)
+
+
+println("Run ...")
+#bench_loop(cell_e)
+bench_loop(cell_r)
+bench_loop(cell_r_auto)
+bench_loop(cell_j)
+bench_loop(cell_j_auto)
+bench_loop(cell_h)
+bench_loop(cell_h_auto)
+
+#using BenchmarkTools
+#
+#cache = array_cache(cell_h_auto)
+#i = 3
+#@btime a = getindex!($cache,$cell_h_auto,$i)
+
+end # module

bench/ArraysBenchs/runbenchs.jl

@@ -4,4 +4,6 @@ include("KernelsBenchs.jl")
 
 include("ApplyBenchs.jl")
 
+include("AutodiffBenchs.jl")
+
 end # module

src/Arrays/Arrays.jl

@@ -18,6 +18,7 @@ using Test
 using FillArrays
 using Base: @propagate_inbounds
 using LinearAlgebra
+using ForwardDiff
 
 export array_cache
 export getindex!
@@ -89,6 +90,10 @@ export lazy_append
 export lazy_split
 export AppendedArray
 
+export autodiff_array_gradient
+export autodiff_array_jacobian
+export autodiff_array_hessian
+
 import Base: size
 import Base: getindex, setindex!
 import Base: similar
@@ -123,4 +128,6 @@ include("ArrayPairs.jl")
 
 include("AppendedArrays.jl")
 
+include("Autodiff.jl")
+
 end # module

src/Arrays/Autodiff.jl

@@ -0,0 +1,114 @@
+
+"""
+"""
+function autodiff_array_gradient(a,i_to_x,j_to_i=IdentityVector(length(i_to_x)))
+
+  i_to_xdual = apply(i_to_x) do x
+    cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{length(x)}())
+    xdual = cfg.duals
+    xdual
+  end
+
+  j_to_f = to_array_of_functions(a,i_to_xdual,j_to_i)
+  j_to_x = reindex(i_to_x,j_to_i)
+
+  k = ForwardDiffGradientKernel()
+  apply(k,j_to_f,j_to_x)
+
+end
+
+struct ForwardDiffGradientKernel <: Kernel end
+
+function kernel_cache(k::ForwardDiffGradientKernel,f,x)
+  cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{length(x)}())
+  r = copy(x)
+  (r, cfg)
+end
+
+@inline function apply_kernel!(cache,k::ForwardDiffGradientKernel,f,x)
+  r, cfg = cache
+  @notimplementedif length(r) != length(x)
+  ForwardDiff.gradient!(r,f,x,cfg)
+  r
+end
+
+"""
+"""
+function autodiff_array_jacobian(a,i_to_x,j_to_i=IdentityVector(length(i_to_x)))
+
+  i_to_xdual = apply(i_to_x) do x
+    cfg = ForwardDiff.JacobianConfig(nothing, x, ForwardDiff.Chunk{length(x)}())
+    xdual = cfg.duals
+    xdual
+  end
+
+  j_to_f = to_array_of_functions(a,i_to_xdual,j_to_i)
+  j_to_x = reindex(i_to_x,j_to_i)
+
+  k = ForwardDiffJacobianKernel()
+  apply(k,j_to_f,j_to_x)
+
+end
+
+struct ForwardDiffJacobianKernel <: Kernel end
+
+function kernel_cache(k::ForwardDiffJacobianKernel,f,x)
+  cfg = ForwardDiff.JacobianConfig(nothing, x, ForwardDiff.Chunk{length(x)}())
+  n = length(x)
+  j = zeros(eltype(x),n,n)
+  (j, cfg)
+end
+
+@inline function apply_kernel!(cache,k::ForwardDiffJacobianKernel,f,x)
+  j, cfg = cache
+  @notimplementedif size(j,1) != length(x)
+  @notimplementedif size(j,2) != length(x)
+  ForwardDiff.jacobian!(j,f,x,cfg)
+  j
+end
+
+"""
+"""
+function autodiff_array_hessian(a,i_to_x,j_to_i=IdentityVector(length(i_to_x)))
+   agrad = i_to_y -> autodiff_array_gradient(a,i_to_y,j_to_i)
+   autodiff_array_jacobian(agrad,i_to_x,j_to_i)
+end
+
+function to_array_of_functions(a,x,ids=IdentityVector(length(x)))
+  k = ArrayOfFunctionsKernel(a,x)
+  j = IdentityVector(length(ids))
+  apply(k,j)
+end
+
+struct ArrayOfFunctionsKernel{A,X} <: Kernel
+  a::A
+  x::X
+end
+
+function kernel_cache(k::ArrayOfFunctionsKernel,j)
+  xi = testitem(k.x)
+  l = length(k.x)
+  x = MutableFill(xi,l)
+  ax = k.a(x)
+  axc = array_cache(ax)
+  (ax, x, axc)
+end
+
+@inline function apply_kernel!(cache,k::ArrayOfFunctionsKernel,j)
+  ax, x, axc = cache
+  @inline function f(xj)
+    x.value = xj
+    axj = getindex!(axc,ax,j)
+  end
+  f
+end
+
+mutable struct MutableFill{T} <: AbstractVector{T}
+  value::T
+  length::Int
+end
+
+Base.size(a::MutableFill) = (a.length,)
+
+@inline Base.getindex(a::MutableFill,i::Integer) = a.value
+

src/Exports.jl

@@ -118,6 +118,7 @@ using Gridap.TensorValues: ⊗; export ⊗
 @publish FESpaces TrialFESpace
 @publish FESpaces TestFESpace
 @publish FESpaces FETerm
+@publish FESpaces FEEnergy
 @publish FESpaces AffineFETerm
 @publish FESpaces LinearFETerm
 @publish FESpaces FESource

src/FESpaces/FEAutodiff.jl

@@ -0,0 +1,52 @@
+
+"""
+"""
+function autodiff_cell_residual_from_energy(uh_to_cell_energy::Function,uh,cell_ids=_default_cell_ids(uh))
+  @assert is_a_fe_function(uh)
+  U = get_fe_space(uh)
+  cell_u_to_cell_energy = _change_argument_to_cell_u(uh_to_cell_energy,U)
+  cell_u = get_cell_values(uh)
+  cell_r = autodiff_array_gradient(cell_u_to_cell_energy,cell_u,cell_ids)
+  cell_r
+end
+
+"""
+"""
+function autodiff_cell_jacobian_from_energy(uh_to_cell_energy::Function,uh,cell_ids=_default_cell_ids(uh))
+  @assert is_a_fe_function(uh)
+  U = get_fe_space(uh)
+  cell_u_to_cell_energy = _change_argument_to_cell_u(uh_to_cell_energy,U)
+  cell_u = get_cell_values(uh)
+  cell_j = autodiff_array_hessian(cell_u_to_cell_energy,cell_u,cell_ids)
+  cell_j
+end
+
+"""
+"""
+function autodiff_cell_jacobian_from_residual(uh_to_cell_residual::Function,uh,cell_ids=_default_cell_ids(uh))
+  @assert is_a_fe_function(uh)
+  U = get_fe_space(uh)
+  cell_u_to_cell_residual = _change_argument_to_cell_u(uh_to_cell_residual,U)
+  cell_u = get_cell_values(uh)
+  cell_j = autodiff_array_jacobian(cell_u_to_cell_residual,cell_u,cell_ids)
+  cell_j
+end
+
+function _default_cell_ids(uh)
+  cell_u = get_cell_values(uh)
+  cell_ids = IdentityVector(length(cell_u))
+end
+
+function _change_argument_to_cell_u(uh_to_cell_energy,U)
+  function f(cell_u)
+    cell_shapefuns = get_cell_shapefuns(U)
+    cell_map = get_cell_map(U)
+    ref_style = RefStyle(get_cell_dof_basis(U))
+    array = lincomb(cell_shapefuns,cell_u)
+    uh = GenericCellField(array,cell_map,ref_style)
+    cell_e = uh_to_cell_energy(uh)
+    cell_e
+  end
+  f
+end
+

src/FESpaces/FESpaces.jl

@@ -162,6 +162,7 @@ export interpolate
 export interpolate_everywhere
 export interpolate_dirichlet
 export get_cell_dof_basis
+export get_cell_shapefuns
 
 export SingleFieldFEFunction
 
@@ -245,6 +246,12 @@ export apply_statelaw
 export CellField
 export update_state_variables!
 
+export autodiff_cell_residual_from_energy
+export autodiff_cell_jacobian_from_energy
+export autodiff_cell_jacobian_from_residual
+
+export FEEnergy
+
 include("CellBases.jl")
 
 include("CellDofBases.jl")
@@ -299,4 +306,8 @@ include("FESpaceFactories.jl")
 
 include("StateLaws.jl")
 
+include("FEAutodiff.jl")
+
+include("FETermsWithAutodiff.jl")
+
 end # module