Inconsistent behavior with multiple GPU

Hello!

I am recently trying to exploit the single-process multi GPU capabilities of CUDA.jl 3.8.0 to run on multiple GPU.

Unfortunately, for some weird reason, my kernel fails with ERROR: LoadError: CUDA error: invalid resource handle (code 400, ERROR_INVALID_HANDLE)
I double checked that the CuArrays are allocated on the correct device. Also, the kernel does not fail at the first run but after iterating a couple of times. This might suggest I have problems of synchronization? On the other hand, using @cuda does not show the same problem.

Here a MWE

using KernelAbstractions, CUDA, CUDAKernels, Test

function cuda_add_1!(c)
    i = threadIdx().x 
    j = blockIdx().x
    @inbounds c[i, j] += 1
    return nothing
end

@kernel function kernel_add_1!(c)
    i, j = @index(Global, NTuple)
    @inbounds c[i, j] += 1
end

function single_gpu_add1!(c, method) 
    if method == :cuda
        @cuda threads=5 blocks=5 cuda_add_1!(c)
        event = Event(CUDAKernels.CUDADevice())
    else
        test! = kernel_add_1!(CUDADevice(), 1)
        event = test!(c, ndrange=size(c))
    end

    return event
end

function multi_gpu_add_1(method, devices)
    r = []
    for dev in devices
        CUDA.device!(dev)
        push!(r, CuArray(rand(5, 5)))
    end
    for dev in devices
        CUDA.device!(dev)
        CUDA.synchronize()
    end
    for (i, dev) in enumerate(devices)
        CUDA.device!(dev)
        event = single_gpu_add1!(r[i], method)
        wait(CUDADevice(), event)
    end
    for (i, dev) in enumerate(devices)
        CUDA.device!(dev)
        CUDA.synchronize()
        @test all(Array(r[i]) .> 1)
    end
end

devices = ((CuDevice(0), ),
           (CuDevice(0), CuDevice(0)), 
           (CuDevice(0), CuDevice(1)))
for dev in devices, method in (:cuda, :kernel)
    @info "Running $method on $dev"
    for test_num in 1:10
        multi_gpu_add_1(method, dev) 
    end
end

Maybe I am doing something stupid but I don't really understand the output

julia> include("test-kernel-abstractions.jl")
[ Info: Running cuda on (CuDevice(0),)
[ Info: Running kernel on (CuDevice(0),)
[ Info: Running cuda on (CuDevice(0), CuDevice(0))
[ Info: Running kernel on (CuDevice(0), CuDevice(0))
[ Info: Running cuda on (CuDevice(0), CuDevice(1))
[ Info: Running kernel on (CuDevice(0), CuDevice(1))
ERROR: LoadError: CUDA error: invalid resource handle (code 400, ERROR_INVALID_HANDLE)
Stacktrace:
  [1] throw_api_error(res::CUDA.cudaError_enum)
    @ CUDA ~/.julia/packages/CUDA/bki2w/lib/cudadrv/error.jl:91
  [2] macro expansion
    @ ~/.julia/packages/CUDA/bki2w/lib/cudadrv/error.jl:101 [inlined]
  [3] cuLaunchKernel(f::CuFunction, gridDimX::UInt32, gridDimY::UInt32, gridDimZ::UInt32, blockDimX::UInt32, blockDimY::UInt32, blockDimZ::UInt32, sharedMemBytes::Int64, hStream::CuStream, kernelParams::Vector{Ptr{Nothing}}, extra::Ptr{Nothing})
    @ CUDA ~/.julia/packages/CUDA/bki2w/lib/utils/call.jl:26
  [4] #35
    @ ~/.julia/packages/CUDA/bki2w/lib/cudadrv/execution.jl:69 [inlined]
  [5] macro expansion
    @ ~/.julia/packages/CUDA/bki2w/lib/cudadrv/execution.jl:33 [inlined]
  [6] macro expansion
    @ ./none:0 [inlined]
  [7] pack_arguments(::CUDA.var"#35#36"{Bool, Int64, CuStream, CuFunction, CuDim3, CuDim3}, ::CUDA.KernelState, ::KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicCheck, Nothing, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, KernelAbstractions.NDIteration.NDRange{2, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.StaticSize{(1, 1)}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, Nothing}}, ::CuDeviceMatrix{Float64, 1})
    @ CUDA ./none:0
  [8] #launch#34
    @ ~/.julia/packages/CUDA/bki2w/lib/cudadrv/execution.jl:62 [inlined]
  [9] #40
    @ ~/.julia/packages/CUDA/bki2w/lib/cudadrv/execution.jl:136 [inlined]
 [10] macro expansion
    @ ~/.julia/packages/CUDA/bki2w/lib/cudadrv/execution.jl:95 [inlined]
 [11] macro expansion
    @ ./none:0 [inlined]
 [12] convert_arguments
    @ ./none:0 [inlined]
 [13] #cudacall#39
    @ ~/.julia/packages/CUDA/bki2w/lib/cudadrv/execution.jl:135 [inlined]
 [14] macro expansion
    @ ~/.julia/packages/CUDA/bki2w/src/compiler/execution.jl:204 [inlined]
 [15] macro expansion
    @ ./none:0 [inlined]
 [16] call(::CUDA.HostKernel{typeof(Cassette.overdub), Tuple{Cassette.Context{nametype(CUDACtx), Nothing, Nothing, KernelAbstractions.var"##PassType#257", Nothing, Cassette.DisableHooks}, typeof(gpu_kernel_add_1!), KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicCheck, Nothing, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, KernelAbstractions.NDIteration.NDRange{2, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.StaticSize{(1, 1)}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, Nothing}}, CuDeviceMatrix{Float64, 1}}}, ::Cassette.Context{nametype(CUDACtx), Nothing, Nothing, KernelAbstractions.var"##PassType#257", Nothing, Cassette.DisableHooks}, ::typeof(gpu_kernel_add_1!), ::KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicCheck, Nothing, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, KernelAbstractions.NDIteration.NDRange{2, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.StaticSize{(1, 1)}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, Nothing}}, ::CuDeviceMatrix{Float64, 1}; call_kwargs::Base.Iterators.Pairs{Symbol, Any, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:threads, :blocks, :stream), Tuple{Int64, Int64, CuStream}}})
    @ CUDA ./none:0
 [17] (::CUDA.HostKernel{typeof(Cassette.overdub), Tuple{Cassette.Context{nametype(CUDACtx), Nothing, Nothing, KernelAbstractions.var"##PassType#257", Nothing, Cassette.DisableHooks}, typeof(gpu_kernel_add_1!), KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicCheck, Nothing, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, KernelAbstractions.NDIteration.NDRange{2, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.StaticSize{(1, 1)}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, Nothing}}, CuDeviceMatrix{Float64, 1}}})(::Cassette.Context{nametype(CUDACtx), Nothing, Nothing, KernelAbstractions.var"##PassType#257", Nothing, Cassette.DisableHooks}, ::Vararg{Any, N} where N; threads::Int64, blocks::Int64, kwargs::Base.Iterators.Pairs{Symbol, CuStream, Tuple{Symbol}, NamedTuple{(:stream,), Tuple{CuStream}}})
    @ CUDA ~/.julia/packages/CUDA/bki2w/src/compiler/execution.jl:462
 [18] (::KernelAbstractions.Kernel{CUDADevice, KernelAbstractions.NDIteration.StaticSize{(1,)}, KernelAbstractions.NDIteration.DynamicSize, typeof(gpu_kernel_add_1!)})(args::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}; ndrange::Tuple{Int64, Int64}, dependencies::CUDAKernels.CudaEvent, workgroupsize::Nothing, progress::Function)
    @ CUDAKernels ~/.julia/packages/CUDAKernels/kCOA4/src/CUDAKernels.jl:223
 [19] single_gpu_add1!(c::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, method::Symbol)
    @ Main ~/test-architectures/test-kernel-abstractions.jl:21
 [20] multi_gpu_add_1(method::Symbol, devices::Tuple{CuDevice, CuDevice})
    @ Main ~/test-architectures/test-kernel-abstractions.jl:39
 [21] top-level scope
    @ ~/test-architectures/test-kernel-abstractions.jl:54
 [22] include(fname::String)
    @ Base.MainInclude ./client.jl:444
 [23] top-level scope
    @ REPL[4]:1
 [24] top-level scope
    @ ~/.julia/packages/CUDA/bki2w/src/initialization.jl:52

I am using

      Status `~/.julia/environments/v1.6/Project.toml`
  [052768ef] CUDA v3.8.0
  [72cfdca4] CUDAKernels v0.3.3
  [63c18a36] KernelAbstractions v0.7.2

Start with the MWE in test-kernel-abstractions.jl and reproduce the multi-device case that fails in CUDAKernels.jl while the direct @cuda path succeeds. Compare the device and stream handling around the KernelAbstractions launch, then verify that repeated runs on CuDevice(0) and CuDevice(1) complete without ERROR_INVALID_HANDLE and that the existing assertions pass.