Type instabilities in `return_cache(::ConfigMap, x)`
- Dominant language
- Julia
- Stars
- 879
- Forks
- 119
- Avg merge
- 6d 3h
- Merged PRs (30d)
- 4
Description
While working on #1303, I realized that methods such as
```julia
function return_cache(k::ConfigMap{typeof(ForwardDiff.gradient)},x)
cfg = ForwardDiff.GradientConfig(k.tag,x,ForwardDiff.Chunk{length(x)}())
# uncomment line below to find un-inferable cases
#@check try @inferred (y -> ForwardDiff.Chunk{length(y)}())(x); true catch _ false end
return cfg
end
```
are not type-inferable since the chunk size, which is included in the type of `cfg`, is not necessarily inferable from the type of `x`.
This violates the contract of `return_cache`, and prevents `array_cache(a::LazyArray)` from being inferrable whenever in contains autodiffed function. As a consequence, `getindex(cache, a, i)` will do dynamic dispatch if `cache` was not passed in the arguments of `getindex!`s caller (often).
This does not apply for `TensorValue`d / `SArray` valued functions, but the case hits in the tests, e.g. in `PLaplacianTests.jl` when using autodiff to define the `FEOperator`. I suspect this should always happen when one differentiate operators with respect to fefunctions / their dof array that are `::Vector{T}` (and might explain part of Multifield's AD tests slowness).
After some tries, it does not look easy to restrict `x` above to have a type like `VectorValue` or `SVector`. Instead, I suggest adding the chunk size in the `ConfigMap` arguments, because the `ConfigMap`s are created in functions like
```julia
function autodiff_array_gradient(a,i_to_x)
tag = x->ForwardDiff.gradient(a, x)
i_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,tag),i_to_x)
i_to_xdual = lazy_map(DualizeMap(),i_to_cfg,i_to_x)
i_to_ydual = a(i_to_xdual)
i_to_result = lazy_map(AutoDiffMap(),i_to_cfg,i_to_ydual)
return i_to_result
end
```
where we should be able to measure the length of `x`. This should achieve always storing the chunk size in the `::LazyArray` type.
Thoughts or suggestions @JordiManyer @zjwegert ?
Contributor guide
Research direction
Start with return_cache(::ConfigMap, x), array_cache(a::LazyArray), and the autodiff_array_gradient path described in the issue. Reproduce the inference problem in PLaplacianTests.jl using the autodiff-defined FEOperator, then trace how ConfigMap carries the chunk size. Done means the relevant cache construction is type-inferable and avoids dynamic dispatch in getindex(cache, a, i).
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100