Non-deterministic zeros in `findall` result on ROCBackend with large arrays

[LasNikas]

On Julia v1.12 with AMDGPU.jl v1.3.4, running on an MI300A, I observe non-deterministic behavior when calling findall on a large array. In particular, the number of returned indices is correct, but some of the indices are 0. The problem seems more likely to appear for very large n.
Note: Seems related to #41 , but still occurs with AMDGPU.jl v1.3.4.

cc @vchuravy

MWE 1:

using KernelAbstractions
using Adapt
using AMDGPU

n = 100_000_000
a_gpu = Adapt.adapt(ROCBackend(), fill(true, n))
any(i -> i == 0, findall(x -> x == true, a_gpu)) # sometimes returns true

MWE 2:

using KernelAbstractions
using Adapt
using AMDGPU

n_true = n_false = 100_000_000
a_gpu = Adapt.adapt(ROCBackend(), vcat(fill(true, n_true), fill(false, n_false)))
indices_true = findall(x -> x == true, a_gpu)

@show length(indices_true) # correct length
any(i -> i == 0, indices_true) # true

[vchuravy]

So crucially this actually seems to lead to segmentation faults / errors.
Running with julia -g2 --check-bounds=2

julia> indices_true = findall(x -> x == true, a_gpu)
67953408-element ROCArray{Int64, 1, AMDGPU.Runtime.Mem.HIPBuffer}:
Error showing value of type ROCArray{Int64, 1, AMDGPU.Runtime.Mem.HIPBuffer}:

SYSTEM (REPL): showing an error caused an error
ERROR: 1-element ExceptionStack:
GPU Kernel Exception
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:44
  [2] throw_if_exception(dev::HIPDevice)
    @ AMDGPU ~/.julia/packages/AMDGPU/cYhJO/src/exception_handler.jl:124
  [3] synchronize(stm::HIPStream; blocking::Bool, stop_hostcalls::Bool)
    @ AMDGPU ~/.julia/packages/AMDGPU/cYhJO/src/highlevel.jl:40
  [4] synchronize
    @ ~/.julia/packages/AMDGPU/cYhJO/src/highlevel.jl:36 [inlined]
  [5] copyto!(dest::Vector{…}, d_offset::Int64, source::ROCArray{…}, s_offset::Int64, amount::Int64; async::Bool)
    @ AMDGPU ~/.julia/packages/AMDGPU/cYhJO/src/array.jl:147
  [6] copyto!
    @ ~/.julia/packages/AMDGPU/cYhJO/src/array.jl:131 [inlined]
  [7] copyto!
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/abstractarray.jl:179 [inlined]
  [8] copyto_axcheck!

julia> for i in 1:10
           let n_true = 100_000_000
               a_gpu = Adapt.adapt(ROCBackend(), fill(true, n_true))
               indices_true = findall(x -> x == true, a_gpu)
               length(indices_true)
               @assert !any(i -> i == 0, indices_true)
           end
       end
Memory access fault by GPU node-4 (Agent handle: 0x45c6d40) on address 0x1467ff423000. Reason: Unknown.

[560922] signal 6 (-6): Aborted
in expression starting at REPL[4]:1
Allocations: 58882826 (Pool: 58879181; Big: 3645); GC: 29
Aborted (core dumped)

[giordano]

To hopefully help with debugging, I attach the compressed archive of the output directory of

julia> using AcceleratedKernels: AcceleratedKernels as AK

julia> using Adapt, AMDGPU

julia> n = 100_000_000
100000000

julia> a_cpu = fill(true, n);

julia> a_gpu = Adapt.adapt(ROCBackend(), fill(true, n));

julia> @device_code dir="accumulate" AK.accumulate(+, a_gpu, ROCBackend(); init=false)

accumulate.tar.gz.

CC @pxl-th.

[giordano]

As requested by @vchuravy on Slack, using the ScanPrefixes alghoritm mentioned at #41 (comment) gives the right result:

julia> sum(AK.cumsum(a_gpu, ROCBackend(); alg=AK.ScanPrefixes()))
5000000050000000

[giordano]

This issue seems to be fixed by #44:

julia> using AcceleratedKernels: AcceleratedKernels as AK

julia> using Adapt, AMDGPU

julia> n = 100_000_000
100000000

julia> a_cpu = fill(true, n);

julia> a_gpu = Adapt.adapt(ROCBackend(), fill(true, n));

julia> sum(AK.accumulate(+, a_gpu, ROCBackend(); init=false))
5000000050000000

julia> ans == (n ^ 2 + n) / 2
true

and also

julia> using AMDGPU

julia> using Adapt

julia> for i in 1:10
           let n_true = 100_000_000
               a_gpu = Adapt.adapt(ROCBackend(), fill(true, n_true))
               indices_true = findall(x -> x == true, a_gpu)
               length(indices_true)
               @assert !any(i -> i == 0, indices_true)
           end
       end

julia>

On stable v0.4.1 instead I get

Memory access fault by GPU node-4 (Agent handle: 0x45c6d40) on address 0x1467ff423000. Reason: Unknown.

[560922] signal 6 (-6): Aborted
in expression starting at REPL[4]:1
Allocations: 58882826 (Pool: 58879181; Big: 3645); GC: 29
Aborted (core dumped)

[vchuravy]

I think the segmentation fault might be #47

[giordano]

Right, with that PR only I get an assertion failure:

julia> for i in 1:10
           let n_true = 100_000_000
               a_gpu = Adapt.adapt(ROCBackend(), fill(true, n_true))
               indices_true = findall(x -> x == true, a_gpu)
               length(indices_true)
               @assert !any(i -> i == 0, indices_true)
           end
       end
ERROR: AssertionError: !(any((i->begin
                #= REPL[4]:6 =#
                i == 0
            end), indices_true))
Stacktrace:
 [1] top-level scope
   @ REPL[4]:6

which is the proper issue reported in this ticket.

[giordano]

For reference, here is the same archive as the one in #46 (comment), but when checking out #44: accumulate.tar.gz. At a quick glance, of the three LLVM optimised modules, the only relevant differences are in gpu__accumulate_previous!_3.opt.ll (the other two are identical when excluding differences in file paths in debug symbols and numerical values of !dbg annotations).

[anicusan]

We made the ScanPrefixes algorithm the default as it is almost always faster on real world data - and doesn't depend on finnicky cross-workgroup communication. The above fix is now the default in AK-0.4.3, so I will close this issue.

Thank you everyone for the amount of debugging effort that went into the whole stack - AMDGPU - AcceleratedKernels - KernelAbstractions - GPUCompiler. Sherlock Holmes got nothing on JuliaGPU.