new indexing approach

Author: yuehhua

lib/NNlibCUDA/src/scatter.jl

@@ -2,37 +2,43 @@ ATM_OPS = Dict((+) => CUDA.atomic_add!, (-) => CUDA.atomic_sub!, (max) => CUDA.a
                (*) => CUDA.atomic_mul!, (/) => CUDA.atomic_div!, (&) => CUDA.atomic_and!, (|) => CUDA.atomic_or!)
 
 for (op, atm_op) in ATM_OPS
-    @eval function scatter_kernel!(op::typeof($(op)), dst, src, idx, dims::Val{0}, pre_rng, post_rng)
-        li = threadIdx().x + (blockIdx().x - 1) * blockDim().x
+    @eval function scatter_kernel!(op::typeof($(op)), dst, src, idx)
+        index = threadIdx().x + (blockIdx().x - 1) * blockDim().x
 
-        @inbounds if li <= post_rng
-            ind = CartesianIndices(idx)[li]
-            dst_i = Base._to_linear_index(dst, idx[li]...)
-            $(atm_op)(pointer(dst, dst_i), src[ind])
+        @inbounds if index <= length(idx)
+            i = Base._to_linear_index(dst, idx[index]...)
+            $(atm_op)(pointer(dst, i), src[index])
         end
         return nothing
     end
 
-    @eval function scatter_kernel!(op::typeof($(op)), dst, src, idx, dims::Val{1}, pre_rng, post_rng)
-        li = threadIdx().y + (blockIdx().y - 1) * blockDim().y
-        i = threadIdx().x + (blockIdx().x - 1) * blockDim().x
+    @eval function scatter_kernel!(op::typeof($(op)), dst, src, idx, dims::Val{N}, max_idx, max_dims_idx, dims_size) where {N}
+        index = threadIdx().x + (blockIdx().x - 1) * blockDim().x
 
-        @inbounds if li <= post_rng && i <= pre_rng
-            j = CartesianIndices(idx)[li]
-            dst_i = Base._to_linear_index(dst, i, idx[li]...)
-            $(atm_op)(pointer(dst, dst_i), src[i, j])
+        @inbounds if index <= max_idx
+            j, k = divrem(index, max_dims_idx)
+            dims_i = CartesianIndices(dims_size)[k]
+            i = Base._to_linear_index(dst, Tuple(dims_i)..., idx[j]...)
+            $(atm_op)(pointer(dst, i), src[index])
         end
         return nothing
     end
 
     @eval function NNlib.scatter!(op::typeof($(op)), dst::CuArray{Tdst}, src::CuArray{Tsrc}, idx::CuArray{<:IntOrIntTuple}, dims::Val{N}) where {Tdst,Tsrc,N}
-        pre_rng = prod(size(dst)[1:N])
-        post_rng = length(idx)
-        thread_x = min(MAX_THREADS, pre_rng)
-        thread_y = min(MAX_THREADS ÷ thread_x, post_rng)
-        threads = (thread_x, thread_y)
-        blocks = ceil.(Int, (pre_rng, post_rng) ./ threads)
-        @cuda blocks=blocks threads=threads scatter_kernel!(op, dst, src, idx, dims, pre_rng, post_rng)
-        return dst
+        if N == 0
+            max_idx = length(idx)
+            threads = min(MAX_THREADS, max_idx)
+            blocks = ceil(Int, max_idx / threads)
+            @cuda blocks=blocks threads=threads scatter_kernel!(op, dst, src, idx)
+            return dst
+        else
+            dims_size = size(dst)[1:N]
+            max_dims_idx = prod(dims_size)
+            max_idx = max_dims_idx * length(idx)
+            threads = min(MAX_THREADS, max_idx)
+            blocks = ceil(Int, max_idx / threads)
+            @cuda blocks=blocks threads=threads scatter_kernel!(op, dst, src, idx, dims, max_idx, max_dims_idx, dims_size)
+            return dst
+        end
     end
 end