diff --git a/Project.toml b/Project.toml index b5162d14..cde1d190 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "AcceleratedKernels" uuid = "6a4ca0a5-0e36-4168-a932-d9be78d558f1" -version = "0.5.1" +version = "0.5.2" authors = ["Andrei-Leonard Nicusan and contributors"] [deps] diff --git a/src/findall.jl b/src/findall.jl index e1aa6e20..e1da08a1 100644 --- a/src/findall.jl +++ b/src/findall.jl @@ -143,13 +143,12 @@ end function findall_gpu( bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, alg::ScanScatter; - bufs, + bufs, like, ) where I block_size = alg.block_size items_per_thread = alg.items_per_thread - # The output on the resolved backend, which a range does not determine - isempty(bools) && return KernelAbstractions.allocate(backend, I, 0) + isempty(bools) && return _similar(backend, like, I, (0,)) elems_per_block = block_size * items_per_thread num_blocks = cld(length(bools), elems_per_block) @@ -163,7 +162,7 @@ function findall_gpu( _accumulate_nested!(+, block_counts, bufs.scan; backend, init=0) n = @allowscalar block_counts[end] - out = KernelAbstractions.allocate(backend, I, n) + out = _similar(backend, like, I, (n,)) if n > 0 kernel!(out, bools, block_counts, input_indices, output_indices, items; ndrange=num_blocks * block_size) @@ -213,11 +212,14 @@ function findall_section!(out, bools, input_indices, output_indices, positions, end +# The result is an array like `like` (the searched array, rather than a mask of it in scratch +# memory), so that it shares properties such as a GPU array's storage mode function findall_impl( bools::AbstractArray{Bool}, ::Type{I}, output_indices, backend::Backend, alg; bufs, + like=bools, ) where I if alg isa ScanScatter - findall_gpu(bools, I, output_indices, backend, alg; bufs) + findall_gpu(bools, I, output_indices, backend, alg; bufs, like) else findall_cpu(bools, I, output_indices, backend, alg; bufs) end @@ -331,5 +333,5 @@ function _findall_run(pred, v, items, s, bufs) else v end - findall_impl(bools, eltype(items), items, backend, a; bufs) + findall_impl(bools, eltype(items), items, backend, a; bufs, like=v) end diff --git a/test/Project.toml b/test/Project.toml index 9ce1e943..71c42482 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -16,6 +16,9 @@ OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" +[sources] +AcceleratedKernels = {path = ".."} + [compat] ParallelTestRunner = "2" oneAPI = "2.9.2"