Optimize Radix Sort.

Project.toml

@@ -5,6 +5,7 @@ version = "0.4.3"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
@@ -18,6 +19,7 @@ AcceleratedKernelsoneAPIExt = "oneAPI"
 
 [compat]
 ArgCheck = "2"
+Atomix = "0.1, 1"
 GPUArraysCore = "0.2.0"
 KernelAbstractions = "0.9.34, 0.10"
 Markdown = "1"

src/accumulate/accumulate_1d_gpu.jl

@@ -22,6 +22,8 @@ end
     len = length(v)
     @uniform block_size = @groupsize()[1]
     temp = @localmem eltype(v) (0x2 * block_size + conflict_free_offset(0x2 * block_size),)
+    s_lastval = @localmem eltype(v) (1,)     # holds v[len]; captured at load time so
+                                             # the inclusive shift needn't re-read v
 
     # NOTE: for many index calculations in this library, computation using zero-indexing leads to
     # fewer operations (also code is transpiled to CUDA / ROCm / oneAPI / Metal code which do zero
@@ -42,16 +44,21 @@ end
     bank_offset_a = conflict_free_offset(ai)
     bank_offset_b = conflict_free_offset(bi)
 
-    if block_offset + ai < len
-        temp[ai + bank_offset_a + 0x1] = v[block_offset + ai + 0x1]
-    else
-        temp[ai + bank_offset_a + 0x1] = neutral
+    # Load both elements, keeping them in registers.  Re-reading global `v` later
+    # from inside the single-thread inclusive-shift branch is a divergent global
+    # read that faults or hangs on some backends, so everything the shift needs is
+    # captured here: each thread's own last element (my_bi_val) and, in a shared
+    # slot, the array's final element v[len] (written by whichever thread owns it).
+    my_ai_val = (block_offset + ai < len) ? v[block_offset + ai + 0x1] : neutral
+    temp[ai + bank_offset_a + 0x1] = my_ai_val
+    my_bi_val = (block_offset + bi < len) ? v[block_offset + bi + 0x1] : neutral
+    temp[bi + bank_offset_b + 0x1] = my_bi_val
+
+    if block_offset + ai == len - 0x1
+        s_lastval[0x1] = my_ai_val
     end
-
-    if block_offset + bi < len
-        temp[bi + bank_offset_b + 0x1] = v[block_offset + bi + 0x1]
-    else
-        temp[bi + bank_offset_b + 0x1] = neutral
+    if block_offset + bi == len - 0x1
+        s_lastval[0x1] = my_bi_val
     end
 
     # Build block reduction down
@@ -102,26 +109,33 @@ end
 
     # For exclusive ScanPrefixes scans, the local exclusive output is already correct.
     # DecoupledLookback still shifts non-first blocks because _accumulate_previous!
-    # expects each block's last value to be globally inclusive.
-    if inclusive || (iblock != 0x0 && !isnothing(flags))
-        # To compute an inclusive scan, shift elements left...
-        @synchronize()
-        t1 = temp[ai + bank_offset_a + 0x1]
-        t2 = temp[bi + bank_offset_b + 0x1]
-        @synchronize()
+    # expects each block's last value to be globally inclusive.  The condition is
+    # uniform across the group, but it is a runtime value: the @synchronize() barriers
+    # of the shift must stay OUTSIDE the branch.  A barrier inside a conditional the
+    # compiler cannot prove uniform deadlocks POCL/SPIR-V (its work-item loop does not
+    # see the barrier reached on every path), so only the shared-memory writes are
+    # guarded while the barriers always execute.
+    do_shift = inclusive || (iblock != 0x0 && !isnothing(flags))
+
+    # To compute an inclusive scan, shift elements left...
+    @synchronize()
+    t1 = temp[ai + bank_offset_a + 0x1]
+    t2 = temp[bi + bank_offset_b + 0x1]
+    @synchronize()
 
+    if do_shift
         if ai > 0x0
             temp[ai - 0x1 + conflict_free_offset(ai - 0x1) + 0x1] = t1
         end
         temp[bi - 0x1 + conflict_free_offset(bi - 0x1) + 0x1] = t2
 
-        # ...and accumulate the last value too
+        # ...and accumulate the last value too, without re-reading global v.  A
+        # full non-last block's last element is this thread's own my_bi_val
+        # (== v[(iblock+1)*block_size*2]); the final block instead needs the
+        # array's last element v[len], captured in s_lastval at load time.
         if bi == 0x2 * block_size - 0x1
-            if iblock < num_blocks - 0x1
-                temp[bi + bank_offset_b + 0x1] = op(t2, v[(iblock + 0x1) * block_size * 0x2])
-            else
-                temp[bi + bank_offset_b + 0x1] = op(t2, v[len])
-            end
+            last_val = iblock < num_blocks - 0x1 ? my_bi_val : s_lastval[0x1]
+            temp[bi + bank_offset_b + 0x1] = op(t2, last_val)
         end
     end