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

**URL:** <https://discourse.julialang.org/t/why-is-the-kwarg-dims-not-being-constant-propagated/83610>\
**Category:** General Usage\
**Tags:** question, constant-propagation\
**Created:** [July 1, 2022, 11:26am UTC](https://discourse.julialang.org/t/why-is-the-kwarg-dims-not-being-constant-propagated/83610 "2022-07-01T11:26:32Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![jishnub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jishnub/32/33620_2.png) [@jishnub](https://discourse.julialang.org/u/jishnub)\
**Post date:** [July 1, 2022, 11:26am UTC](https://discourse.julialang.org/t/why-is-the-kwarg-dims-not-being-constant-propagated/83610/1 "2022-07-01T11:26:32Z")

</div>

This doesn’t seem to be inferred

```julia
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
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?

---

<div class="post-metadata">

**Author:** ![N5N3](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/n5n3/32/17663_2.png) [@N5N3](https://discourse.julialang.org/u/N5N3)\
**Post date:** [July 1, 2022, 12:46pm UTC](https://discourse.julialang.org/t/why-is-the-kwarg-dims-not-being-constant-propagated/83610/2 "2022-07-01T12:46:21Z")

</div>

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

```julia
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
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}

```

---

<div class="post-metadata">

**Author:** ![jishnub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jishnub/32/33620_2.png) [@jishnub](https://discourse.julialang.org/u/jishnub)\
**Post date:** [July 1, 2022, 1:50pm UTC](https://discourse.julialang.org/t/why-is-the-kwarg-dims-not-being-constant-propagated/83610/3 "2022-07-01T13:50:58Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![N5N3](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/n5n3/32/17663_2.png) [@N5N3](https://discourse.julialang.org/u/N5N3)\
**Post date:** [July 2, 2022, 2:07am UTC](https://discourse.julialang.org/t/why-is-the-kwarg-dims-not-being-constant-propagated/83610/4 "2022-07-02T02:07:52Z")

</div>

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](https://github.com/JuliaLang/julia/pull/38080#issue-723856855).

> 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)

---

<div class="post-metadata">

**Author:** ![jishnub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jishnub/32/33620_2.png) [@jishnub](https://discourse.julialang.org/u/jishnub)\
**Post date:** [July 4, 2022, 6:42am UTC](https://discourse.julialang.org/t/why-is-the-kwarg-dims-not-being-constant-propagated/83610/5 "2022-07-04T06:42:09Z")

</div>

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

```julia
@eval Base Base.@constprop :aggressive function

```

works,

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

```

leads to inference failing as before.

---

<div class="post-metadata">

**Author:** ![aviatesk](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/aviatesk/32/7610_2.png) [@aviatesk](https://discourse.julialang.org/u/aviatesk)\
**Post date:** [July 5, 2022, 3:03am UTC](https://discourse.julialang.org/t/why-is-the-kwarg-dims-not-being-constant-propagated/83610/6 "2022-07-05T03:03:13Z")

</div>

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](https://github.com/JuliaLang/julia/pull/45041) and then this kind of difficulty around constprop/inlining on kwfuncs will be solved.