Add fix & tests for new mapreduce_scalar bypass

Project.toml

@@ -1,6 +1,6 @@
 name = "Strided"
 uuid = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
-version = "2.6.3"
+version = "2.6.4"
 authors = ["Lukas Devos <lukas.devos@ugent.be>", "Maarten Van Damme <maartenvd1994@gmail.com>", "Jutho Haegeman <jutho.haegeman@ugent.be>"]
 
 [deps]

ext/StridedGPUArraysExt.jl

@@ -105,6 +105,26 @@ function Strided._mapreduce(
     return Array(out)[1]
 end
 
+# 0-dimensional fast path: bypass @generated kernel
+# needs GPU specific extension to avoid scalar indexing error
+function Strided._mapreduce_scalar!(
+        @nospecialize(f), @nospecialize(op), @nospecialize(initop),
+        arrays::Tuple{GPUStridedView{TO, 0}, Vararg{GPUStridedView{<:Any, 0}}}
+    ) where {TO}
+    out = arrays[1]
+    iout = ParentIndex(Strided.offset(out) + 1)
+    @allowscalar begin
+        v = f(map(a -> a[ParentIndex(Strided.offset(a) + 1)], Base.tail(arrays))...)
+        if op === nothing
+            out[iout] = v
+        else
+            o = initop === nothing ? out[iout] : initop(out[iout])
+            out[iout] = op(o, v)
+        end
+    end
+    return nothing
+end
+
 function Strided._mapreduce_block!(
         f, op, initop,
         dims::Dims{N},

test/gpu.jl

@@ -184,3 +184,64 @@ end
         @test compare(x -> prod(exp, x), AT, A3)
     end
 end
+
+@testset "0-dimensional (scalar) StridedView ($AT)" for AT in ATs
+    @testset for T in (Float32, ComplexF32)
+        R = fill(rand(T)) # 0-dimensional Array
+        A = StridedView(AT(R))
+        @test ndims(A) == 0
+
+        # full reductions
+        @test sum(A) == sum(R)
+        @test prod(A) == prod(R)
+        @test mapreduce(abs2, +, A) == mapreduce(abs2, +, R)
+        @test maximum(abs, A) == maximum(abs, R)
+        @test minimum(abs, A) == minimum(abs, R)
+        @test sum(abs2, A) == sum(abs2, R)
+        @test mapreduce(identity, +, A; init = one(T)) ==
+            mapreduce(identity, +, R; init = one(T))
+
+        # map / map! / copy! / fill!
+        mapx = map(x -> 2x, A)
+        GPUArrays.@allowscalar begin
+            @test mapx[] == 2 * R[]
+        end
+        B = StridedView(AT(fill(zero(T))))
+        map!(x -> x + one(T), B, A)
+        GPUArrays.@allowscalar begin
+            @test B[] == collect(R)[] + one(T)
+        end
+        copy!(B, A)
+        GPUArrays.@allowscalar begin
+            @test B[] == R[]
+        end
+        fill!(B, one(T))
+        GPUArrays.@allowscalar begin
+            @test B[] == one(T)
+        end
+
+        # offset handling: 0-dim views into a larger parent
+        Psrc = AT(rand(T, 5))
+        Pdst = AT(rand(T, 5))
+        s = sreshape(StridedView(Psrc)[4:4], ())
+        d = sreshape(StridedView(Pdst)[3:3], ())
+        GPUArrays.@allowscalar begin
+            @test sum(s) == Psrc[4]
+        end
+        copy!(d, s)
+        GPUArrays.@allowscalar begin
+            @test Pdst[3] == Psrc[4]
+        end
+
+        # low-level in-place reduction with a custom initop
+        Pd = AT(rand(T, 5))
+        d2 = sreshape(StridedView(Pd)[2:2], ())
+        GPUArrays.@allowscalar begin
+            prev = Pd[2]
+        end
+        Strided._mapreducedim!(sin, +, identity, (), (d2, A))
+        GPUArrays.@allowscalar begin
+            @test Pd[2] == prev + sin(R[])
+        end
+    end
+end