diff --git a/src/reverse.jl b/src/reverse.jl index 02f5723..70cec27 100644 --- a/src/reverse.jl +++ b/src/reverse.jl @@ -1,7 +1,138 @@ +# Reversing. Two regimes: +# * `dims=:` (default) reverses the whole array, a reversal of the column-major linear order, +# taken by the fast flat path below (swap mirrored pairs, no index arithmetic). +# * `dims=d` reverses only along those dimensions, via the general ND kernel below. + + +# Index math for the `dims` reversal (same mapping as `Base.reverse`): maps element `i` to its +# mirror. Layout-agnostic via LinearIndices/CartesianIndices, so reshaped/strided arrays work. +@inline function _reverse_out_index(i, nd_idx, lin_idx, rev_dims, ref) + idx = Tuple(nd_idx[i]) + idx_mirror = ifelse.(rev_dims, ref .- idx, idx) + lin_idx[idx_mirror...] +end + +@inline function _reverse_swap_indices(i, nd_idx, lin_idx, rev_dims, ref) + idx = Tuple(nd_idx[i]) + index_in = lin_idx[idx...] + idx_mirror = ifelse.(rev_dims, ref .- idx, idx) + index_out = lin_idx[idx_mirror...] + index_in, index_out +end + + +# GPU kernels for the `dims` reversal, one thread per element. The index math runs directly in the +# kernel body, rather than through a `foreachindex` closure, so it inlines fully. +@kernel inbounds=true cpu=false unsafe_indices=true function _reverse_oop_kernel!( + dst, src, nd_idx, lin_idx, rev_dims, ref, len, +) + block_size = @groupsize()[1] + i = @index(Local, Linear) + (@index(Group, Linear) - 0x1) * block_size + if i <= len + dst[_reverse_out_index(i, nd_idx, lin_idx, rev_dims, ref)] = src[i] + end +end + +@kernel inbounds=true cpu=false unsafe_indices=true function _reverse_inplace_kernel!( + v, nd_idx, lin_idx, rev_dims, ref, len, +) + block_size = @groupsize()[1] + i = @index(Local, Linear) + (@index(Group, Linear) - 0x1) * block_size + if i <= len + index_in, index_out = _reverse_swap_indices(i, nd_idx, lin_idx, rev_dims, ref) + if index_in < index_out + temp = v[index_out] + v[index_out] = v[index_in] + v[index_in] = temp + end + end +end + + +# Validate `dims`: a `Colon`, an integer, or an iterable of integers within `1:ndims(A)`. +function _check_reverse_dims(A, dims) + dims isa Colon && return + applicable(iterate, dims) || throw(ArgumentError("dimension $dims is not iterable")) + for d in dims # an integer iterates once + d isa Integer || + throw(ArgumentError("reversed dimension(s) must be integers, got $dims")) + 1 <= d <= ndims(A) || + throw(ArgumentError("dimension $dims is not 1 ≤ dims ≤ $(ndims(A))")) + end + return +end + + +# In-place: split along the last non-singleton reversed dim so only ~half the elements need a +# thread, each swapping with its mirror. +function _reverse_dims!( + v::AbstractArray{T, N}, dims, backend; + max_tasks=Threads.nthreads(), min_elems=1, prefer_threads=true, block_size=256, +) where {T, N} + rev_dims = ntuple(d -> (d in dims) && size(v, d) > 1, N) + half_dim = findlast(rev_dims) + isnothing(half_dim) && return v # all reversed dims are singletons + + ref = size(v) .+ 1 + lin_idx = LinearIndices(v) + reduced_size = ntuple(d -> ifelse(d == half_dim, cld(size(v, d), 2), size(v, d)), N) + nd_idx = CartesianIndices(reduced_size) + len = Base.prod(reduced_size) + len == 0 && return v + + if use_gpu_algorithm(backend, prefer_threads) + _reverse_inplace_kernel!(backend, block_size)( + v, nd_idx, lin_idx, rev_dims, ref, len, + ndrange = block_size * cld(len, block_size), + ) + else + foreachindex(1:len, backend; max_tasks, min_elems, prefer_threads) do i + index_in, index_out = _reverse_swap_indices(i, nd_idx, lin_idx, rev_dims, ref) + @inbounds if index_in < index_out + temp = v[index_out] + v[index_out] = v[index_in] + v[index_in] = temp + end + end + end + + v +end + + +# Out-of-place: one thread per element copies `src[i]` to its mirror slot. +function _reverse_dims!( + dst::AbstractArray{T, N}, src::AbstractArray{T, N}, dims, backend; + max_tasks=Threads.nthreads(), min_elems=1, prefer_threads=true, block_size=256, +) where {T, N} + rev_dims = ntuple(d -> (d in dims) && size(src, d) > 1, N) + ref = size(src) .+ 1 + lin_idx = LinearIndices(src) + nd_idx = CartesianIndices(src) + len = length(src) + len == 0 && return dst + + if use_gpu_algorithm(backend, prefer_threads) + _reverse_oop_kernel!(backend, block_size)( + dst, src, nd_idx, lin_idx, rev_dims, ref, len, + ndrange = block_size * cld(len, block_size), + ) + else + foreachindex(src, backend; max_tasks, min_elems, prefer_threads) do i + @inbounds dst[_reverse_out_index(i, nd_idx, lin_idx, rev_dims, ref)] = src[i] + end + end + + dst +end + + """ reverse!( v::AbstractArray, backend::Backend=get_backend(v); + dims=:, + # CPU settings max_tasks=Threads.nthreads(), min_elems=1, @@ -10,12 +141,15 @@ block_size=256, ) -Reverse `v` in-place and return it. The CPU and GPU settings are the same as for -[`foreachindex`](@ref). +Reverse `v` in-place and return it. With `dims=:` (the default) the whole array is reversed; pass +`dims=d` (an integer or an iterable of integers) to reverse only along those dimensions, matching +`Base.reverse!`. The CPU and GPU settings are the same as for [`foreachindex`](@ref). + +For the whole-array case each thread swaps one symmetric pair `v[i] <-> v[end - i + 1]`, so only +`length(v) ÷ 2` threads are launched and no temporary array is allocated. Arrays of odd length keep +their middle element in place. -Each thread swaps one symmetric pair `v[i] <-> v[end - i + 1]`, so only `length(v) ÷ 2` threads -are launched and no temporary array is allocated. Arrays of odd length keep their middle element -in place. +To reverse a contiguous sub-range of a vector, reverse a view: `AK.reverse!(@view v[lo:hi])`. # Examples ```julia @@ -24,19 +158,27 @@ import AcceleratedKernels as AK v = CUDA.CuArray(1:100_000) AK.reverse!(v) + +m = CUDA.CuArray(reshape(1:12, 3, 4)) +AK.reverse!(m; dims=2) # reverse the columns ``` """ function reverse!( v::AbstractArray, backend::Backend=get_backend(v); - kwargs... + dims=:, kwargs... ) + _check_reverse_dims(v, dims) + if !(dims isa Colon) + return _reverse_dims!(v, dims, backend; kwargs...) + end + len = length(v) len <= 1 && return v lo = firstindex(v) hi = lastindex(v) - # Only the lower half needs threads - each one swaps its mirrored partner too; for odd + # Only the lower half needs threads, each swapping its mirrored partner too; for odd # lengths the middle element is its own mirror, so it is correctly left untouched foreachindex(1:(len ÷ 2), backend; kwargs...) do i left = lo + i - 1 @@ -56,6 +198,8 @@ end reverse!( dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(src); + dims=:, + # CPU settings max_tasks=Threads.nthreads(), min_elems=1, @@ -65,13 +209,21 @@ end ) Write the reverse of `src` into `dst` and return `dst`; `src` is left unchanged. `dst` and `src` -must have the same length and must not alias. The CPU and GPU settings are the same as for +must have the same size and must not alias. With `dims=:` (the default) the whole array is reversed; +pass `dims=d` to reverse only along those dimensions. The CPU and GPU settings are the same as for [`foreachindex`](@ref). """ function reverse!( dst::AbstractArray, src::AbstractArray, backend::Backend=get_backend(src); - kwargs... + dims=:, kwargs... ) + _check_reverse_dims(src, dims) + if !(dims isa Colon) + @argcheck size(dst) == size(src) + length(src) == 0 && return dst + return _reverse_dims!(dst, src, dims, backend; kwargs...) + end + @argcheck length(dst) == length(src) length(src) == 0 && return dst @@ -90,6 +242,8 @@ end reverse( v::AbstractArray, backend::Backend=get_backend(v); + dims=:, + # CPU settings max_tasks=Threads.nthreads(), min_elems=1, @@ -98,15 +252,15 @@ end block_size=256, ) -Return a reversed copy of `v`, leaving `v` unchanged. The CPU and GPU settings are the same as for -[`foreachindex`](@ref). +Return a reversed copy of `v`, leaving `v` unchanged. With `dims=:` (the default) the whole array is +reversed; pass `dims=d` to reverse only along those dimensions, matching `Base.reverse`. The CPU and +GPU settings are the same as for [`foreachindex`](@ref). -Prefer [`reverse!`](@ref) when you do not need to keep `v` - it avoids the allocation. +Prefer [`reverse!`](@ref) when you do not need to keep `v`; it avoids the allocation. """ function reverse( v::AbstractArray, backend::Backend=get_backend(v); kwargs... ) - dst = similar(v) - reverse!(dst, v, backend; kwargs...) + reverse!(similar(v), v, backend; kwargs...) end diff --git a/test/generic/reverse.jl b/test/generic/reverse.jl index 4f32aad..42abc97 100644 --- a/test/generic/reverse.jl +++ b/test/generic/reverse.jl @@ -91,4 +91,78 @@ @test Array(v) == reverse(h) end end + + # N-dimensional reversal along a subset of dimensions (Base.reverse parity) + @testset "dims" begin + # Single dimension, including a degenerate size-1 dim and a large 3-D array + for shape in ([1, 2, 4, 3], [4, 2], [5], [8, 8, 8]), + dim in 1:length(shape) + + h = rand(Float32, shape...) + + v = array_from_host(h) + AK.reverse!(v; dims=dim, prefer_threads) + @test Array(v) == reverse(h; dims=dim) + + src = array_from_host(h) + out = AK.reverse(src; dims=dim, prefer_threads) + @test Array(out) == reverse(h; dims=dim) + @test Array(src) == h # source left untouched + + dst = array_from_host(zeros(Float32, shape...)) + AK.reverse!(dst, src; dims=dim, prefer_threads) + @test Array(dst) == reverse(h; dims=dim) + end + + # Multiple dimensions at once, plus dims=: (dispatches to the flat whole-array path). + # The odd sizes of [7, 6, 5] exercise the in-place middle-plane swaps, where only the + # index ordering guard stops a pair from being swapped twice + for shape in ([1, 2, 4, 3], [8, 8, 8], [7, 6, 5]), + dims in ((1, 2), (2, 3), (1, 3), :) + + h = rand(Float32, shape...) + + v = array_from_host(h) + AK.reverse!(v; dims=dims, prefer_threads) + @test Array(v) == reverse(h; dims=dims) + + out = AK.reverse(array_from_host(h); dims=dims, prefer_threads) + @test Array(out) == reverse(h; dims=dims) + + src = array_from_host(h) + dst = array_from_host(zeros(Float32, shape...)) + AK.reverse!(dst, src; dims=dims, prefer_threads) + @test Array(dst) == reverse(h; dims=dims) + end + + # Any iterable of integers works, e.g. a Vector (Base only accepts tuples) + h = rand(Float32, 4, 5, 6) + out = AK.reverse(array_from_host(h); dims=[1, 3], prefer_threads) + @test Array(out) == reverse(h; dims=(1, 3)) + + # Empty arrays are returned unchanged + h = zeros(Float32, 0, 5) + for dims in (1, 2, (1, 2)) + v = array_from_host(h) + @test Array(AK.reverse!(v; dims, prefer_threads)) == reverse(h; dims) + + dst = array_from_host(copy(h)) + @test Array(AK.reverse!(dst, v; dims, prefer_threads)) == reverse(h; dims) + + @test Array(AK.reverse(v; dims, prefer_threads)) == reverse(h; dims) + end + end + + # Invalid dims arguments throw, matching Base/CUDA + @testset "dims errors" begin + v = array_from_host(rand(Float32, 2, 3, 4)) + @test_throws ArgumentError AK.reverse!(v; dims=0, prefer_threads) + @test_throws ArgumentError AK.reverse!(v; dims=4, prefer_threads) + @test_throws ArgumentError AK.reverse(v; dims=0, prefer_threads) + @test_throws ArgumentError AK.reverse(v; dims=4, prefer_threads) + + # Non-integer dims must throw rather than silently do nothing + @test_throws ArgumentError AK.reverse!(v; dims=1.5, prefer_threads) + @test_throws ArgumentError AK.reverse(v; dims=(1, 2.5), prefer_threads) + end end