preliminary support for multi-arg map and mapreduce with breaking changes

Project.toml

@@ -1,7 +1,7 @@
 name = "AcceleratedKernels"
 uuid = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
 authors = ["Andrei-Leonard Nicusan <leonard@evophase.co.uk> and contributors"]
-version = "0.4.3"
+version = "0.5.0"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

src/arithmetics.jl

@@ -47,8 +47,8 @@ function sum(
     kwargs...
 )
     reduce(
-        +, src, backend;
-        init,
+        +, src;
+        backend, init,
         kwargs...
     )
 end
@@ -103,8 +103,8 @@ function prod(
     kwargs...
 )
     reduce(
-        *, src, backend;
-        init,
+        *, src;
+        backend, init,
         kwargs...
     )
 end
@@ -159,8 +159,8 @@ function maximum(
     kwargs...
 )
     reduce(
-        max, src, backend;
-        init,
+        max, src;
+        backend, init,
         kwargs...
     )
 end
@@ -215,8 +215,8 @@ function minimum(
     kwargs...
 )
     reduce(
-        min, src, backend;
-        init,
+        min, src;
+        backend, init,
         kwargs...
     )
 end
@@ -277,8 +277,8 @@ function count(
     kwargs...
 )
     mapreduce(
-        x -> x ? one(typeof(init)) : zero(typeof(init)), +, src, backend;
-        init,
+        x -> x ? one(typeof(init)) : zero(typeof(init)), +, src;
+        backend, init,
         neutral=zero(typeof(init)),
         kwargs...
     )
@@ -291,8 +291,8 @@ function count(
     kwargs...
 )
     mapreduce(
-        x -> f(x) ? one(typeof(init)) : zero(typeof(init)), +, src, backend;
-        init,
+        x -> f(x) ? one(typeof(init)) : zero(typeof(init)), +, src;
+        backend, init,
         neutral=zero(typeof(init)),
         kwargs...
     )

src/map.jl

@@ -1,6 +1,7 @@
 """
     map!(
-        f, dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(src);
+        f, dst::AbstractArray, src::AbstractArray...;
+        backend::Backend=get_backend(src);
 
         # CPU settings
         max_tasks=Threads.nthreads(),
@@ -32,23 +33,25 @@ end
 ```
 """
 function map!(
-    f, dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(src);
+    f, dst::AbstractArray, src::AbstractArray...;
+    backend::Backend=get_backend(src[1]),
     kwargs...
 )
-    @argcheck length(dst) == length(src)
+    @argcheck lengthcheck(dst, src...)
     foreachindex(
-        src, backend;
+        dst, backend;
         kwargs...
     ) do idx
-        dst[idx] = f(src[idx])
+        dst[idx] = f(indextuple(src, idx)...)
     end
     dst
 end
 
 
 """
     map(
-        f, src::AbstractArray, backend::Backend=get_backend(src);
+        f, src::AbstractArray;
+        backend::Backend=get_backend(src),
 
         # CPU settings
         max_tasks=Threads.nthreads(),
@@ -63,12 +66,13 @@ changes the `eltype`, allocate `dst` separately and call [`map!`](@ref)). The CP
 settings are the same as for [`foreachindex`](@ref).
 """
 function map(
-    f, src::AbstractArray, backend::Backend=get_backend(src);
+    f, src::AbstractArray...;
+    backend::Backend=get_backend(src[1]),
     kwargs...
 )
-    dst = similar(src)
+    dst = similar(src[1])
     map!(
-        f, dst, src, backend;
-        kwargs...
+        f, dst, src...;
+        backend, kwargs...
     )
 end

src/reduce/mapreduce_1d_cpu.jl

@@ -1,5 +1,6 @@
 function mapreduce_1d_cpu(
-    f, op, src::AbstractArray, backend::Backend;
+    f, op, src::AbstractArray;
+    backend::Backend,
     init,
     neutral,
 

src/reduce/mapreduce_1d_gpu.jl

@@ -47,7 +47,8 @@ end
 
 
 function mapreduce_1d_gpu(
-    f, op, src::AbstractArray, backend::Backend;
+    f, op, src::AbstractArray;
+    backend::Backend,
     init,
     neutral,
 

src/reduce/mapreduce_nd.jl

@@ -1,5 +1,6 @@
 function mapreduce_nd(
-    f, op, src::AbstractArray, backend::Backend;
+    f, op, src::AbstractArray;
+    backend::Backend,
     init,
     neutral=neutral_element(op, eltype(src)),
     dims::Int,

src/reduce/reduce.jl

@@ -7,7 +7,8 @@ include("mapreduce_nd.jl")
 
 """
     reduce(
-        op, src::AbstractArray, backend::Backend=get_backend(src);
+        op, src::AbstractArray;
+        backend::Backend=get_backend(src),
         init,
         neutral=neutral_element(op, eltype(src)),
         dims::Union{Nothing, Int}=nothing,
@@ -74,13 +75,14 @@ mcolsum = AK.reduce(+, m; init=zero(eltype(m)), dims=2)
 ```
 """
 function reduce(
-    op, src::AbstractArray, backend::Backend=get_backend(src);
+    op, src::AbstractArray;
+    backend::Backend=get_backend(src),
     init,
     kwargs...
 )
     _mapreduce_impl(
-        identity, op, src, backend;
-        init,
+        identity, op, src;
+        backend, init,
         kwargs...
     )
 end
@@ -90,7 +92,8 @@ end
 
 """
     mapreduce(
-        f, op, src::AbstractArray, backend::Backend=get_backend(src);
+        f, op, src::AbstractArray;
+        backend::Backend=get_backend(src),
         init,
         neutral=neutral_element(op, eltype(src)),
         dims::Union{Nothing, Int}=nothing,
@@ -154,20 +157,22 @@ mcolsumsq = AK.mapreduce(f, +, m; init=zero(eltype(m)), dims=2)
 ```
 """
 function mapreduce(
-    f, op, src::AbstractArray, backend::Backend=get_backend(src);
+    f, op, src::AbstractArray;
+    backend::Backend=get_backend(src),
     init,
     kwargs...
 )
     _mapreduce_impl(
-        f, op, src, backend;
-        init,
+        f, op, src;
+        backend, init,
         kwargs...
     )
 end
 
 
 function _mapreduce_impl(
-    f, op, src::AbstractArray, backend::Backend;
+    f, op, src::AbstractArray;
+    backend::Backend,
     init,
     neutral=neutral_element(op, eltype(src)),
     dims::Union{Nothing, Int}=nothing,
@@ -185,25 +190,25 @@ function _mapreduce_impl(
     if isnothing(dims)
         if use_gpu_algorithm(backend, prefer_threads)
             mapreduce_1d_gpu(
-                f, op, src, backend;
-                init, neutral,
+                f, op, src;
+                backend, init, neutral,
                 max_tasks, min_elems,
                 block_size, temp,
                 switch_below
             )
         else
             mapreduce_1d_cpu(
-                f, op, src, backend;
-                init, neutral,
+                f, op, src;
+                backend, init, neutral,
                 max_tasks, min_elems,
                 block_size, temp,
                 switch_below
             )
         end
     else
         return mapreduce_nd(
-            f, op, src, backend;
-            init, neutral, dims,
+            f, op, src; 
+            backend, init, neutral, dims,
             max_tasks, prefer_threads,
             min_elems, block_size,
             temp,

src/utils.jl

@@ -126,3 +126,13 @@ module DocHelpers
         return Markdown.parse(join(lines[captured_istart:captured_iend], '\n'))
     end
 end
+
+@inline indextuple(arrays::Tuple, idx)::Tuple = Base.map(a -> a[idx], arrays)
+
+@inline function lengthcheck(dst::AbstractArray, src::AbstractArray...)
+    n = length(dst)
+    for a ∈ src
+        length(a) == n || return false
+    end
+    return true
+end

test/map.jl

@@ -16,12 +16,23 @@
         end
         @test y == map(i -> i^2, x)
 
+        z = ones(Int, 1000)
+        AK.map!(y, x, z; prefer_threads) do a, b
+            a^2 + b
+        end
+        @test y == map((a, b) -> a^2 + b, x, z)
+
         x = rand(Float32, 1000)
         y = AK.map(x; prefer_threads, max_tasks=2, min_elems=100) do i
             i > 0.5 ? i : 0
         end
         @test y == map(i -> i > 0.5 ? i : 0, x)
 
+        y = AK.map(x, z; prefer_threads, max_tasks=2, min_elems=100) do a, b
+            a > 0.5 ? a+b : -b
+        end
+        @test y == map((a, b) -> a > 0.5 ? a+b : -b, x, z)
+
         x = rand(Float32, 1000)
         y = AK.map(x; prefer_threads, max_tasks=4, min_elems=500) do i
             i > 0.5 ? i : 0
@@ -45,12 +56,23 @@
         end
         @test Array(y) == map(i -> i^2, 1:1000)
 
+        z = array_from_host(ones(Int, 1000))
+        AK.map!(y, x, z; prefer_threads) do a, b
+            a^2 + b
+        end
+        @test Array(y) == map((a, b) -> a^2 + b, x, z)
+
         x = array_from_host(rand(Float32, 1000))
         y = AK.map(x; prefer_threads, block_size=64) do i
             i > 0.5 ? i : 0
         end
         @test Array(y) == map(i -> i > 0.5 ? i : 0, Array(x))
 
+        y = AK.map(x, z; prefer_threads, block_size=64) do a, b
+            a > 0.5 ? a+b : -b
+        end
+        @test Array(y) == map((a, b) -> a > 0.5 ? a+b : -b, x, z)
+
         # Test that undefined kwargs are not accepted
         @test_throws MethodError AK.map(x -> x^2, x; prefer_threads, bad=:kwarg)
     end

test/reduce.jl

@@ -115,7 +115,7 @@ Base.zero(::Type{Point}) = Point(0.0f0, 0.0f0)
     for _ in 1:100
         num_elems = rand(1:1000)
         v = 1:num_elems
-        s = AK.reduce(+, v, BACKEND; prefer_threads, init=Int32(0))
+        s = AK.reduce(+, v; backend=BACKEND, prefer_threads, init=Int32(0))
         vh = Array(v)
         @test s == reduce(+, vh)
     end
@@ -337,7 +337,7 @@ end
     for _ in 1:100
         num_elems = rand(1:1000)
         v = 1:num_elems
-        s = AK.mapreduce(abs, +, v, BACKEND; prefer_threads, init=Int32(0))
+        s = AK.mapreduce(abs, +, v; backend=BACKEND, prefer_threads, init=Int32(0))
         vh = Array(v)
         @test s == mapreduce(abs, +, vh)
     end