Instability detected in numerical field construction

[henry2004y]

Consider this mixed type field case:

using TestParticle
using StaticArrays
using Meshes

x = range(-10, 10, length=15);
y = range(-10, 10, length=20);
z = range(-10, 10, length=25);
B = fill(0.0, 3, length(x), length(y), length(z)); # [T]
F = fill(0.0, 3, length(x), length(y), length(z)); # [N]

B[3,:,:,:] .= 1e-11;
E_field(r) = SA[0, 5e-11, 0]  # [V/M]
F[1,:,:,:] .= 9.10938356e-42;

Δx = x[2] - x[1]
Δy = y[2] - y[1]
Δz = z[2] - z[1]

mesh = CartesianGrid((length(x)-1, length(y)-1, length(z)-1),
   (x[1], y[1], z[1]),
   (Δx, Δy, Δz))

 @code_warntype prepare(mesh, E_field, B, F; species=Electron)

MethodInstance for Core.kwcall(::@NamedTuple{species::TestParticle.Species}, ::typeof(prepare), ::CartesianGrid{3, Float64}, ::typeof(E_field), ::Array{Float64, 4}, ::Array{Float64, 4})
  from kwcall(::NamedTuple, ::typeof(prepare), grid::CartesianGrid, E::TE, B::TB, F::TF) where {TE, TB, TF} @ TestParticle D:\Computer\TestParticle.jl\src\TestParticle.jl:336
Static Parameters
  TE = typeof(E_field)
  TB = Array{Float64, 4}
  TF = Array{Float64, 4}
Arguments
  _::Core.Const(Core.kwcall)
  @_2::@NamedTuple{species::TestParticle.Species}
  @_3::Core.Const(TestParticle.prepare)
  grid::CartesianGrid{3, Float64}
  E::Core.Const(E_field)
  B::Array{Float64, 4}
  F::Array{Float64, 4}
Locals
  species::Union{}
  q::Union{}
  m::Union{}
  order::Union{}
  @_12::Union{Float64, Int64, TestParticle.Species}
Body::Tuple{Float64, Float64, TestParticle.Field{false, typeof(E_field)}, TestParticle.Field{false, TestParticle.var"#get_field#3"{interpz, interpy, interpx}} where {interpz, interpy, interpx}, TestParticle.Field{false, TestParticle.var"#get_field#3"{interpz, interpy, interpx}} where {interpz, interpy, interpx}}
1 ──       ......
15 ─       Core.Const(:(Base.kwerr(@_2, @_3, grid, E, B, F)))
16 ┄ %62 = TestParticle.:(var"#prepare#7")(%17, %29, %41, %53, @_3, grid, E, B, F)::Tuple{Float64, Float64, TestParticle.Field{false, typeof(E_field)}, TestParticle.Field{false, TestParticle.var"#get_field#3"{interpz, interpy, interpx}} where {interpz, interpy, interpx}, TestParticle.Field{false, TestParticle.var"#get_field#3"{interpz, interpy, interpx}} where {interpz, interpy, interpx}}
└───       return %62

TestParticle.var"#get_field#3"{interpz, interpy, interpx}} where {interpz, interpy, interpx} indicates type instability.

After investigating the issue further, I realized that it is related to the ScaledInterpolation type from Interpolations.jl:

   interpx = scale(itpx, gridx, gridy, gridz)
   interpy = scale(itpy, gridx, gridy, gridz)
   interpz = scale(itpz, gridx, gridy, gridz)

   # Return field value at a given location.
   function get_field(xu)
      r = @view xu[1:3]

      return SA[interpx(r...), interpy(r...), interpz(r...)]
   end

If I replace interpx with a function, it will be type stable. For example, as a minimal demo,

function foo1(a)
   x = a + 1
   y = a + 2

   foo2(x) = x + 3
   foo3(x) = x + 4

   function inner(b)
      return [b, foo2(b), foo3(b)]
   end

   return inner
end

f = foo1(1)
@code_warntype f(1)

This is type stable!.

My feeling is that this issue is related to the captured problem. I have no idea how to solve this...

[henry2004y]

Maybe this is not an issue for performance in our case, but somehow @code_warntype list it as a warning?

I tried to manually indicate the type:

interpx::ScaledInterpolation{TA, 3, FilledExtrapolation{TA, 3, BSplineInterpolation{TA, 3, Array{TA, 3}, BSpline{Linear{Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, BSpline{Linear{Throw{OnGrid}}}, TA}, BSpline{Linear{Throw{OnGrid}}}, Tuple{TG, TG, TG}} = scale(itpx, gridx, gridy, gridz)
interpy::ScaledInterpolation{TA, 3, FilledExtrapolation{TA, 3, BSplineInterpolation{TA, 3, Array{TA, 3}, BSpline{Linear{Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, BSpline{Linear{Throw{OnGrid}}}, TA}, BSpline{Linear{Throw{OnGrid}}}, Tuple{TG, TG, TG}} = scale(itpy, gridx, gridy, gridz)
interpz::ScaledInterpolation{TA, 3, FilledExtrapolation{TA, 3, BSplineInterpolation{TA, 3, Array{TA, 3}, BSpline{Linear{Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, BSpline{Linear{Throw{OnGrid}}}, TA}, BSpline{Linear{Throw{OnGrid}}}, Tuple{TG, TG, TG}} = scale(itpz, gridx, gridy, gridz)

Now @code_warntype does not show anything red, but the performance is almost the same:

# Manually mark the types
julia> @benchmark solve($prob_ip, Tsit5(); save_idxs=[1,2,3])
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  18.500 μs …  2.191 ms  ┊ GC (min … max): 0.00% … 97.82%
 Time  (median):     20.900 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   23.986 μs ± 36.758 μs  ┊ GC (mean ± σ):  2.59% ±  1.69%

  ▃▆██▇▅▄▁▁▁▄▅▆▆▆▄▂▁▂▂▂▂▁                    ▁ ▁              ▂
  █████████████████████████▇▆▅▅▅▆███▆▆▅▅▄▄▆▇████▇▆▅▅▅▃▆▇▇█▆▅▅ █
  18.5 μs      Histogram: log(frequency) by time        47 μs <

 Memory estimate: 20.62 KiB, allocs estimate: 129.
# Current master
julia> @benchmark solve($prob_ip, Tsit5(); save_idxs=[1,2,3])
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  19.500 μs …  2.222 ms  ┊ GC (min … max): 0.00% … 97.71%
 Time  (median):     22.800 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   25.032 μs ± 37.235 μs  ┊ GC (mean ± σ):  2.52% ±  1.69%

  ▄▇█▇▄▂▁▁▂▄▆▇▆▅▄▂▂▂▂▂▂▁▁▁▁▁▁                                 ▂
  █████████████████████████████▆██▇▇▆▄▆▅▄▅▅▇▇▇▆▅▅▅▄▅▅▄▆▆▅▅▃▅▅ █
  19.5 μs      Histogram: log(frequency) by time      49.1 μs <

 Memory estimate: 20.62 KiB, allocs estimate: 129.