Why this function involving eachslice is type unstable?

General Usage Performance
1

I need to perform iteratively some operations slicewise in a multidimensional array, and I cannot make it type stable.

Here is an MWE in follow-up to the examples from eachslice is type unstable · Issue #39639 · JuliaLang/julia · GitHub

julia> arr = reshape(collect(1:27),3,3,3);

julia> function f3(arr, d)
           for (i,x) in enumerate(eachslice(arr; dims=d))
           @show  sum(abs2,x)
           end
       end
f3 (generic function with 1 method)

julia> function f4(arr)
           for (i,x) in enumerate(eachslice(arr; dims=(1,3)))
           @show sum(abs2,x)
           end
       end
f4 (generic function with 1 method)

julia> @code_warntype f3(arr,(1,3))
MethodInstance for f3(::Array{Int64, 3}, ::Tuple{Int64, Int64})
  from f3(arr, d) @ Main REPL[28]:1
Arguments
  #self#::Core.Const(Main.f3)
  arr::Array{Int64, 3}
  d::Tuple{Int64, Int64}
Locals
  @_4::Union{Nothing, Tuple{Tuple{Int64, SubArray{Int64, N, Array{Int64, 3}} where N}, Tuple{Int64, Tuple{CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, CartesianIndex{2}}}}}
  value::Any
  @_6::Int64
  x::SubArray{Int64, N, Array{Int64, 3}} where N
  i::Int64
Body::Nothing
...

julia> @code_warntype f4(arr)
MethodInstance for f4(::Array{Int64, 3})
  from f4(arr) @ Main REPL[29]:1
Arguments
  #self#::Core.Const(Main.f4)
  arr::Array{Int64, 3}
Locals
  @_3::Union{Nothing, Tuple{Tuple{Int64, SubArray{Int64, 1, Array{Int64, 3}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Tuple{Int64, Tuple{CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, CartesianIndex{2}}}}}
  value::Int64
  @_5::Int64
  x::SubArray{Int64, 1, Array{Int64, 3}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}, Int64}, true}
  i::Int64
...

As you can see, function f3 is type unstable, and function f4 is OK, and I cannot understand this.

How can I make f3 type-stable?

1 Like
2

eachslice with multiple dims currently depends upon constant propagation to infer the resulting dimensionality — it internally uses an idiom that could skip a repeated dimension, but first asserts that no dimensions are repeated. And that’s just one step too far for inference to see without actually knowing the constant values (and that no dims are repeated).

So that’s why the literal (1,3) behaves differently from the arg dims. How to fix it? I think eachslice itself could probably be further improved internally, but without that, at a end-user-level, you’d “just” need to get dims to constant-propagate — often inlining can make that happen, but it can be touchy.

2 Likes
3

Thanks a lot! I haven’t succeeded in experiments with @inline or the tricks I’ve seen in other branches here or in the implementation of _eachslice itself, though. Hence, I now use different clumsy workarounds involving reshaping, permuted dimensions and CartesianIndexings, depending on the situation.

Some side note: after 5 years of using julia for my needs, I’m slowly starting to lose that initial feeling of it being “as easy as Python and as fast as C”, primarily because of some unexpected small things, like (type stable) tuple differences, or the situation above, which make the straightforward implementation quite slow, and optimised code less understandable for people without computer-science background and more similar to a piece of C code.