Reduce allocations for in-place reductions and fix mixed-dimensionali…

…ty edge-cases

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

base/reducedim.jl

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

test/offsetarray.jl

@@ -489,3 +489,22 @@ 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 @inferred(minimum!(R, B)) == reshape(minimum(B, dims=(2,3)), axes(R)) == reshape(1:4, axes(R))
+    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 @inferred(minimum!(R, B)) == reshape(minimum(B, dims=(1,3)), axes(R)) == reshape(1:4:9, axes(R))
+    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