diff --git a/Project.toml b/Project.toml index a0f9ab3..c483aaa 100644 --- a/Project.toml +++ b/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 ", "Maarten Van Damme ", "Jutho Haegeman "] [deps] diff --git a/ext/StridedGPUArraysExt.jl b/ext/StridedGPUArraysExt.jl index b65313d..7525b44 100644 --- a/ext/StridedGPUArraysExt.jl +++ b/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}, diff --git a/test/gpu.jl b/test/gpu.jl index 596ca00..cf08ee5 100644 --- a/test/gpu.jl +++ b/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