shfl_down intrinsics

Co-Authored-By: Anton Smirnov <tonysmn97@gmail.com>

Author: christiangnrd

src/intrinsics.jl

@@ -188,6 +188,32 @@ Declare memory that is local to a workgroup.
 """
 localmemory(::Type{T}, dims) where {T} = localmemory(T, Val(dims))
 
+"""
+    shfl_down(val::T, offset::Integer) where T
+
+Read `val` from a lane with higher id given by `offset`.
+
+!!! note
+    Backend implementations **must** implement:
+    ```
+    @device_override shfl_down(val::T, offset::Integer) where T
+    ```
+    As well as the on-device functionality.
+"""
+function shfl_down end
+
+"""
+    shfl_down_types(::Backend)::Vector{DataType}
+
+Returns a vector of `DataType`s supported on `backend`
+
+!!! note
+    Backend implementations **must** implement this function
+    only if they support `shfl_down` for any types.
+"""
+shfl_down_types(::Backend) = DataType[]
+
+
 """
     barrier()
 

test/intrinsics.jl

@@ -45,6 +45,25 @@ function test_subgroup_kernel(results)
     return
 end
 
+function shfl_down_test_kernel(a, b, ::Val{N}) where {N}
+    idx = KI.get_sub_group_local_id()
+
+    val = a[idx]
+
+    offset = 0x00000001
+    while offset < N
+        val += KI.shfl_down(val, offset)
+        offset <<= 1
+    end
+
+    KI.sub_group_barrier()
+
+    if idx == 1
+        b[idx] = val
+    end
+    return
+end
+
 function intrinsics_testsuite(backend, AT)
     @testset "KernelIntrinsics Tests" begin
         @testset "Launch parameters" begin
@@ -177,6 +196,23 @@ function intrinsics_testsuite(backend, AT)
                 @test sg_data.sub_group_local_id == expected_sg_local
             end
         end
+        @testset "shfl_down" begin
+            @test !isempty(KI.shfl_down_types(backend()))
+            types_to_test = setdiff(KI.shfl_down_types(backend()), [Bool])
+            @testset "$T" for T in types_to_test
+                N = KI.sub_group_size(backend())
+                a = zeros(T, N)
+                rand!(a, (0:1))
+
+                dev_a = AT(a)
+                dev_b = AT(zeros(T, N))
+
+                KI.@kernel backend() workgroupsize = N shfl_down_test_kernel(dev_a, dev_b, Val(N))
+
+                b = Array(dev_b)
+                @test sum(a) ≈ b[1]
+            end
+        end
     end
     return nothing
 end