Add workaround for system atomics

Co-authored-by: Valentin Churavy <v.churavy@gmail.com>

Author: efaulhaber

ext/PointNeighborsCUDAExt.jl

@@ -2,11 +2,12 @@ module PointNeighborsCUDAExt
 
 using PointNeighbors: PointNeighbors, CUDAMultiGPUBackend, DynamicVectorOfVectors,
                       GridNeighborhoodSearch, FullGridCellList, extract_svector,
-                      cell_coords, cell_index, check_cell_bounds, pushat_atomic!
+                      cell_coords, cell_index, @threaded, generic_kernel
 using CUDA: CUDA, CuArray, CUDABackend, cu
 using UnsafeAtomics: UnsafeAtomics
 using PointNeighbors.KernelAbstractions: KernelAbstractions, @kernel, @index
 using PointNeighbors.Adapt: Adapt
+using Base: @propagate_inbounds
 
 const UnifiedCuArray = CuArray{<:Any, <:Any, CUDA.UnifiedMemory}
 
@@ -24,19 +25,6 @@ end
 end
 
 @inline function PointNeighbors.parallel_foreach(f::T, iterator, x::CUDAMultiGPUBackend) where {T}
-    parallel_foreach2((i, gpu) -> @inline f(i), iterator, x)
-end
-
-@inline function parallel_foreach2(f, iterator, x::UnifiedCuArray)
-    parallel_foreach2(f, iterator, CUDAMultiGPUBackend())
-end
-
-KernelAbstractions.@kernel function generic_kernel2(f, gpu)
-    i = KernelAbstractions.@index(Global)
-    @inline f(i, gpu)
-end
-
-@inline function parallel_foreach2(f, iterator, x::CUDAMultiGPUBackend)
     # On the GPU, we can only loop over `1:N`. Therefore, we loop over `1:length(iterator)`
     # and index with `iterator[eachindex(iterator)[i]]`.
     # Note that this only works with vector-like iterators that support arbitrary indexing.
@@ -64,8 +52,8 @@ end
         CUDA.device!(i - 1)
 
         # Call the generic kernel, which only calls a function with the global GPU index
-        generic_kernel2(backend)(i, ndrange = length(indices_)) do j, gpu
-            @inbounds @inline f(iterator[indices_[j]], gpu)
+        generic_kernel(backend)(ndrange = length(indices_)) do j
+            @inbounds @inline f(iterator[indices_[j]])
         end
     end
 
@@ -79,93 +67,140 @@ end
     CUDA.device!(0)
 end
 
-struct MultiGPUVectorOfVectors{T, VOV, V} <: AbstractVector{Array{T, 1}}
-    vov::VOV
-    vov_per_gpu::V
-end
-
-function MultiGPUVectorOfVectors(vov, vov_per_gpu)
-    MultiGPUVectorOfVectors{eltype(vov.backend), typeof(vov), typeof(vov_per_gpu)}(vov, vov_per_gpu)
+function atomic_system_add(ptr::CUDA.LLVMPtr{Int32, CUDA.AS.Global}, val::Int32)
+    CUDA.LLVM.Interop.@asmcall(
+        "atom.sys.global.add.u32 \$0, [\$1], \$2;",
+        "=r,l,r,~{memory}",
+        true, Int32, Tuple{CUDA.LLVMPtr{Int32, CUDA.AS.Global}, Int32},
+        ptr, val
+    )
 end
 
-Adapt.@adapt_structure(MultiGPUVectorOfVectors)
+@inline function pushat_atomic_system!(vov, i, value)
+    (; backend, lengths) = vov
 
-function Adapt.adapt_structure(to::CUDAMultiGPUBackend, vov::DynamicVectorOfVectors{T}) where {T}
-    max_inner_length, max_outer_length = size(vov.backend)
+    @boundscheck checkbounds(vov, i)
 
-    n_gpus = length(CUDA.devices())
-    vov_per_gpu = [DynamicVectorOfVectors{T}(; max_outer_length, max_inner_length) for _ in 1:n_gpus]
+    # Increment the column length with an atomic add to avoid race conditions.
+    # Store the new value since it might be changed immediately afterwards by another
+    # thread.
+    # new_length = Atomix.@atomic lengths[i] += 1
+    new_length = atomic_system_add(pointer(lengths, i), Int32(1)) + 1
 
-    vov_ = DynamicVectorOfVectors(Adapt.adapt(to, vov.backend), Adapt.adapt(to, vov.length_), Adapt.adapt(to, vov.lengths))
-    vov_per_gpu_ = ntuple(i -> Adapt.adapt(CuArray, vov_per_gpu[i]), n_gpus)
-    for vov__ in vov_per_gpu_
-        vov__.length_[] = vov.length_[]
-    end
+    # We can write here without race conditions, since the atomic add guarantees
+    # that `new_length` is different for each thread.
+    backend[new_length, i] = value
 
-    return MultiGPUVectorOfVectors{T, typeof(vov_), typeof(vov_per_gpu_)}(vov_, vov_per_gpu_)
+    return vov
 end
 
-const MultiGPUNeighborhoodSearch = GridNeighborhoodSearch{<:Any, <:Any, <:FullGridCellList{<:MultiGPUVectorOfVectors}}
+const MultiGPUNeighborhoodSearch = GridNeighborhoodSearch{<:Any, <:Any, <:FullGridCellList{<:DynamicVectorOfVectors{<:Any, <:CuArray{<:Any, <:Any, CUDA.UnifiedMemory}}}}
 
 function PointNeighbors.initialize_grid!(neighborhood_search::MultiGPUNeighborhoodSearch, y::AbstractMatrix)
     (; cell_list) = neighborhood_search
-    (; cells) = cell_list
-    (; vov, vov_per_gpu) = cells
 
-    for vov_ in vov_per_gpu
-        vov_.lengths .= 0
-    end
+    cell_list.cells.lengths .= 0
 
     if neighborhood_search.search_radius < eps()
         # Cannot initialize with zero search radius.
         # This is used in TrixiParticles when a neighborhood search is not used.
         return neighborhood_search
     end
 
-    # Fill cell lists per GPU
-    # TODO split range into chunks for each GPU
-    parallel_foreach2(axes(y, 2), y) do point, gpu
+    # Faster on a single GPU
+    @threaded CUDABackend() for point in axes(y, 2)
         # Get cell index of the point's cell
-        point_coords = extract_svector(y, Val(ndims(neighborhood_search)), point)
-        cell = cell_coords(point_coords, neighborhood_search)
+        point_coords = PointNeighbors.extract_svector(y, Val(ndims(neighborhood_search)), point)
+        cell = PointNeighbors.cell_coords(point_coords, neighborhood_search)
 
         # Add point to corresponding cell
-        @boundscheck check_cell_bounds(cell_list, cell)
-
-        # The atomic version of `pushat!` uses atomics to avoid race conditions when
-        # `pushat!` is used in a parallel loop.
-        @inbounds pushat_atomic!(vov_per_gpu[gpu], cell_index(cell_list, cell), point)
-    end
-
-    lengths = ntuple(gpu -> vov_per_gpu[gpu].lengths, length(vov_per_gpu))
-    CUDA.synchronize()
-    offsets_ = cu(cumsum(lengths), unified = true)
-    CUDA.synchronize()
-    vov.lengths .= offsets_[end]
-    CUDA.synchronize()
-    offsets = offsets_ .- lengths
-
-    # Merge cell lists
-    parallel_foreach2(axes(vov, 2), y) do cell, gpu
-        offset = offsets[gpu][cell]
-
-        points = vov_per_gpu[gpu][cell]
-        for i in eachindex(points)
-            vov.backend[offset + i, cell] = points[i]
-        end
+        pushat_atomic_system!(cell_list.cells, PointNeighbors.cell_index(cell_list, cell), point)
     end
 
     return neighborhood_search
 end
 
-function PointNeighbors.update!(neighborhood_search::MultiGPUNeighborhoodSearch,
-                                x::AbstractMatrix, y::AbstractMatrix;
-                                points_moving = (true, true), parallelization_backend = x)
-    # The coordinates of the first set of points are irrelevant for this NHS.
-    # Only update when the second set is moving.
-    points_moving[2] || return neighborhood_search
+# This might be slightly faster, but causes problems with synchronization in the interaction
+# kernel because we are carrying around device memory.
+# struct MultiGPUVectorOfVectors{T, VU, VD} <: AbstractVector{Array{T, 1}}
+#     vov_unified :: VU
+#     vov_device  :: VD
+# end
 
-    PointNeighbors.initialize_grid!(neighborhood_search, y)
-end
+# function MultiGPUVectorOfVectors(vov_unified, vov_device)
+#     MultiGPUVectorOfVectors{eltype(vov_unified.backend), typeof(vov_unified), typeof(vov_device)}(vov_unified, vov_device)
+# end
+
+# # Adapt.@adapt_structure(MultiGPUVectorOfVectors)
+
+# function Adapt.adapt_structure(to, vov::MultiGPUVectorOfVectors)
+#     @info "" to
+#     return MultiGPUVectorOfVectors(Adapt.adapt(to, vov.vov_unified), Adapt.adapt(to, vov.vov_device))
+# end
+
+# @propagate_inbounds function Base.getindex(vov::MultiGPUVectorOfVectors, i)
+#     return getindex(vov.vov_unified, i)
+# end
+
+# function Adapt.adapt_structure(to::CUDAMultiGPUBackend, vov::DynamicVectorOfVectors{T}) where {T}
+#     max_inner_length, max_outer_length = size(vov.backend)
+
+#     # Make sure the vector of vectors in device memory lives on the first GPU
+#     CUDA.device!(0)
+
+#     vov_unified = DynamicVectorOfVectors(Adapt.adapt(to, vov.backend), Adapt.adapt(to, vov.length_), Adapt.adapt(to, vov.lengths))
+#     vov_device = Adapt.adapt(CuArray, vov)
+
+#     return MultiGPUVectorOfVectors(vov_unified, vov_device)
+# end
+
+# const MultiGPUNeighborhoodSearch = GridNeighborhoodSearch{<:Any, <:Any, <:FullGridCellList{<:MultiGPUVectorOfVectors}}
+
+# function PointNeighbors.initialize_grid!(neighborhood_search::MultiGPUNeighborhoodSearch, y::AbstractMatrix)
+#     (; cell_list) = neighborhood_search
+#     (; cells) = cell_list
+#     (; vov_unified, vov_device) = cells
+
+#     if neighborhood_search.search_radius < eps()
+#         # Cannot initialize with zero search radius.
+#         # This is used in TrixiParticles when a neighborhood search is not used.
+#         return neighborhood_search
+#     end
+
+#     vov_device.lengths .= 0
+
+#     CUDA.device!(0)
+
+#     # Fill vector of vectors in device memory (on a single GPU)
+#     @threaded CUDABackend() for point in axes(y, 2)
+#         # Get cell index of the point's cell
+#         point_coords = extract_svector(y, Val(ndims(neighborhood_search)), point)
+#         cell = cell_coords(point_coords, neighborhood_search)
+
+#         # Add point to corresponding cell
+#         @boundscheck PointNeighbors.check_cell_bounds(cell_list, cell)
+
+#         # The atomic version of `pushat!` uses atomics to avoid race conditions when
+#         # `pushat!` is used in a parallel loop.
+#         @inbounds pushat_atomic!(vov_device, cell_index(cell_list, cell), point)
+#     end
+
+#     # Copy vector of vectors to unified memory
+#     vov_unified.backend .= vov_device.backend
+#     vov_unified.lengths .= vov_device.lengths
+#     vov_unified.length_[] = vov_device.length_[]
+
+#     return neighborhood_search
+# end
+
+# function PointNeighbors.update!(neighborhood_search::MultiGPUNeighborhoodSearch,
+#                                 x::AbstractMatrix, y::AbstractMatrix;
+#                                 points_moving = (true, true), parallelization_backend = x)
+#     # The coordinates of the first set of points are irrelevant for this NHS.
+#     # Only update when the second set is moving.
+#     points_moving[2] || return neighborhood_search
+
+#     PointNeighbors.initialize_grid!(neighborhood_search, y)
+# end
 
 end # module