Detangle Slice and fix mixed-dimension in-place reductions (#28941)

* Detangle Slice between whole dimensions and axes

We use axes in many downstream computations that may not re-index directly into the original array, so we add a second type parameter to `Slice` that is turned on when converting `:` in indexing expressions -- and really only `SubArray` cares about.

* Reduce allocations for in-place reductions and fix mixed-dimensionality edge-cases

Alleviates #28928 but does not completely remove allocations due to the allocation of the view that gets passed to `map!`.

* Introduce a whole new IdentityUnitRange

that we will encourage offset array implementations to use instead of Base.Slice

* Test that maximum! allocates less

Author: mbauman

base/indices.jl

@@ -303,7 +303,8 @@ _maybetail(t::Tuple) = tail(t)
 """
    Slice(indices)
 
-Represent an AbstractUnitRange of indices as a vector of the indices themselves.
+Represent an AbstractUnitRange of indices as a vector of the indices themselves,
+with special handling to signal they represent a complete slice of a dimension (:).
 
 Upon calling `to_indices`, Colons are converted to Slice objects to represent
 the indices over which the Colon spans. Slice objects are themselves unit
@@ -315,9 +316,9 @@ struct Slice{T<:AbstractUnitRange} <: AbstractUnitRange{Int}
     indices::T
 end
 Slice(S::Slice) = S
-axes(S::Slice) = (S,)
-unsafe_indices(S::Slice) = (S,)
-axes1(S::Slice) = S
+axes(S::Slice) = (IdentityUnitRange(S.indices),)
+unsafe_indices(S::Slice) = (IdentityUnitRange(S.indices),)
+axes1(S::Slice) = IdentityUnitRange(S.indices)
 axes(S::Slice{<:OneTo}) = (S.indices,)
 unsafe_indices(S::Slice{<:OneTo}) = (S.indices,)
 axes1(S::Slice{<:OneTo}) = S.indices
@@ -333,6 +334,38 @@ getindex(S::Slice, i::StepRange{<:Integer}) = (@_inline_meta; @boundscheck check
 show(io::IO, r::Slice) = print(io, "Base.Slice(", r.indices, ")")
 iterate(S::Slice, s...) = iterate(S.indices, s...)
 
+
+"""
+   IdentityUnitRange(range::AbstractUnitRange)
+
+Represent an AbstractUnitRange `range` as an offset vector such that `range[i] == i`.
+
+`IdentityUnitRange`s are frequently used as axes for offset arrays.
+"""
+struct IdentityUnitRange{T<:AbstractUnitRange} <: AbstractUnitRange{Int}
+    indices::T
+end
+IdentityUnitRange(S::IdentityUnitRange) = S
+# IdentityUnitRanges are offset and thus have offset axes, so they are their own axes... but
+# we need to strip the wholedim marker because we don't know how they'll be used
+axes(S::IdentityUnitRange) = (S,)
+unsafe_indices(S::IdentityUnitRange) = (S,)
+axes1(S::IdentityUnitRange) = S
+axes(S::IdentityUnitRange{<:OneTo}) = (S.indices,)
+unsafe_indices(S::IdentityUnitRange{<:OneTo}) = (S.indices,)
+axes1(S::IdentityUnitRange{<:OneTo}) = S.indices
+
+first(S::IdentityUnitRange) = first(S.indices)
+last(S::IdentityUnitRange) = last(S.indices)
+size(S::IdentityUnitRange) = (length(S.indices),)
+length(S::IdentityUnitRange) = length(S.indices)
+unsafe_length(S::IdentityUnitRange) = unsafe_length(S.indices)
+getindex(S::IdentityUnitRange, i::Int) = (@_inline_meta; @boundscheck checkbounds(S, i); i)
+getindex(S::IdentityUnitRange, i::AbstractUnitRange{<:Integer}) = (@_inline_meta; @boundscheck checkbounds(S, i); i)
+getindex(S::IdentityUnitRange, i::StepRange{<:Integer}) = (@_inline_meta; @boundscheck checkbounds(S, i); i)
+show(io::IO, r::IdentityUnitRange) = print(io, "Base.IdentityUnitRange(", r.indices, ")")
+iterate(S::IdentityUnitRange, s...) = iterate(S.indices, s...)
+
 """
     LinearIndices(A::AbstractArray)
 

base/multidimensional.jl

@@ -472,7 +472,6 @@ index_dimsum() = ()
 index_lengths() = ()
 @inline index_lengths(::Real, rest...) = (1, index_lengths(rest...)...)
 @inline index_lengths(A::AbstractArray, rest...) = (length(A), index_lengths(rest...)...)
-@inline index_lengths(A::Slice, rest...) = (length(axes1(A)), index_lengths(rest...)...)
 
 # shape of array to create for getindex() with indices I, dropping scalars
 # returns a Tuple{Vararg{AbstractUnitRange}} of indices

base/reducedim.jl

@@ -4,7 +4,7 @@
 
 # for reductions that expand 0 dims to 1
 reduced_index(i::OneTo) = OneTo(1)
-reduced_index(i::Slice) = first(i):first(i)
+reduced_index(i::Union{Slice, IdentityUnitRange}) = first(i):first(i)
 reduced_index(i::AbstractUnitRange) =
     throw(ArgumentError(
 """
@@ -218,17 +218,20 @@ Extract first entry of slices of array A into existing array R.
 """
 copyfirst!(R::AbstractArray, A::AbstractArray) = mapfirst!(identity, R, A)
 
-function mapfirst!(f, R::AbstractArray, A::AbstractArray)
+function mapfirst!(f, R::AbstractArray, A::AbstractArray{<:Any,N}) where {N}
     lsiz = check_reducedims(R, A)
-    iA = axes(A)
-    iR = axes(R)
-    t = []
-    for i in 1:length(iR)
-        iAi = iA[i]
-        push!(t, iAi == iR[i] ? iAi : first(iAi))
-    end
+    t = _firstreducedslice(axes(R), axes(A))
     map!(f, R, view(A, t...))
 end
+# We know that the axes of R and A are compatible, but R might have a different number of
+# dimensions than A, which is trickier than it seems due to offset arrays and type stability
+_firstreducedslice(::Tuple{}, a::Tuple{}) = ()
+_firstreducedslice(::Tuple, ::Tuple{}) = ()
+@inline _firstreducedslice(::Tuple{}, a::Tuple) = (_firstslice(a[1]), _firstreducedslice((), tail(a))...)
+@inline _firstreducedslice(r::Tuple, a::Tuple) = (length(r[1])==1 ? _firstslice(a[1]) : r[1], _firstreducedslice(tail(r), tail(a))...)
+_firstslice(i::OneTo) = OneTo(1)
+_firstslice(i::Slice) = Slice(_firstslice(i.indices))
+_firstslice(i) = i[firstindex(i):firstindex(i)]
 
 function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArray)
     lsiz = check_reducedims(R,A)

base/show.jl

@@ -1831,7 +1831,7 @@ dims2string(d) = isempty(d) ? "0-dimensional" :
 
 inds2string(inds) = join(map(_indsstring,inds), '×')
 _indsstring(i) = string(i)
-_indsstring(i::Slice) = string(i.indices)
+_indsstring(i::Union{IdentityUnitRange, Slice}) = string(i.indices)
 
 # anything array-like gets summarized e.g. 10-element Array{Int64,1}
 summary(io::IO, a::AbstractArray) = summary(io, a, axes(a))
@@ -1906,7 +1906,7 @@ function showarg(io::IO, v::SubArray, toplevel)
     print(io, ')')
     toplevel && print(io, " with eltype ", eltype(v))
 end
-showindices(io, ::Slice, inds...) =
+showindices(io, ::Union{Slice,IdentityUnitRange}, inds...) =
     (print(io, ", :"); showindices(io, inds...))
 showindices(io, ind1, inds...) =
     (print(io, ", ", ind1); showindices(io, inds...))

base/sort.jl

@@ -19,7 +19,6 @@ import .Base:
     sort!,
     issorted,
     sortperm,
-    Slice,
     to_indices
 
 export # also exported by Base

base/subarray.jl

@@ -316,7 +316,7 @@ compute_offset1(parent::AbstractVector, stride1::Integer, I::Tuple{AbstractRange
 # linear indexing always starts with 1.
 compute_offset1(parent, stride1::Integer, I::Tuple) =
     (@_inline_meta; compute_offset1(parent, stride1, find_extended_dims(1, I...), find_extended_inds(I...), I))
-compute_offset1(parent, stride1::Integer, dims::Tuple{Int}, inds::Tuple{Slice}, I::Tuple) =
+compute_offset1(parent, stride1::Integer, dims::Tuple{Int}, inds::Tuple{Union{Slice, IdentityUnitRange}}, I::Tuple) =
     (@_inline_meta; compute_linindex(parent, I) - stride1*first(axes(parent, dims[1])))  # index-preserving case
 compute_offset1(parent, stride1::Integer, dims, inds, I::Tuple) =
     (@_inline_meta; compute_linindex(parent, I) - stride1)  # linear indexing starts with 1
@@ -384,7 +384,7 @@ function parentdims(s::SubArray)
     j = 1
     for i = 1:ndims(s.parent)
         r = s.indices[i]
-        if j <= nd && (isa(r,Union{Slice,AbstractRange}) ? sp[i]*step(r) : sp[i]) == sv[j]
+        if j <= nd && (isa(r,AbstractRange) ? sp[i]*step(r) : sp[i]) == sv[j]
             dimindex[j] = i
             j += 1
         end

stdlib/LinearAlgebra/test/adjtrans.jl

@@ -489,7 +489,7 @@ using .Main.OffsetArrays
 
 @testset "offset axes" begin
     s = Base.Slice(-3:3)'
-    @test axes(s) === (Base.OneTo(1), Base.Slice(-3:3))
+    @test axes(s) === (Base.OneTo(1), Base.IdentityUnitRange(-3:3))
     @test collect(LinearIndices(s)) == reshape(1:7, 1, 7)
     @test collect(CartesianIndices(s)) == reshape([CartesianIndex(1,i) for i = -3:3], 1, 7)
     @test s[1] == -3

test/arrayops.jl

@@ -1767,7 +1767,7 @@ end
     @test CartesianIndices(CartesianIndex{3}(1,2,3)) == CartesianIndices((1, 2, 3))
     @test Tuple{}(CartesianIndices{0,Tuple{}}(())) == ()
 
-    R = CartesianIndices(map(Base.Slice, (2:5, 3:5)))
+    R = CartesianIndices(map(Base.IdentityUnitRange, (2:5, 3:5)))
     @test eltype(R) <: CartesianIndex{2}
     @test eltype(typeof(R)) <: CartesianIndex{2}
     @test eltype(CartesianIndices{2}) <: CartesianIndex{2}

test/offsetarray.jl

@@ -81,7 +81,7 @@ for i = 1:9 @test A_3_3[i] == i end
 @test_throws BoundsError A[CartesianIndex(1,1)]
 @test_throws BoundsError S[CartesianIndex(1,1)]
 @test eachindex(A) == 1:4
-@test eachindex(S) == CartesianIndices(axes(S)) == CartesianIndices(map(Base.Slice, (0:1,3:4)))
+@test eachindex(S) == CartesianIndices(axes(S)) == CartesianIndices(map(Base.IdentityUnitRange, (0:1,3:4)))
 
 # LinearIndices
 # issue 27986
@@ -122,13 +122,13 @@ S = view(A, :, 3)
 @test S[0] == 1
 @test S[1] == 2
 @test_throws BoundsError S[2]
-@test axes(S) === (Base.Slice(0:1),)
+@test axes(S) === (Base.IdentityUnitRange(0:1),)
 S = view(A, 0, :)
 @test S == OffsetArray([1,3], (A.offsets[2],))
 @test S[3] == 1
 @test S[4] == 3
 @test_throws BoundsError S[1]
-@test axes(S) === (Base.Slice(3:4),)
+@test axes(S) === (Base.IdentityUnitRange(3:4),)
 S = view(A, 0:0, 4)
 @test S == [3]
 @test S[1] == 3
@@ -147,17 +147,17 @@ S = view(A, :, :)
 @test S[0,4] == S[3] == 3
 @test S[1,4] == S[4] == 4
 @test_throws BoundsError S[1,1]
-@test axes(S) === Base.Slice.((0:1, 3:4))
+@test axes(S) === Base.IdentityUnitRange.((0:1, 3:4))
 # https://github.com/JuliaArrays/OffsetArrays.jl/issues/27
 g = OffsetArray(Vector(-2:3), (-3,))
 gv = view(g, -1:2)
 @test axes(gv, 1) === Base.OneTo(4)
 @test collect(gv) == -1:2
 gv = view(g, OffsetArray(-1:2, (-2,)))
-@test axes(gv, 1) === Base.Slice(-1:2)
+@test axes(gv, 1) === Base.IdentityUnitRange(-1:2)
 @test collect(gv) == -1:2
 gv = view(g, OffsetArray(-1:2, (-1,)))
-@test axes(gv, 1) === Base.Slice(0:3)
+@test axes(gv, 1) === Base.IdentityUnitRange(0:3)
 @test collect(gv) == -1:2
 
 # iteration
@@ -234,26 +234,26 @@ B = similar(A, (3,4))
 @test axes(B) === (Base.OneTo(3), Base.OneTo(4))
 B = similar(A, (-3:3,1:4))
 @test isa(B, OffsetArray{Int,2})
-@test axes(B) === Base.Slice.((-3:3, 1:4))
+@test axes(B) === Base.IdentityUnitRange.((-3:3, 1:4))
 B = similar(parent(A), (-3:3,1:4))
 @test isa(B, OffsetArray{Int,2})
-@test axes(B) === Base.Slice.((-3:3, 1:4))
+@test axes(B) === Base.IdentityUnitRange.((-3:3, 1:4))
 
 # Indexing with OffsetArray indices
 i1 = OffsetArray([2,1], (-5,))
 i1 = OffsetArray([2,1], -5)
 b = A0[i1, 1]
-@test axes(b) === (Base.Slice(-4:-3),)
+@test axes(b) === (Base.IdentityUnitRange(-4:-3),)
 @test b[-4] == 2
 @test b[-3] == 1
 b = A0[1,i1]
-@test axes(b) === (Base.Slice(-4:-3),)
+@test axes(b) === (Base.IdentityUnitRange(-4:-3),)
 @test b[-4] == 3
 @test b[-3] == 1
 v = view(A0, i1, 1)
-@test axes(v) === (Base.Slice(-4:-3),)
+@test axes(v) === (Base.IdentityUnitRange(-4:-3),)
 v = view(A0, 1:1, i1)
-@test axes(v) === (Base.OneTo(1), Base.Slice(-4:-3))
+@test axes(v) === (Base.OneTo(1), Base.IdentityUnitRange(-4:-3))
 
 # copyto! and fill!
 a = OffsetArray{Int}(undef, (-3:-1,))
@@ -340,7 +340,7 @@ a = OffsetArray(a0, (-1,2,3,4,5))
 v = OffsetArray(v0, (-3,))
 @test lastindex(v) == 1
 @test v ≈ v
-@test axes(v') === (Base.OneTo(1),Base.Slice(-2:1))
+@test axes(v') === (Base.OneTo(1),Base.IdentityUnitRange(-2:1))
 @test parent(v) == collect(v)
 rv = reverse(v)
 @test axes(rv) == axes(v)
@@ -356,7 +356,7 @@ A = OffsetArray(rand(4,4), (-3,5))
 @test lastindex(A, 1) == 1
 @test lastindex(A, 2) == 9
 @test A ≈ A
-@test axes(A') === Base.Slice.((6:9, -2:1))
+@test axes(A') === Base.IdentityUnitRange.((6:9, -2:1))
 @test parent(copy(A')) == copy(parent(A)')
 @test collect(A) == parent(A)
 @test maximum(A) == maximum(parent(A))
@@ -498,3 +498,25 @@ end
     A = OffsetArray(reshape(16:-1:1, (4, 4)), (-3,5))
     @test maximum(A, dims=1) == OffsetArray(maximum(parent(A), dims=1), A.offsets)
 end
+
+@testset "in-place reductions with mismatched dimensionalities" begin
+    B = OffsetArray(reshape(1:24, 4, 3, 2), -5, 6, -7)
+    for R in (fill(0, -4:-1), fill(0, -4:-1, 7:7), fill(0, -4:-1, 7:7, -6:-6))
+        @test @inferred(maximum!(R, B)) == reshape(maximum(B, dims=(2,3)), axes(R)) == reshape(21:24, axes(R))
+        @test @allocated(maximum!(R, B)) <= 400
+        @test @inferred(minimum!(R, B)) == reshape(minimum(B, dims=(2,3)), axes(R)) == reshape(1:4, axes(R))
+        @test @allocated(minimum!(R, B)) <= 400
+    end
+    for R in (fill(0, -4:-4, 7:9), fill(0, -4:-4, 7:9, -6:-6))
+        @test @inferred(maximum!(R, B)) == reshape(maximum(B, dims=(1,3)), axes(R)) == reshape(16:4:24, axes(R))
+        @test @allocated(maximum!(R, B)) <= 400
+        @test @inferred(minimum!(R, B)) == reshape(minimum(B, dims=(1,3)), axes(R)) == reshape(1:4:9, axes(R))
+        @test @allocated(minimum!(R, B)) <= 400
+    end
+    @test_throws DimensionMismatch maximum!(fill(0, -4:-1, 7:7, -6:-6, 1:1), B)
+    @test_throws DimensionMismatch minimum!(fill(0, -4:-1, 7:7, -6:-6, 1:1), B)
+    @test_throws DimensionMismatch maximum!(fill(0, -4:-4, 7:9, -6:-6, 1:1), B)
+    @test_throws DimensionMismatch minimum!(fill(0, -4:-4, 7:9, -6:-6, 1:1), B)
+    @test_throws DimensionMismatch maximum!(fill(0, -4:-4, 7:7, -6:-5, 1:1), B)
+    @test_throws DimensionMismatch minimum!(fill(0, -4:-4, 7:7, -6:-5, 1:1), B)
+end

test/reducedim.jl

@@ -380,3 +380,21 @@ end
     @test B[argmax(B, dims=[2, 3])] == maximum(B, dims=[2, 3])
     @test B[argmin(B, dims=[2, 3])] == minimum(B, dims=[2, 3])
 end
+
+@testset "in-place reductions with mismatched dimensionalities" begin
+    B = reshape(1:24, 4, 3, 2)
+    for R in (fill(0, 4), fill(0, 4, 1), fill(0, 4, 1, 1))
+        @test @inferred(maximum!(R, B)) == reshape(21:24, size(R))
+        @test @inferred(minimum!(R, B)) == reshape(1:4, size(R))
+    end
+    for R in (fill(0, 1, 3), fill(0, 1, 3, 1))
+        @test @inferred(maximum!(R, B)) == reshape(16:4:24, size(R))
+        @test @inferred(minimum!(R, B)) == reshape(1:4:9, size(R))
+    end
+    @test_throws DimensionMismatch maximum!(fill(0, 4, 1, 1, 1), B)
+    @test_throws DimensionMismatch minimum!(fill(0, 4, 1, 1, 1), B)
+    @test_throws DimensionMismatch maximum!(fill(0, 1, 3, 1, 1), B)
+    @test_throws DimensionMismatch minimum!(fill(0, 1, 3, 1, 1), B)
+    @test_throws DimensionMismatch maximum!(fill(0, 1, 1, 2, 1), B)
+    @test_throws DimensionMismatch minimum!(fill(0, 1, 1, 2, 1), B)
+end