Write direct dispatches for axpy! & axpby! (#225)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "ComponentArrays"
 uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 authors = ["Jonnie Diegelman <[email protected]>"]
-version = "0.15.3"
+version = "0.15.4"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/linear_algebra.jl

@@ -42,4 +42,18 @@ for op in [:*, :\, :/]
             end
         end
     end
-end
+end
+
+# Common Accumulation Operations
+## Needed for CUDA to work properly
+function LinearAlgebra.axpy!(α::Number, x::ComponentArray, y::ComponentArray)
+    getaxes(x) != getaxes(y) && throw(ArgumentError("Axes of `x` and `y` must match"))
+    axpy!(α, getdata(x), getdata(y))
+    return ComponentArray(y, getaxes(y))
+end
+
+function LinearAlgebra.axpby!(α::Number, x::ComponentArray, β::Number, y::ComponentArray)
+    getaxes(x) != getaxes(y) && throw(ArgumentError("Axes of `x` and `y` must match"))
+    axpby!(α, getdata(x), β, getdata(y))
+    return ComponentArray(y, getaxes(y))
+end

test/runtests.jl

@@ -123,7 +123,7 @@ end
     # Issue #116
     # Part 2: Arrays of arrays
     @test_throws Exception ComponentVector(a = [[3], [4, 5]], b = 1)
-    
+
     x = ComponentVector(a = [[3, 3], [4, 5]], b = 1)
     @test x.a[1] == [3, 3]
     @test x.b == 1
@@ -668,6 +668,28 @@ end
     @test ndims(dropdims(ones(1,1), dims=(1,2))) == 0
 end
 
+@testset "axpy! / axpby!" begin
+    y = ComponentArray(a = rand(4), b = rand(4))
+    x = ComponentArray(a = rand(4), b = rand(4))
+    ydata = copy(getdata(y))
+
+    axpy!(2, x, y)
+    @test getdata(y) == 2 .* getdata(x) .+ ydata
+
+    x = ComponentArray(a = rand(4), c = rand(4))
+    @test_throws ArgumentError axpy!(2, x, y)
+
+    y = ComponentArray(a = rand(4), b = rand(4))
+    x = ComponentArray(a = rand(4), b = rand(4))
+    ydata = copy(getdata(y))
+
+    axpby!(2, x, 3, y)
+    @test getdata(y) == 2 .* getdata(x) .+ 3 .* ydata
+
+    x = ComponentArray(a = rand(4), c = rand(4))
+    @test_throws ArgumentError axpby!(2, x, 3, y)
+end
+
 @testset "Autodiff" begin
     include("autodiff_tests.jl")
 end