Fix histogram implementation

The final part of the loop expects every thread to exists, so we cannot not launch them. Avoid work on extra threads until then.

Also use Int32 since some backends lack Int64 atomics, and make one of the tests have weird groupsize since that's when the errors used to pop up.

Author: christiangnrd

examples/histogram.jl

@@ -5,31 +5,29 @@ include(joinpath(dirname(pathof(KernelAbstractions)), "../examples/utils.jl")) #
 
 # Function to use as a baseline for CPU metrics
 function create_histogram(input)
-    histogram_output = zeros(Int, maximum(input))
+    histogram_output = zeros(eltype(input), maximum(input))
     for i in input
         histogram_output[i] += 1
     end
     return histogram_output
 end
 
 # This a 1D histogram kernel where the histogramming happens on shmem
-@kernel function histogram_kernel!(histogram_output, input)
-    tid = @index(Global, Linear)
+@kernel unsafe_indices = true function histogram_kernel!(histogram_output, input)
+    gid = @index(Group, Linear)
     lid = @index(Local, Linear)
 
-    @uniform warpsize = Int(32)
-
-    @uniform gs = @groupsize()[1]
+    @uniform gs = prod(@groupsize())
+    tid = (gid - 1) * gs + lid
     @uniform N = length(histogram_output)
 
-    shared_histogram = @localmem Int (gs)
+    shared_histogram = @localmem eltype(input) (gs)
 
     # This will go through all input elements and assign them to a location in
     # shmem. Note that if there is not enough shem, we create different shmem
     # blocks to write to. For example, if shmem is of size 256, but it's
     # possible to get a value of 312, then we will have 2 separate shmem blocks,
     # one from 1->256, and another from 256->512
-    @uniform max_element = 1
     for min_element in 1:gs:N
 
         # Setting shared_histogram to 0
@@ -42,7 +40,7 @@ end
         end
 
         # Defining bin on shared memory and writing to it if possible
-        bin = input[tid]
+        bin = tid <= length(input) ? input[tid] : 0
         if bin >= min_element && bin < max_element
             bin -= min_element - 1
             @atomic shared_histogram[bin] += 1
@@ -58,10 +56,10 @@ end
 
 end
 
-function histogram!(histogram_output, input)
+function histogram!(histogram_output, input, groupsize = 256)
     backend = get_backend(histogram_output)
     # Need static block size
-    kernel! = histogram_kernel!(backend, (256,))
+    kernel! = histogram_kernel!(backend, (groupsize,))
     kernel!(histogram_output, input, ndrange = size(input))
     return
 end
@@ -74,9 +72,10 @@ function move(backend, input)
 end
 
 @testset "histogram tests" begin
-    rand_input = [rand(1:128) for i in 1:1000]
-    linear_input = [i for i in 1:1024]
-    all_two = [2 for i in 1:512]
+    # Use Int32 as some backends don't support 64-bit atomics
+    rand_input = Int32.(rand(1:128, 1000))
+    linear_input = Int32.(1:1024)
+    all_two = fill(Int32(2), 512)
 
     histogram_rand_baseline = create_histogram(rand_input)
     histogram_linear_baseline = create_histogram(linear_input)
@@ -86,14 +85,14 @@ end
     linear_input = move(backend, linear_input)
     all_two = move(backend, all_two)
 
-    rand_histogram = KernelAbstractions.zeros(backend, Int, 128)
-    linear_histogram = KernelAbstractions.zeros(backend, Int, 1024)
-    two_histogram = KernelAbstractions.zeros(backend, Int, 2)
+    rand_histogram = KernelAbstractions.zeros(backend, eltype(rand_input), Int(maximum(rand_input)))
+    linear_histogram = KernelAbstractions.zeros(backend, eltype(linear_input), Int(maximum(linear_input)))
+    two_histogram = KernelAbstractions.zeros(backend, eltype(all_two), Int(maximum(all_two)))
 
-    histogram!(rand_histogram, rand_input)
+    histogram!(rand_histogram, rand_input, 6)
     histogram!(linear_histogram, linear_input)
     histogram!(two_histogram, all_two)
-    KernelAbstractions.synchronize(CPU())
+    KernelAbstractions.synchronize(backend)
 
     @test isapprox(Array(rand_histogram), histogram_rand_baseline)
     @test isapprox(Array(linear_histogram), histogram_linear_baseline)