From dd03c1bb1d2c1ea50712c4b0d99bf6b6a5f3246a Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 7 Oct 2026 13:40:40 -0400 Subject: [PATCH 1/3] Skip layout planning for single-element arrays --- ext/StridedGPUArraysExt.jl | 11 +++++++++++ src/mapreduce.jl | 5 +++-- test/gpu.jl | 20 ++++++++++---------- test/othertests.jl | 20 ++++++++++---------- 4 files changed, 34 insertions(+), 22 deletions(-) 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..7585da9 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) diff --git a/test/gpu.jl b/test/gpu.jl index 219cd69..c2327e6 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,7 +206,7 @@ end GPUArrays.@allowscalar begin @test mapx[] == 2 * R[] end - B = StridedView(AT(fill(zero(T)))) + B = StridedView(AT(fill(zero(T), dims))) map!(x -> x + one(T), B, A) GPUArrays.@allowscalar begin @test B[] == collect(R)[] + one(T) @@ -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,11 +235,11 @@ 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[]) 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 From db110c20b2ce327ef5a05ab1ca9c07321bc4c85a Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Wed, 7 Oct 2026 18:01:00 -0400 Subject: [PATCH 2/3] Skip stride ranking when the memory estimate fits (#92) --- src/mapreduce.jl | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 7585da9..0f10caf 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -200,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) @@ -595,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)) From 4f4589b21a34f8ccf5132e0d7878e508b9872eaa Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 8 Oct 2026 11:06:04 -0400 Subject: [PATCH 3/3] Fix single-element GPU tests on Metal and CUDA The `map!` closure captured the loop variable `T`, which is not isbits and so cannot be passed to a GPU kernel; use `one(x)` instead. The in-place reduction check compared GPU and CPU `cos` results exactly, which can differ by an ulp; compare approximately. Co-Authored-By: Claude Opus 5.5 --- test/gpu.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/gpu.jl b/test/gpu.jl index c2327e6..8c21e82 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -207,7 +207,7 @@ end @test mapx[] == 2 * R[] end B = StridedView(AT(fill(zero(T), dims))) - map!(x -> x + one(T), B, A) + map!(x -> x + one(x), B, A) GPUArrays.@allowscalar begin @test B[] == collect(R)[] + one(T) end @@ -241,7 +241,7 @@ end end 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