Make GPU ComponentArray boradcasting work. Fixes #155

Author: jonniedie

src/ComponentArrays.jl

@@ -59,6 +59,7 @@ function __init__()
     @require RecursiveArrayTools="731186ca-8d62-57ce-b412-fbd966d074cd" required("recursivearraytools.jl")
     @require StaticArrays="90137ffa-7385-5640-81b9-e52037218182" required("staticarrays.jl")
     @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" required("reversediff.jl")
+    @require GPUArrays="0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" required("gpuarrays.jl")
 end
 
 end

src/compat/gpuarrays.jl

@@ -0,0 +1,49 @@
+const GPUComponentArray = ComponentArray{T,N,<:GPUArrays.AbstractGPUArray,Ax} where {T,N,Ax}
+
+GPUArrays.backend(x::ComponentArray) = GPUArrays.backend(getdata(x))
+
+function GPUArrays.Adapt.adapt_structure(to, x::ComponentArray)
+    data = GPUArrays.Adapt.adapt_structure(to, getdata(x))
+    return ComponentArray(data, getaxes(x))
+end
+
+function Base.map(f, x::GPUComponentArray, args...)
+    data = map(f, getdata(x), getdata.(args)...)
+    return ComponentArray(data, getaxes(x))
+end
+function Base.map(f, x::GPUComponentArray, args::Vararg{Union{Base.AbstractBroadcasted, AbstractArray}})
+    data = map(f, getdata(x), map(getdata, args)...)
+    return ComponentArray(data, getaxes(x))
+end
+
+# We need all of these to avoid method ambiguities
+function Base.mapreduce(f, op, x::GPUComponentArray; kwargs...)
+    return mapreduce(f, op, getdata(x); kwargs...)
+end
+function Base.mapreduce(f, op, x::GPUComponentArray, args...; kwargs...)
+    return mapreduce(f, op, getdata(x), map(getdata, args)...; kwargs...)
+end
+function Base.mapreduce(f, op, x::GPUComponentArray, args::Vararg{Union{Base.AbstractBroadcasted, AbstractArray}}; kwargs...)
+    return mapreduce(f, op, getdata(x), map(getdata, args)...; kwargs...)
+end
+
+# These are all stolen from GPUArrays.j;
+Base.any(A::GPUComponentArray{Bool}) = mapreduce(identity, |, getdata(A))
+Base.all(A::GPUComponentArray{Bool}) = mapreduce(identity, &, getdata(A))
+
+Base.any(f::Function, A::GPUComponentArray) = mapreduce(f, |, getdata(A))
+Base.all(f::Function, A::GPUComponentArray) = mapreduce(f, &, getdata(A))
+
+Base.count(pred::Function, A::GPUComponentArray; dims=:, init=0) =
+    mapreduce(pred, Base.add_sum, getdata(A); init=init, dims=dims)
+
+# avoid calling into `initarray!`
+for (fname, op) in [(:sum, :(Base.add_sum)), (:prod, :(Base.mul_prod)),
+                    (:maximum, :(Base.max)), (:minimum, :(Base.min)),
+                    (:all, :&),              (:any, :|)]
+    fname! = Symbol(fname, '!')
+    @eval begin
+        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

test/Project.toml

@@ -4,6 +4,7 @@ ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 InvertedIndices = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"

test/gpu_tests.jl

@@ -0,0 +1,24 @@
+using JLArrays
+
+JLArrays.allowscalar(false)
+
+jla = jl(collect(1:4))
+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

test/runtests.jl

@@ -608,4 +608,8 @@ end
 
 @testset "Autodiff" begin
     include("autodiff_tests.jl")
+end
+
+@testset "GPU" begin
+    include("gpu_tests.jl")
 end