Allocate findall's result like the searched array

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 <leonard@evophase.co.uk> and contributors"]
 
 [deps]

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

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"