Skip to content

Skip layout planning for single-element arrays #90

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lkdvos merged 3 commits into from
Oct 8, 2026
Merged

Skip layout planning for single-element arrays #90

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
dd03c1b
Skip layout planning for single-element arrays
lkdvos Oct 7, 2026
db110c2
Skip stride ranking when the memory estimate fits (#92)
lkdvos Oct 7, 2026
4f4589b
Fix single-element GPU tests on Metal and CUDA
lkdvos Oct 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions ext/StridedGPUArraysExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down
23 changes: 14 additions & 9 deletions src/mapreduce.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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))
Expand Down
24 changes: 12 additions & 12 deletions test/gpu.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
20 changes: 10 additions & 10 deletions test/othertests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -231,28 +231,28 @@ 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)
@test B[] == R[]
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
Expand Down
Loading