Fix linearly indexed array math by reshaping arrays (#60164)

This was broken in #59961, as
`map` deals with trailing singleton axes differently from broadcasting:
```julia
julia> map(+, ones(1), ones(1,1)) |> size
(1,)

julia> broadcast(+, ones(1), ones(1,1)) |> size
(1, 1)
```
This PR limits the new method to the case where the ndims match, in
which case there are no trailing axes and the two are equivalent. The
alternate approach suggested in
#59961 (comment) is
to reshape the arrays, but this adds overhead that nullifies the
performance improvement for small arrays.

Author: jishnub

base/arraymath.jl

@@ -9,8 +9,19 @@ end
 # Using map over broadcast enables vectorization for wide matrices with few rows.
 # This is because we use linear indexing in `map` as opposed to Cartesian indexing in broadcasting.
 # https://github.com/JuliaLang/julia/issues/47873#issuecomment-1352472461
-function _broadcast_preserving_zero_d(f, A::Array, B::Array)
-    map(f, A, B)
+function _broadcast_preserving_zero_d(f, A::Array{<:Any,N}, B::Array{<:Any,N}, Cs::Array{<:Any,N}...) where {N}
+    map(f, A, B, Cs...)
+end
+
+function _broadcast_preserving_zero_d(f, A::Array, B::Array, Cs::Array...)
+    # we already know that the shapes are compatible.
+    # We just need to select the size corresponding to the higest ndims
+    # and reshape all the arrays to that size
+    arrays = (A, B, Cs...)
+    sz = mapreduce(size, (x,y) -> length(x) > length(y) ? x : y, arrays)
+    # Skip reshaping where possible to avoid the overhead
+    arrays_sameshape = map(x -> length(sz) == ndims(x) ? x : reshape(x, sz), arrays)
+    map(f, arrays_sameshape...)
 end
 
 function _broadcast_preserving_zero_d(f, A::Array, B::Number)
@@ -28,11 +39,12 @@ for f in (:+, :-)
     end
 end
 
-function +(A::Array, Bs::Array...)
-    for B in Bs
-        promote_shape(A, B) # check size compatibility
+function +(A::Array, B::Array, Cs::Array...)
+    promote_shape(A, B)
+    for C in Cs
+        promote_shape(A, C) # check size compatibility
     end
-    map(+, A, Bs...)
+    _broadcast_preserving_zero_d(+, A, B, Cs...)
 end
 
 for f in (:/, :\, :*)

test/abstractarray.jl

@@ -2376,3 +2376,16 @@ end
     show(io, m2)
     @test String(take!(io)) == "Any[#= circular reference @-1 =# 3; 2 4;;; 5 7; 6 8]"
 end
+
+@testset "size promotion in addition/subtraction" begin
+    for A in Any[ones(), ones(1), ones(1,1,1)]
+        @test +(A) == A
+        for B in Any[ones(), ones(1), ones(1,1,1)]
+            sz = ndims(A) > ndims(B) ? size(A) : size(B)
+            @test A + B == fill(2.0,sz)
+            @test A - B == zeros(sz)
+            @test A + B + zeros() == A + B
+            @test A - B - zeros() == A - B
+        end
+    end
+end