Inference/constant propagation regression in Julia 1.7

[mateuszbaran]

I've run into a problem that affects Interpolations.jl. Type inference (or more likely constant propagation?) in Julia 1.7.0 gives up earlier than in previous versions of Julia.
Fairly minimal example:

using Interpolations
nx = 10
f1(x) = sin((x-3)*2pi/(nx-1) - 1)
A = Float64[f1(x) for x in 1:nx]
itp = interpolate(A, BSpline(Constant()))
etp = extrapolate(itp, 0.0)
@code_warntype ndims(etp)

With the result:

julia> @code_warntype ndims(etp)
MethodInstance for ndims(::Interpolations.FilledExtrapolation{Float64, 1, Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}}}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Float64})
  from ndims(itp::AbstractInterpolation) in Interpolations at /home/mateusz/.julia/dev/Interpolations/src/Interpolations.jl:158
Arguments
  #self#::Core.Const(ndims)
  itp::Interpolations.FilledExtrapolation{Float64, 1, Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}}}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Float64}
Body::Any
1 ─ %1 = Interpolations.typeof(itp)::Core.Const(Interpolations.FilledExtrapolation{Float64, 1, Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}}}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Float64})
│   %2 = Interpolations.ndims(%1)::Any
└──      return %2

While for example using Julia 1.6.1:

julia> @code_warntype ndims(etp)
Variables
  #self#::Core.Const(ndims)
  itp::Interpolations.FilledExtrapolation{Float64, 1, Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}}}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Float64}

Body::Int64
1 ─ %1 = Interpolations.typeof(itp)::Core.Const(Interpolations.FilledExtrapolation{Float64, 1, Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}}}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Float64})
│   %2 = Interpolations.ndims(%1)::Core.Const(1)
└──      return %2

Forcing specialization for one more method of ndims solves the problem but I don't know if this is the right solution.

Versioninfo for reference:

julia> versioninfo()
Julia Version 1.7.0
Commit 3bf9d17731 (2021-11-30 12:12 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Core(TM) i7-4800MQ CPU @ 2.70GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-12.0.1 (ORCJIT, haswell)
Environment:
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 

ref: JuliaMath/Interpolations.jl#468

[timholy]

Worth checking if #43347 fixes this too, CC @aviatesk

[aviatesk]

Optimization isn't able to derive new inter-procedural information like this (return type), so it seems to be a regression in inference. Will look into this when I have time.

[N5N3]

Replacing

ndims(::Type{AbstractInterpolation{T,N,IT}}) where {T,N,IT<:DimSpec{InterpolationType}} = N
ndims(::Type{ITP}) where {ITP<:AbstractInterpolation} = ndims(supertype(ITP))

with

ndims(::Type{<:AbstractInterpolation{T,N,<:DimSpec{InterpolationType}}}) where {T,N} = N

could fix it.
The original definition leads to a recursive call of:

julia> t = typeof(etp)
Interpolations.FilledExtrapolation{Float64, 1, Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}}}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Float64}
julia> t = supertype(t)
AbstractExtrapolation{Float64, 1, Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}}}, BSpline{Constant{Nearest, Throw{OnGrid}}}}
julia> t = supertype(t)
Interpolations.AbstractInterpolationWrapper{Float64, 1, Interpolations.BSplineInterpolation{Float64, 1, Vector{Float64}, BSpline{Constant{Nearest, Throw{OnGrid}}}, Tuple{Base.OneTo{Int64}}}, BSpline{Constant{Nearest, Throw{OnGrid}}}}
julia> t = supertype(t)
AbstractInterpolation{Float64, 1, BSpline{Constant{Nearest, Throw{OnGrid}}}}

So this seems also related with the type-more-complex check?

[maleadt]

Bisected to:

e8caa2750969b1f4e8e41431ca8d1840aa8ac51b is the first bad commit
commit e8caa2750969b1f4e8e41431ca8d1840aa8ac51b
Author: Jameson Nash <vtjnash@gmail.com>
Date:   Mon Nov 1 15:10:32 2021 -0400

    inference: relax type_more_complex for Type
    
    If the object is not a potential Type{...}, we can allow greater
    flexibility in recursion.
    
    (cherry picked from commit 1daaf473d1c268ce8b54f47bd7b30c2f74d66b99)

Same as #43296. cc @vtjnash

[maleadt]

MWE:

abstract type e{a,j} <: AbstractArray{a,j} end
abstract type b{a,j,c,d} <: e{a,j} end
struct h{a,j,f,d,i} <: b{a,j,f,d} end
Base.ndims(::Type{f}) where f<:e = ndims(supertype(f))
Base.ndims(g::e) = ndims(typeof(g))
Base.return_types(ndims, (h{Any, 0, Any, Int64, Any},))

Int64 on 1.6, Any on 1.7

[vtjnash]

Note that OP is fixed, due to changes to ndims, but the MWE above still works.