diff --git a/ext/StridedGPUArraysExt.jl b/ext/StridedGPUArraysExt.jl index 7525b44..d34465d 100644 --- a/ext/StridedGPUArraysExt.jl +++ b/ext/StridedGPUArraysExt.jl @@ -125,6 +125,17 @@ function Strided._mapreduce_scalar!( return nothing end +# Singleton views with axes still execute on the device, without layout planning. +function Strided._mapreduce_scalar!( + f, op, initop, + arrays::Tuple{GPUStridedView{TO, N}, Vararg{GPUStridedView{<:Any, N}}} + ) where {TO, N} + return Strided._mapreduce_block!( + f, op, initop, size(arrays[1]), map(strides, arrays), + map(Strided.offset, arrays), nothing, arrays + ) +end + function Strided._mapreduce_block!( f, op, initop, dims::Dims{N}, diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 48b3e66..0f10caf 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -173,14 +173,15 @@ _computecosts(strides) = map(a -> ifelse(iszero(a), 1, a << 1), map(min, strides function _mapreduce_order!( f::F, op::OP, initop::I, dims::Dims, arrays::Tuple{Vararg{StridedView}} ) where {F, OP, I} - isempty(dims) && return _mapreduce_scalar!(f, op, initop, arrays) + # A single element needs no layout planning, regardless of the number of axes. + all(isone, dims) && return _mapreduce_scalar!(f, op, initop, arrays) dims, allstrides = order_and_fuse_dims(dims, map(strides, arrays)) offsets = map(offset, arrays) costs = _computecosts(allstrides) return _mapreduce_block!(f, op, initop, dims, allstrides, offsets, costs, arrays) end -# 0-dimensional fast path: bypass @generated kernel +# Single-element fast path: bypass @generated kernel function _mapreduce_scalar!(f::F, op::OP, initop::I, arrays) where {F, OP, I} out = arrays[1] iout = ParentIndex(offset(out) + 1) @@ -199,11 +200,14 @@ function _mapreduce_block!( f::F, op::OP, initop::I, dims, strides, offsets, costs, arrays ) where {F, OP, I} bytestrides = map((s, stride) -> s .* stride, sizeof.(eltype.(arrays)), strides) - strideorders = map(indexorder, strides) - blocks = _computeblocks(dims, costs, bytestrides, strideorders) - - # t = @elapsed _computeblocks(dims, costs, bytestrides, strideorders) - # println("_computeblocks time: $t") + memoryregion = totalmemoryregion(dims, bytestrides) + # When the estimated footprint fits, the full dimensions are already optimal. + blocks = if memoryregion <= BLOCKMEMORYSIZE + dims + else + strideorders = map(indexorder, strides) + _computeblocks(dims, costs, bytestrides, strideorders, BLOCKMEMORYSIZE, memoryregion) + end if get_num_threads() == 1 || prod(dims) <= MINTHREADLENGTH _mapreduce_kernel!(f, op, initop, dims, blocks, arrays, strides, offsets) @@ -594,9 +598,10 @@ function _computeblocks( dims::NTuple{N, Int}, costs::NTuple{N, Int}, bytestrides::Tuple{Vararg{NTuple{N, Int}}}, strideorders::Tuple{Vararg{NTuple{N, Int}}}, - blocksize::Int = BLOCKMEMORYSIZE + blocksize::Int = BLOCKMEMORYSIZE, + memoryregion::Int = totalmemoryregion(dims, bytestrides) ) where {N} - if totalmemoryregion(dims, bytestrides) <= blocksize + if memoryregion <= blocksize return dims end minstrideorder = minimum(minimum.(strideorders)) diff --git a/test/gpu.jl b/test/gpu.jl index 219cd69..8c21e82 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -185,11 +185,11 @@ end 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 +@testset "single-element StridedView ($AT)" for AT in ATs + @testset for T in (Float32, ComplexF32), dims in ((), (1,), (1, 1), (1, 1, 1)) + R = fill(rand(T), dims) A = StridedView(AT(R)) - @test ndims(A) == 0 + @test size(A) == dims # full reductions @test sum(A) == sum(R) @@ -206,8 +206,8 @@ end GPUArrays.@allowscalar begin @test mapx[] == 2 * R[] end - B = StridedView(AT(fill(zero(T)))) - map!(x -> x + one(T), B, A) + B = StridedView(AT(fill(zero(T), dims))) + map!(x -> x + one(x), B, A) GPUArrays.@allowscalar begin @test B[] == collect(R)[] + one(T) end @@ -220,11 +220,11 @@ end @test B[] == one(T) end - # offset handling: 0-dim views into a larger parent + # offset handling: singleton 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], ()) + s = sreshape(StridedView(Psrc)[4:4], dims) + d = sreshape(StridedView(Pdst)[3:3], dims) GPUArrays.@allowscalar begin @test sum(s) == Psrc[4] end @@ -235,13 +235,13 @@ end # low-level in-place reduction with a custom initop Pd = AT(rand(T, 5)) - d2 = sreshape(StridedView(Pd)[2:2], ()) + d2 = sreshape(StridedView(Pd)[2:2], dims) GPUArrays.@allowscalar begin prev = Pd[2] end - Strided._mapreducedim!(cos, +, identity, (), (d2, A)) + Strided._mapreducedim!(cos, +, identity, dims, (d2, A)) GPUArrays.@allowscalar begin - @test Pd[2] == prev + cos(R[]) + @test Pd[2] ≈ prev + cos(R[]) end end end diff --git a/test/othertests.jl b/test/othertests.jl index 44eaee9..6240dea 100644 --- a/test/othertests.jl +++ b/test/othertests.jl @@ -213,11 +213,11 @@ end end end -@testset "0-dimensional (scalar) StridedView" begin - @testset for T in (Float32, Float64, ComplexF32, ComplexF64) - R = fill(rand(T)) # 0-dimensional Array +@testset "single-element StridedView" begin + @testset for T in (Float32, Float64, ComplexF32, ComplexF64), dims in ((), (1,), (1, 1), (1, 1, 1)) + R = fill(rand(T), dims) A = StridedView(R) - @test ndims(A) == 0 + @test size(A) == dims # full reductions @test sum(A) == sum(R) @@ -231,7 +231,7 @@ end # map / map! / copy! / fill! @test map(x -> 2x, A)[] == 2 * R[] - B = StridedView(fill(zero(T))) + B = StridedView(fill(zero(T), dims)) map!(x -> x + one(T), B, A) @test B[] == R[] + one(T) copy!(B, A) @@ -239,20 +239,20 @@ end fill!(B, one(T)) @test B[] == one(T) - # offset handling: 0-dim views into a larger parent + # offset handling: singleton views into a larger parent Psrc = rand(T, 5) Pdst = rand(T, 5) - s = sreshape(StridedView(Psrc)[4:4], ()) - d = sreshape(StridedView(Pdst)[3:3], ()) + s = sreshape(StridedView(Psrc)[4:4], dims) + d = sreshape(StridedView(Pdst)[3:3], dims) @test sum(s) == Psrc[4] copy!(d, s) @test Pdst[3] == Psrc[4] # low-level in-place reduction with a custom initop Pd = rand(T, 5) - d2 = sreshape(StridedView(Pd)[2:2], ()) + d2 = sreshape(StridedView(Pd)[2:2], dims) prev = Pd[2] - Strided._mapreducedim!(sin, +, identity, (), (d2, A)) + Strided._mapreducedim!(sin, +, identity, dims, (d2, A)) @test Pd[2] == prev + sin(R[]) end end