Mutable sum of array to array (#278)

* Mutable sum of array to array

* Add tests

* Fix format

Author: blegat

src/implementations/LinearAlgebra.jl

@@ -94,6 +94,25 @@ function operate!(op::Union{typeof(+),typeof(-)}, A::Array, B::AbstractArray)
     return broadcast!(op, A, B)
 end
 
+function operate_to!(
+    output::Array,
+    op::Union{typeof(+),typeof(-)},
+    A::AbstractArray,
+    B::AbstractArray,
+)
+    if axes(output) != promote_shape(A, B)
+        throw(
+            DimensionMismatch(
+                "Cannot sum or substract matrices of axes `$(axes(A))` and" *
+                " `$(axes(B))` into a matrix of axes `$(axes(output))`," *
+                " expected axes `$(promote_shape(A, B))`.",
+            ),
+        )
+    end
+    # We don't have `MA.broadcast_to!` as it would be exactly `Base.broadcast!`.
+    return Base.broadcast!(op, output, A, B)
+end
+
 # We call `scaling_to_number` as `UniformScaling` do not support broadcasting
 function operate!(
     op::AddSubMul,

test/matmul.jl

@@ -119,6 +119,13 @@ end
             "Cannot sum or substract a matrix of axes `$(axes(B))` into matrix of axes `$(axes(A))`, expected axes `$(axes(B))`.",
         )
         @test_throws err MA.operate!(+, A, B)
+        output = zeros(2)
+        A = zeros(2, 1)
+        B = zeros(2, 1)
+        err = DimensionMismatch(
+            "Cannot sum or substract matrices of axes `$(axes(A))` and `$(axes(B))` into a matrix of axes `$(axes(output))`, expected axes `$(axes(B))`.",
+        )
+        @test_throws err MA.operate_to!(output, +, A, B)
     end
     @testset "unsupported_product" begin
         unsupported_product()
@@ -435,3 +442,18 @@ end
     @test B == [1 2]
     @test D == B * A
 end
+
+function test_array_sum(::Type{T}) where {T}
+    x = zeros(T, 2)
+    y = copy(x)
+    z = copy(y)
+    alloc_test(() -> MA.operate!(+, y, z), 0)
+    alloc_test(() -> MA.add!!(y, z), 0)
+    alloc_test(() -> MA.operate_to!(x, +, y, z), 0)
+    alloc_test(() -> MA.add_to!!(x, y, z), 0)
+    return
+end
+
+@testset "Array sum" begin
+    test_array_sum(Int)
+end