use StaticArraysCore

[mcabbott]

Closes #591. Should reduce loading time by a factor of 5 or so. Needs JuliaDiff/DiffResults.jl#22

Needs Julia 1.6. What thoughts on dropping Julia 1.0 support that in a patch release?

The version with #481 is not yet tagged. If that does throw up any surprises, then best to be able to make a patch release correcting them... so we should require Julia 1.6 before tagging that, rather than just after?

Master is 0.11-DEV, and requires Julia 1.6.

[mcabbott]

With these two PRs:

julia> @time_imports using ForwardDiff
      6.3 ms  IrrationalConstants
      0.9 ms  DiffRules
      2.7 ms  StaticArraysCore
      2.3 ms  DiffResults
     19.3 ms  Preferences
      1.2 ms  OpenLibm_jll
      0.3 ms  NaNMath
      0.1 ms  Compat
     82.6 ms  ChainRulesCore
      3.6 ms  DocStringExtensions 68.58% compilation time
      0.7 ms  ChangesOfVariables
      0.6 ms  InverseFunctions
      0.7 ms  LogExpFunctions
      0.3 ms  JLLWrappers
      2.9 ms  OpenSpecFun_jll 76.75% compilation time
     15.5 ms  SpecialFunctions
     27.6 ms  MacroTools
      0.3 ms  CommonSubexpressions
    101.4 ms  ForwardDiff

Before, StaticArrays completely dominates:

julia> @time_imports using ForwardDiff
      6.4 ms  IrrationalConstants
      0.7 ms  DiffRules
      2.4 ms  StaticArraysCore
    608.7 ms  StaticArrays
      2.4 ms  DiffResults
     20.7 ms  Preferences
      1.2 ms  OpenLibm_jll
...

[mcabbott]

CI gives errors like this, which don't show up when running simple examples locally:

  ...testing specialized StaticArray codepaths
Hessians: Error During Test at /home/runner/work/ForwardDiff.jl/ForwardDiff.jl/test/HessianTest.jl:106
  Test threw exception
  Expression: ForwardDiff.hessian(prod, sx) == actual
  MethodError: no method matching length(::Type{StaticArraysCore.SMatrix{3, 3, Float64, 9}})
  The applicable method may be too new: running in world age 32433, while current world is 32506.
  Closest candidates are:
    length(::Type{SA}) where SA<:(Union{LinearAlgebra.Adjoint{T, <:Union{StaticArraysCore.StaticArray{Tuple{var"#s2"}, T, 1} where var"#s2", ...) at ~/.julia/packages/StaticArrays/NOLon/src/abstractarray.jl:2 (method too new to be called from this world context.)
    length(::Union{LinearAlgebra.Adjoint{T, <:Union{StaticArraysCore.StaticArray{...) at ~/.julia/packages/StaticArrays/NOLon/src/abstractarray.jl:1 (method too new to be called from this world context.)
    ...
  Stacktrace:
   [1] #s24#84
     @ ~/work/ForwardDiff.jl/ForwardDiff.jl/src/apiutils.jl:22 [inlined]
   [2] var"#s24#84"(T::Any, ::Any, #unused#::Any, x::Any)
     @ ForwardDiff ./none:0
   [3] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
     @ Core ./boot.jl:582
   [4] static_dual_eval(#unused#::Type{ForwardDiff.Tag{ForwardDiff.var"#118#119"{typeof(prod)}, Float64}}, f::ForwardDiff.var"#118#119"{typeof(prod)}, x::StaticArraysCore.SMatrix{3, 3, Float64, 9})
     @ ForwardDiff ~/work/ForwardDiff.jl/ForwardDiff.jl/src/apiutils.jl:32
   [5] vector_mode_jacobian(f::Function, x::StaticArraysCore.SMatrix{3, 3, Float64, 9})
     @ ForwardDiff ~/work/ForwardDiff.jl/ForwardDiff.jl/src/jacobian.jl:185
   [6] jacobian
     @ ~/work/ForwardDiff.jl/ForwardDiff.jl/src/jacobian.jl:85 [inlined]
   [7] hessian(f::Function, x::StaticArraysCore.SMatrix{3, 3, Float64, 9})
     @ ForwardDiff ~/work/ForwardDiff.jl/ForwardDiff.jl/src/hessian.jl:71
   [8] macro expansion
     @ /opt/hostedtoolcache/julia/1.8.1/x64/share/julia/stdlib/v1.8/Test/src/Test.jl:464 [inlined]
     ```

[devmotion]

My suspicion is that the type piracy in StaticArrays is actually a problem here: it defines functions (that it does not own) for types in StaticArraysCore (that it does not own) which are not available when the @generated functions in ForwardDiff are defined. IIRC I saw similar issues related to generated functions in DiffEqBase a while ago. It might be fine locally if you e.g. load StaticArrays before ForwardDiff.

[mcabbott]

Yes that's my guess too. I did not try very hard to replicate locally, but I did use ForwardDiff first, before loading StaticArrays.

Scrolling through it looks like all are from length(::Type{...}), so just defining that in SACore ought to solve this?

[devmotion]

Scrolling through it looks like all are from length(::Type{...}), so just defining that in SACore ought to solve this?

Yes, seems that might be sufficient for ForwardDiff. However, I guess these issues could also occur in other packages for different functions and hence ideally StaticArrays would fix these type piracies.

[fredrikekre]

The only way to fix (all of) them would be to merge the two packages again, no?

[devmotion]

Hmm, I am not familiar with the internals of StaticArrays but from a quick look it indeed seems that most of StaticArrays is type piracy currently - and hence, as far as I understand, potentially problematic when working with @generated functions. If one actually tries to fix that, maybe one approach could be to move basic types and definitions for Base functions to something like StaticArraysCore and definitions for stdlibs to separate packages (e.g. StaticArraysLinearAlgebra etc.)? Of course, StaticArraysLinearAlgebra would still contain type piracy but maybe the split would be clearer and it would be easier to avoid issues as in this PR here since downstream users would know if they need StaticArraysLinearAlgebra (but not the whole StaticArrays package) or if StaticArraysCore would be sufficient.

[thchr]

This seems to run counter to most of the point of having a *Core.jl package: the parent package must commit a great deal of type-piracy. For StaticArrays, it must be nearly 100% since most of its functionality comes from extending Base/stdlib functions.

We could in principle put all @generated methods into StaticArraysCore.jl but that will come at the cost of increasing its load-time.
It seems a better path to try to find out how many of the relevant @generated methods indeed still need to be @generated.

[devmotion]

Sure, I see that one has to commit type piracy somewhere, so the ideal I mentioned above won't be possible.

However, with the current setup to me it seems that one can't use Base functionality of static arrays inside of @generated functions when only depending on StaticArraysCore. AFAICT the problem are not the generated functions in StaticArrays but generated functions in downstream packages, similar to the behaviour observed in SciML/DiffEqBase.jl#507 (comment).

Maybe that could be avoided for Base functions by moving all definitions with Base functions to StaticArraysCore (without stdlibs) but maybe already that would increase compilation and loading times of StaticArraysCore too much...

[mcabbott]

seems that one can't use Base functionality of static arrays inside of @generated functions

It might be narrower than that. Functions which act on the instance can't be called until you have one, but functions which accept the type can be called outside the returned quote, and these are the problem. Minimal example:

# Core package
struct M1 end
foo(x::Int) = x

# Foreign package
@generated function gen(x::T) where T
  # :(foo(x))  # no problem
  foo(x)     # problem
  # :(foo($T)) # no problem
end
gen(1)

# Big package
foo(::M1) = 33
foo(::Type{M1}) = 99

# User
gen(M1())

The problem code here uses similar_type and length (although I think Base defines length on a type, so this is not a generic use):

ForwardDiff.jl/src/apiutils.jl

Lines 21 to 25 in c6dc209

	@generated function dualize(::Type{T}, x::StaticArray) where T
	N = length(x)
	dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...)
	V = StaticArrays.similar_type(x, Dual{T,eltype(x),N})
	return quote

[mcabbott]

I think this is done. The above @generated world-age problems are solved by not calling length, and by moving similat_type out of the generated body & into the quote.

[KristofferC]

This could be made into an extension which would also solve the loading time.

[mcabbott]

Maybe there would be less reason to create StaticArraysCore now that extensions exist. But unless there's a plan to deprecate the core package, maybe ForwardDiff may as well use it?

The next heaviest thing is ChainRulesCore, which isn't directly loaded. Perhaps all the packages for which ForwardDiff defines rules out to extensions would be a good idea.

[fredrikekre]

Maybe there would be less reason to create StaticArraysCore now that extensions exist. But unless there's a plan to deprecate the core package, maybe ForwardDiff may as well use it?

Hah, just posted JuliaArrays/StaticArraysCore.jl#15. If people can agree on that plan, it might be easier to not port more packages to using StaticArraysCore.

[mcabbott]

Is it certain the extension feature will make 1.9?

I see there are 22 packages using StaticArraysCore, but I suppose no problem to have some mixture of mechanisms for a while.

[KristofferC]

Is it certain the extension feature will make 1.9?

It's on the backport branch at least.