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

General Usage
,
1

This doesn’t seem to be inferred

julia> a = rand(1,1); 

julia> @code_warntype (a -> eachslice(a, dims=1))(a)
MethodInstance for (::var"#13#14")(::Matrix{Float64})
  from (::var"#13#14")(a)
     @ Main REPL[23]:1
Arguments
  #self#::Core.Const(var"#13#14"())
  a::Matrix{Float64}
Body::Slices{Matrix{Float64}, _A, Tuple{Base.OneTo{Int64}}, _B, 1} where {_A, _B}
1 ─ %1 = (:dims,)::Core.Const((:dims,))
│   %2 = Core.apply_type(Core.NamedTuple, %1)::Core.Const(NamedTuple{(:dims,)})
│   %3 = Core.tuple(1)::Core.Const((1,))
│   %4 = (%2)(%3)::Core.Const((dims = 1,))
│   %5 = Core.kwfunc(Main.eachslice)::Core.Const(Base.var"#eachslice##kw"())
│   %6 = (%5)(%4, Main.eachslice, a)::Slices{Matrix{Float64}, _A, Tuple{Base.OneTo{Int64}}, _B, 1} where {_A, _B}
└──      return %6

whereas this internal function is

julia> @code_warntype (a -> Base._eachslice(a, 1, true))(a)
MethodInstance for (::var"#15#16")(::Matrix{Float64})
  from (::var"#15#16")(a)
     @ Main REPL[24]:1
Arguments
  #self#::Core.Const(var"#15#16"())
  a::Matrix{Float64}
Body::RowSlices{Matrix{Float64}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}
1 ─ %1 = Base._eachslice::Core.Const(Base._eachslice)
│   %2 = (%1)(a, 1, true)::Core.PartialStruct(RowSlices{Matrix{Float64}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}, Any[Matrix{Float64}, Core.Const((1, Colon())), Tuple{Base.OneTo{Int64}}])
└──      return %2

Since the former calls the latter, shouldn’t constant-propagation of dims allow this to be inferred?

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
3

Intresting, I wonder if this should be added to Base? I don’t really know the consequences of this

4

I’m also not familiar with it. But based on the following motivation from Add a macro to opt into aggressive constprop by Keno · Pull Request #38080 · JuliaLang/julia · GitHub.

Right now aggressive constprop is essentially tied to the inlining
threshold (or to function’s name being getproperty or setproperty!
respectively), which can be both somewhat brittle if the inlining cost
changes and insufficient when you do really know that const prop
would be beneficial even if the function is not inlineable.

The fragility of @inline based const propagation has been predicted. So this seems a good example why we need @constprop (especially when you know the function would be inlined automaticly when const prop happen)

3 Likes
5

It’s a bit strange that adding @inline to the function definition makes things worse. While

@eval Base Base.@constprop :aggressive function

works,

@eval Base @inline Base.@constprop :aggressive function

leads to inference failing as before.

6

Yes, this is very tricky. Here we actually need both eachslice and _eachslice (and their kwsorter versions as well) to propagate constant information, but if we force inlining on _eachslice then the body of eachslice (or, its kwsorter func actually) turns out to be complex enough to prohibit the constant propagation on eachslice. Ironically, however, if we force constant prop’ on eachslice, then _eachslice will be simplified further using the constant information, and then we will get really successful return type inference and inlining optimization. For now, I’d suggest adding @constprop annotation on and remove @inline from _eachslice.

For the meanwhile, I will try to finish propagate method metadata to keyword sorter methods by JeffBezanson · Pull Request #45041 · JuliaLang/julia · GitHub and then this kind of difficulty around constprop/inlining on kwfuncs will be solved.

4 Likes