Merge pull request #806 from kishore-nori/AD-for-SkeletonIntegration

Added a dummy tag for ForwardDiff configs being constructed in `Gridap` `Autodiff.jl` to fix issue #805

Author: amartinhuertas

NEWS.md

@@ -11,6 +11,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Extended the `MultiField` functional gradient (with respect to degrees-of-freedom of `MultiFieldFEFunction`) to functionals involving Skeleton integration. The interface remains the same `gradient(f,xh)`. Since PR [#799](https://github.com/gridap/Gridap.jl/pull/799)
 - Added functionality to take jacobian of functional involving integration (`DomainContribution`) over Skeleton faces (obtained from testing bilinear form with the whole set of test `fe_basis`), with respect to the degrees-of-freedom of `FEFunction`.  The interface remains the same - `jacobian(f,uh)`. Since PR [#803](https://github.com/gridap/Gridap.jl/pull/803)
 
+### Changed 
+- Added a dummy tag for ForwardDiff configs being constructed in Gridap at `src/Arrays/Autodiff.jl` to fix issue [#805](https://github.com/gridap/Gridap.jl/issues/805). Since PR [#806](https://github.com/gridap/Gridap.jl/pull/806)
+
 ### Fixed
 - Fixed the behavior of `gradient` for functionals involving operations of `CellFields` inside `mean` and `jump` terms of Skeleton Integration terms. Since PR [#800](https://github.com/gridap/Gridap.jl/pull/800)
 - Fixed the behavior `SkeletonCellFieldPair` at the Boundary integration terms. Since PR [#800](https://github.com/gridap/Gridap.jl/pull/800)

src/Arrays/Autodiff.jl

@@ -1,15 +1,17 @@
 
 function autodiff_array_gradient(a,i_to_x)
-  i_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient),i_to_x)
-  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient),i_to_x)
+  dummy_forwarddiff_tag = ()->()
+  i_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x)
+  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x)
   i_to_ydual = a(i_to_xdual)
   i_to_result = lazy_map(AutoDiffMap(ForwardDiff.gradient),i_to_ydual,i_to_x,i_to_cfg)
   i_to_result
 end
 
 function autodiff_array_jacobian(a,i_to_x)
-  i_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian),i_to_x)
-  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian),i_to_x)
+  dummy_forwarddiff_tag = ()->()
+  i_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x)
+  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x)
   i_to_ydual = a(i_to_xdual)
   i_to_result = lazy_map(AutoDiffMap(ForwardDiff.jacobian),i_to_ydual,i_to_x,i_to_cfg)
   i_to_result
@@ -21,19 +23,21 @@ function autodiff_array_hessian(a,i_to_x)
 end
 
 function autodiff_array_gradient(a,i_to_x,j_to_i)
-  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient),i_to_x)
+  dummy_forwarddiff_tag = ()->()
+  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x)
   j_to_ydual = a(i_to_xdual)
   j_to_x = lazy_map(Reindex(i_to_x),j_to_i)
-  j_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient),j_to_x)
+  j_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),j_to_x)
   j_to_result = lazy_map(AutoDiffMap(ForwardDiff.gradient),j_to_ydual,j_to_x,j_to_cfg)
   j_to_result
 end
 
 function autodiff_array_jacobian(a,i_to_x,j_to_i)
-  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian),i_to_x)
+  dummy_forwarddiff_tag = ()->()
+  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x)
   j_to_ydual = a(i_to_xdual)
   j_to_x = lazy_map(Reindex(i_to_x),j_to_i)
-  j_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian),j_to_x)
+  j_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),j_to_x)
   j_to_result = lazy_map(AutoDiffMap(ForwardDiff.jacobian),j_to_ydual,j_to_x,j_to_cfg)
   j_to_result
 end
@@ -43,31 +47,43 @@ function autodiff_array_hessian(a,i_to_x,i_to_j)
   autodiff_array_jacobian(agrad,i_to_x,i_to_j)
 end
 
-struct ConfigMap{F} <: Map
- f::F
+struct ConfigMap{
+  F <: Union{typeof(ForwardDiff.gradient),typeof(ForwardDiff.jacobian)},
+  T <: Union{<:Function,Nothing}} <: Map
+
+  f::F # ForwardDiff operation
+  tag::T # function for config tag name
 end
+# constructor with nothing as the tag, for backwards compatibility
+ConfigMap(f) = ConfigMap(f,nothing)
 
 # TODO Prescribing long chunk size can lead to slow compilation times!
 function return_cache(k::ConfigMap{typeof(ForwardDiff.gradient)},x)
-  cfg = ForwardDiff.GradientConfig(nothing,x,ForwardDiff.Chunk{length(x)}())
+  cfg = ForwardDiff.GradientConfig(k.tag,x,ForwardDiff.Chunk{length(x)}())
   cfg
 end
 
 function return_cache(k::ConfigMap{typeof(ForwardDiff.jacobian)},x)
-  cfg = ForwardDiff.JacobianConfig(nothing,x,ForwardDiff.Chunk{length(x)}())
+  cfg = ForwardDiff.JacobianConfig(k.tag,x,ForwardDiff.Chunk{length(x)}())
   cfg
 end
 
 function evaluate!(cfg,k::ConfigMap,x)
   cfg
 end
 
-struct DualizeMap{F} <: Map
- f::F
+struct DualizeMap{
+  F <: Union{typeof(ForwardDiff.gradient),typeof(ForwardDiff.jacobian)},
+  T <: Union{<:Function,Nothing}} <: Map
+
+  f::F # ForwardDiff operation
+  tag::T # function for config tag name
 end
+# constructor with nothing as the tag, for backwards compatibility
+DualizeMap(f) = DualizeMap(f,nothing)
 
 function return_cache(k::DualizeMap,x)
-  return_cache(ConfigMap(k.f),x)
+  return_cache(ConfigMap(k.f,k.tag),x)
 end
 
 function evaluate!(cfg,k::DualizeMap,x)

src/FESpaces/FEAutodiff.jl

@@ -166,20 +166,21 @@ end
   modules without any circular dependency
 =#
 function autodiff_array_gradient(a, i_to_x, j_to_i::SkeletonPair)
-  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient),i_to_x)
+  dummy_forwarddiff_tag = ()->()
+  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x)
 
   # dual output of both sides at once
   j_to_ydual_plus, j_to_ydual_minus = a(i_to_xdual)
 
   # Work for plus side
   j_to_x_plus = lazy_map(Reindex(i_to_x),j_to_i.plus)
-  j_to_cfg_plus = lazy_map(ConfigMap(ForwardDiff.gradient),j_to_x_plus)
+  j_to_cfg_plus = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),j_to_x_plus)
   j_to_result_plus = lazy_map(AutoDiffMap(ForwardDiff.gradient),
                               j_to_ydual_plus,j_to_x_plus,j_to_cfg_plus)
 
   # Work for minus side
   j_to_x_minus = lazy_map(Reindex(i_to_x),j_to_i.minus)
-  j_to_cfg_minus = lazy_map(ConfigMap(ForwardDiff.gradient),j_to_x_minus)
+  j_to_cfg_minus = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),j_to_x_minus)
   j_to_result_minus = lazy_map(AutoDiffMap(ForwardDiff.gradient),
                                j_to_ydual_minus,j_to_x_minus,j_to_cfg_minus)
 
@@ -190,7 +191,8 @@ function autodiff_array_gradient(a, i_to_x, j_to_i::SkeletonPair)
 end
 
 function autodiff_array_jacobian(a,i_to_x,j_to_i::SkeletonPair)
-  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian),i_to_x)
+  dummy_forwarddiff_tag = ()->()
+  i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x)
   j_to_ydual_plus, j_to_ydual_minus = a(i_to_xdual)
 
   # Bilinear form when tested with test basis functions dv results in
@@ -205,13 +207,13 @@ function autodiff_array_jacobian(a,i_to_x,j_to_i::SkeletonPair)
 
   # Work for plus side
   j_to_x_plus = lazy_map(Reindex(i_to_x),j_to_i.plus)
-  j_to_cfg_plus = lazy_map(ConfigMap(ForwardDiff.jacobian),j_to_x_plus)
+  j_to_cfg_plus = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),j_to_x_plus)
   j_to_result_plus_dense = lazy_map(AutoDiffMap(ForwardDiff.jacobian),
                                     j_to_ydual_plus_dense,j_to_x_plus,j_to_cfg_plus)
 
   # Work for minus side
   j_to_x_minus = lazy_map(Reindex(i_to_x),j_to_i.minus)
-  j_to_cfg_minus = lazy_map(ConfigMap(ForwardDiff.jacobian),j_to_x_minus)
+  j_to_cfg_minus = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),j_to_x_minus)
   j_to_result_minus_dense = lazy_map(AutoDiffMap(ForwardDiff.jacobian),
                                      j_to_ydual_minus_dense,j_to_x_minus,j_to_cfg_minus)