Improve GPU functionality

Project.toml

@@ -25,10 +25,12 @@ Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [weakdeps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [extensions]
 DimensionalDataCategoricalArraysExt = "CategoricalArrays"
+DimensionalDataGPUArraysCoreExt = "GPUArraysCore"
 DimensionalDataMakie = "Makie"
 
 [compat]
@@ -54,6 +56,7 @@ Interfaces = "0.3"
 IntervalSets = "0.5, 0.6, 0.7"
 InvertedIndices = "1"
 IteratorInterfaceExtensions = "1"
+JLArrays = "0.1"
 LinearAlgebra = "1"
 Makie = "0.19, 0.20, 0.21"
 OffsetArrays = "1"
@@ -85,6 +88,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ImageFiltering = "6a3955dd-da59-5b1f-98d4-e7296123deb5"
 ImageTransformations = "02fcd773-0e25-5acc-982a-7f6622650795"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
@@ -95,4 +99,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["Aqua", "ArrayInterface", "BenchmarkTools", "CategoricalArrays", "ColorTypes", "Combinatorics", "CoordinateTransformations", "DataFrames", "Distributions", "Documenter", "ImageFiltering", "ImageTransformations", "CairoMakie", "OffsetArrays", "Plots", "Random", "SafeTestsets", "StatsPlots", "Test", "Unitful"]
+test = ["Aqua", "ArrayInterface", "BenchmarkTools", "CategoricalArrays", "ColorTypes", "Combinatorics", "CoordinateTransformations", "DataFrames", "Distributions", "Documenter", "ImageFiltering", "ImageTransformations", "JLArrays", "CairoMakie", "OffsetArrays", "Plots", "Random", "SafeTestsets", "StatsPlots", "Test", "Unitful"]

ext/DimensionalDataGPUArraysCoreExt.jl

@@ -0,0 +1,12 @@
+module DimensionalDataGPUArraysCoreExt
+
+using DimensionalData: AbstractDimArray
+using GPUArraysCore: AbstractGPUArrayStyle
+using Base.Broadcast: Broadcasted
+
+function Base.copyto!(des::AbstractDimArray, bc::Broadcasted{<:AbstractGPUArrayStyle})
+    copyto!(parent(des), bc)
+    return des
+end
+
+end

src/array/broadcast.jl

@@ -67,6 +67,20 @@ function Base.copyto!(dest::AbstractDimArray, bc::Broadcasted{DimensionalStyle{S
     end
 end
 
+# This is needed for GPUs to prevent scalar indexing problems for things like
+# d .= 1:10
+function Base.copyto!(dest::AbstractDimArray, bc::Broadcasted{Nothing})
+    copyto!(parent(dest), bc)
+    dest
+end
+
+# Needed for things like d .= 0 when on the GPU
+function Base.copyto!(dest::AbstractDimArray, bc::Broadcasted{<:Broadcast.AbstractArrayStyle{0}})
+    copyto!(parent(dest), bc)
+    dest
+end
+
+
 function Base.similar(bc::Broadcast.Broadcasted{DimensionalStyle{S}}, ::Type{T}) where {S,T}
     A = _firstdimarray(bc)
     rebuildsliced(A, similar(_unwrap_broadcasted(bc), T, axes(bc)...), axes(bc), Symbol(""))

src/array/methods.jl

@@ -53,26 +53,33 @@ for (m, f) in ((:Statistics, :median), (:Base, :any), (:Base, :all))
     end
 end
 
-# These are not exported but it makes a lot of things easier using them
-function Base._mapreduce_dim(f, op, nt::NamedTuple{(),<:Tuple}, A::AbstractDimArray, dims)
-    rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, _astuple(dims))), reducedims(A, dims))
-end
-function Base._mapreduce_dim(f, op, nt::NamedTuple{(),<:Tuple}, A::AbstractDimArray, dims::Colon)
-    Base._mapreduce_dim(f, op, nt, parent(A), dims)
-end
-function Base._mapreduce_dim(f, op, nt, A::AbstractDimArray, dims)
-    rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
-end
-function Base._mapreduce_dim(f, op, nt, A::AbstractDimArray, dims::Colon)
-    rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
+function Base.mapreduce(f, op, A::AbstractDimArray; dims=Base.Colon(), kwargs...)
+    dims === Colon() && return mapreduce(f, op, parent(A); kwargs...)
+    rebuild(A, mapreduce(f, op, parent(A); dims=dimnum(A, dims), kwargs...), reducedims(A, dims))
 end
 
-function Base._mapreduce_dim(f, op, nt::Base._InitialValue, A::AbstractDimArray, dims)
-    rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
-end
-function Base._mapreduce_dim(f, op, nt::Base._InitialValue, A::AbstractDimArray, dims::Colon)
-    Base._mapreduce_dim(f, op, nt, parent(A), dims)
-end
+# These methods prevent mapreduce from working on the GPU so we will directly overload it.
+# # These are not exported but it makes a lot of things easier using them
+# function Base._mapreduce_dim(f, op, nt::NamedTuple{(),<:Tuple}, A::AbstractDimArray, dims)
+#     rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, _astuple(dims))), reducedims(A, dims))
+# end
+# function Base._mapreduce_dim(f, op, nt::NamedTuple{(),<:Tuple}, A::AbstractDimArray, dims::Colon)
+#     Base._mapreduce_dim(f, op, nt, parent(A), dims)
+# end
+# function Base._mapreduce_dim(f, op, nt, A::AbstractDimArray, dims)
+#     rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
+# end
+# function Base._mapreduce_dim(f, op, nt, A::AbstractDimArray, dims::Colon)
+#     rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
+# end
+
+# function Base._mapreduce_dim(f, op, nt::Base._InitialValue, A::AbstractDimArray, dims)
+#     rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
+# end
+# function Base._mapreduce_dim(f, op, nt::Base._InitialValue, A::AbstractDimArray, dims::Colon)
+#     Base._mapreduce_dim(f, op, nt, parent(A), dims)
+# end
+
 
 
 # TODO: Unfortunately Base/accumulate.jl kw methods all force dims to be Integer.

test/broadcast.jl

@@ -1,5 +1,5 @@
 using DimensionalData, Test
-
+using JLArrays
 using DimensionalData: NoLookup
 
 # Tests taken from NamedDims. Thanks @oxinabox
@@ -168,6 +168,19 @@ end
     @test A[DimSelectors(sub)] == C[DimSelectors(sub)]
 end
 
+@testset "GPUArray broadcast" begin
+    arr = JLArray(rand(64, 64))
+    A = DimArray(arr, (X(1.0:1.0:64), Y(-32.0:1.0:31)))
+    @test arr.^2 ≈ parent(A.^2)
+    x = 1.0:1.0:64
+    A .= x.^2 .+ x'
+    @test parent(A) ≈ x.^2 .+ x'
+
+    A .= 1.0
+    # all gives scalar indexing so we use mapreduce
+    @test mapreduce(==(1.0), *, A)
+end
+
 # @testset "Competing Wrappers" begin
 #     da = DimArray(ones(4), X)
 #     ta = TrackedArray(5 * ones(4))

test/methods.jl

@@ -1,6 +1,6 @@
 using DimensionalData, Statistics, Test, Unitful, SparseArrays, Dates
 using DimensionalData.Lookups, DimensionalData.Dimensions
-
+using JLArrays
 using LinearAlgebra: Transpose
 
 xs = (1, X, X(), :X)
@@ -709,3 +709,81 @@ end
         end
     end
 end
+
+
+@testset "mapreduce" begin
+    @testset "Array 2D" begin
+        y = Y(['a', 'b', 'c'])
+        ti = Ti(DateTime(2021, 1):Month(1):DateTime(2021, 4))
+        ys = (1, Y, Y(), :Y, y)
+        tis = (2, Ti, Ti(), :Ti, ti)
+        data = [-87  -49  107  -18
+                24   44  -62  124
+                122  -11   48   -7]
+        A = DimArray(data, (y, ti))
+
+
+        @test mapreduce(identity, +, A) ≈ mapreduce(identity, +, parent(A))
+        @test mapreduce(x->x^3+5, +, A) ≈ mapreduce(x->x^3+5, +, parent(A))
+
+        for dims in ys
+            @test mapreduce(identity, +, A; dims) ≈ mapreduce(identity, +, parent(A); dims=1)
+        end
+
+        for dims in tis
+            @test mapreduce(identity, +, A; dims) ≈ mapreduce(identity, +, parent(A); dims=2)
+        end
+
+        @test mapreduce(identity, +, A; dims=Y) ≈ mapreduce(identity, +, parent(A); dims=1)
+        @test mapreduce(identity, +, A; dims=Ti) ≈ mapreduce(identity, +, parent(A); dims=2)
+        @test mapreduce(identity, +, A; dims=(Y, Ti)) ≈ mapreduce(identity, +, parent(A); dims=(1, 2))
+
+        init = 5.0
+        @test mapreduce(identity, +, A; init) ≈ mapreduce(identity, +, parent(A); init)
+        @test mapreduce(x->x^3+5, +, A; init) ≈ mapreduce(x->x^3+5, +, parent(A); init)
+        @test mapreduce(identity, +, A; dims=Y, init) ≈ mapreduce(identity, +, parent(A); dims=1, init)
+        @test mapreduce(identity, +, A; dims=Ti, init) ≈ mapreduce(identity, +, parent(A); dims=2, init)
+        @test mapreduce(identity, +, A; dims=(Y, Ti), init) ≈ mapreduce(identity, +, parent(A); dims=(1, 2), init)
+    end
+    @testset "Vector" begin
+        x = DimArray([56, -123, -60, -44, -64, 70, 52, -48, -74, 86], X(2:2:20))
+        @test mapreduce(x->x^2, +, x) ≈ mapreduce(x->x^2, +, parent(x))
+        @test mapreduce(identity, +, x) ≈ mapreduce(identity, +, parent(x))
+        @test mapreduce(identity, +, x; dims=X) ≈ mapreduce(identity, +, parent(x); dims=1)
+        @test mapreduce(x->x^2, +, x; dims=X) ≈ mapreduce(x->x^2, +, parent(x); dims=1)
+        @test mapreduce(identity, +, x; init=5.0) ≈ mapreduce(identity, +, parent(x); init=5.0)
+    end
+
+    @testset "JLArray" begin
+        y = Y(['a', 'b', 'c'])
+        ti = Ti(DateTime(2021, 1):Month(1):DateTime(2021, 4))
+        ys = (1, Y, Y(), :Y, y)
+        tis = (2, Ti, Ti(), :Ti, ti)
+        data = JLArray([-87  -49  107  -18
+                24   44  -62  124
+                122  -11   48   -7])
+        A = DimArray(data, (y, ti))
+
+        @test mapreduce(identity, +, A) ≈ mapreduce(identity, +, parent(A))
+        @test mapreduce(x->x^3+5, +, A) ≈ mapreduce(x->x^3+5, +, parent(A))
+        # Using parent since JLArray errors
+        for dims in ys
+            @test parent(mapreduce(identity, +, A; dims)) ≈ mapreduce(identity, +, parent(A); dims=1)
+        end
+
+        for dims in tis
+            @test parent(mapreduce(identity, +, A; dims)) ≈ mapreduce(identity, +, parent(A); dims=2)
+        end
+
+        @test parent(mapreduce(identity, +, A; dims=Y)) ≈ mapreduce(identity, +, parent(A); dims=1)
+        @test parent(mapreduce(identity, +, A; dims=Ti)) ≈ mapreduce(identity, +, parent(A); dims=2)
+        @test parent(mapreduce(identity, +, A; dims=(Y, Ti))) ≈ mapreduce(identity, +, parent(A); dims=(1, 2))
+
+        init = 5.0
+        @test mapreduce(identity, +, A; init) ≈ mapreduce(identity, +, parent(A); init)
+        @test mapreduce(x->x^3+5, +, A; init) ≈ mapreduce(x->x^3+5, +, parent(A); init)
+        @test parent(mapreduce(identity, +, A; dims=Y, init)) ≈ mapreduce(identity, +, parent(A); dims=1, init)
+        @test parent(mapreduce(identity, +, A; dims=Ti, init)) ≈ mapreduce(identity, +, parent(A); dims=2, init)
+        @test parent(mapreduce(identity, +, A; dims=(Y, Ti), init)) ≈ mapreduce(identity, +, parent(A); dims=(1, 2), init)
+    end
+end