Overload more GPU operations (#164)

* make adapt work on GPU

* overload more gpu operations

* update

Author: Shambles

src/compat/gpuarrays.jl

@@ -7,6 +7,23 @@ function GPUArrays.Adapt.adapt_structure(to, x::ComponentArray)
     return ComponentArray(data, getaxes(x))
 end
 
+GPUArrays.Adapt.adapt_storage(::Type{ComponentArray{T,N,A,Ax}}, xs::AT) where {T,N,A,Ax,AT<:AbstractArray} =
+    GPUArrays.Adapt.adapt_storage(A, xs)
+
+function Base.fill!(A::GPUComponentArray{T}, x) where {T}
+    length(A) == 0 && return A
+    GPUArrays.gpu_call(A, convert(T, x)) do ctx, a, val
+        idx = GPUArrays.@linearidx(a)
+        @inbounds a[idx] = val
+        return
+    end
+    A
+end
+
+LinearAlgebra.dot(x::GPUComponentArray, y::GPUComponentArray) = dot(getdata(x), getdata(y))
+LinearAlgebra.norm(ca::GPUComponentArray, p::Real) = norm(getdata(ca), p)
+LinearAlgebra.rmul!(ca::GPUComponentArray, b::Number) = GPUArrays.generic_rmul!(ca, b)
+
 function Base.map(f, x::GPUComponentArray, args...)
     data = map(f, getdata(x), getdata.(args)...)
     return ComponentArray(data, getaxes(x))
@@ -46,4 +63,4 @@ for (fname, op) in [(:sum, :(Base.add_sum)), (:prod, :(Base.mul_prod)),
         Base.$(fname!)(f::Function, r::GPUComponentArray, A::GPUComponentArray{T}) where T =
             GPUArrays.mapreducedim!(f, $(op), getdata(r), getdata(A); init=neutral_element($(op), T))
     end
-end
+end

test/gpu_tests.jl

@@ -7,18 +7,42 @@ jlca = ComponentArray(jla, Axis(a=1:2, b=3:4))
 
 @testset "Broadcasting" begin
     @test identity.(jlca + jla) ./ 2 == jlca
-    
+
     @test getdata(map(identity, jlca)) isa JLArray
     @test all(==(0), map(-, jlca, jla))
     @test all(map(-, jlca, jlca) .== 0)
     @test all(==(0), map(-, jla, jlca))
 
     @test any(==(1), jlca)
     @test count(>(2), jlca) == 2
-    
+
     # Make sure mapreducing multiple arrays works
     @test mapreduce(==, +, jlca, jla) == 4
     @test mapreduce(abs2, +, jlca) == 30
 
     @test all(map(sin, jlca) .== sin.(jlca) .== sin.(jla) .≈ sin.(1:4))
 end
+
+@testset "adapt" begin
+    x = [1 2; 3 4]
+    jlx = JLArrays.Adapt.adapt(typeof(jlca), x)
+    @test jlx isa JLArray
+end
+
+@testset "Linear Algebra" begin
+    @testset "fill!" begin
+        jlca2 = deepcopy(jlca)
+        jlca2 = fill!(jlca2, 2)
+        @test jlca2 == ComponentArray(jl([2,2,2,2]), Axis(a=1:2, b=3:4))
+    end
+
+    @testset "norm" begin
+        @test norm(jlca, 2) == norm(jla,2)
+        @test norm(jlca, Inf) == norm(jla,Inf)
+    end
+
+    @testset "rmul!" begin
+        jlca3 = deepcopy(jlca)
+        @test rmul!(jlca3, 2) == ComponentArray(jla .* 2, Axis(a=1:2, b=3:4))
+    end
+end