Why is the kwarg `dims` not being constant-propagated?

General Usage
,
2

Interesting, if we add call site @inline then everything works well.

julia> f(a) = @inline eachslice(a; dims = 1)
f (generic function with 1 method)

julia> @code_warntype f(a)
MethodInstance for f(::Array{Float64, 3})
  from f(a)
     @ Main REPL[37]:1
Arguments
  #self#::Core.Const(f)
  a::Array{Float64, 3}
Locals
  val::Slices{Array{Float64, 3}, Tuple{Int64, Colon, Colon}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 2, Array{Float64, 3}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true}, 1}
Body::Slices{Array{Float64, 3}, Tuple{Int64, Colon, Colon}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 2, Array{Float64, 3}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true}, 1}
1 ─      nothing
│   %2 = (:dims,)::Core.Const((:dims,))
│   %3 = Core.apply_type(Core.NamedTuple, %2)::Core.Const(NamedTuple{(:dims,)})
│   %4 = Core.tuple(1)::Core.Const((1,))
│   %5 = (%3)(%4)::Core.Const((dims = 1,))
│   %6 = Core.kwfunc(Main.eachslice)::Core.Const(Base.var"#eachslice##kw"())
│        (val = (%6)(%5, Main.eachslice, a))
│        nothing
└──      return val::Core.PartialStruct(Slices{Array{Float64, 3}, Tuple{Int64, Colon, Colon}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 2, Array{Float64, 3}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, 
true}, 1}, Any[Array{Float64, 3}, Core.Const((1, Colon(), Colon())), Tuple{Base.OneTo{Int64}}])

Edit: Maybe we should not force inline the kernal function. A single Base.@constprop :aggressive would make our compiler much happier. On a fresh REPL

julia> f(a) = eachslice(a; dims = 1)
f (generic function with 1 method)

julia> a = randn(5,3,2);

julia> Base.return_types(f, Base.typesof(a))
1-element Vector{Any}:
 Slices{Array{Float64, 3}, _A, Tuple{Base.OneTo{Int64}}, _B, 1} where {_A, _B}

julia> @eval Base Base.@constprop :aggressive function _eachslice(A::AbstractArray{T,N}, dims::NTuple{M,Integer}, drop::Bool) where {T,N,M}
           _slice_check_dims(N,dims...)
               if drop
               # if N = 4, dims = (3,1) then
               # axes = (axes(A,3), axes(A,1))
               # slicemap = (2, :, 1, :)
               ax = map(dim -> axes(A,dim), dims)
               slicemap = ntuple(dim -> something(findfirst(isequal(dim), dims), (:)), N)
               return Slices(A, slicemap, ax)
           else
               # if N = 4, dims = (3,1) then
               # axes = (axes(A,1), OneTo(1), axes(A,3), OneTo(1))
               # slicemap = (1, :, 3, :)
                       ax = ntuple(dim -> dim in dims ? axes(A,dim) : unitaxis(A), N)
               slicemap = ntuple(dim -> dim in dims ? dim : (:), N)
               return Slices(A, slicemap, ax)
           end
       end
_eachslice (generic function with 2 methods)

julia> Base.return_types(f, Base.typesof(a))
1-element Vector{Any}:
 Slices{Array{Float64, 3}, Tuple{Int64, Colon, Colon}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 2, Array{Float64, 3}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true}, 1}
1 Like