optimized mapreduce using sub group shuffle

ref #352

Unfortunately, I don't really see any performance improvements with
this, any ideas why? I expected this to be quite a bit faster.

Author: simeonschaub

lib/intrinsics/src/work_item.jl

@@ -34,6 +34,50 @@ for (julia_name, (spirv_name, julia_type, offset)) in [
     end
 end
 
+
+# Sub-group shuffle intrinsics using a loop and @eval, matching the style of the 1D/3D value loops above
+export sub_group_shuffle, sub_group_shuffle_xor
+
+for (jltype, llvmtype, julia_type_str) in [
+        (Int8,    "i8",    :Int8),
+        (UInt8,   "i8",    :UInt8),
+        (Int16,   "i16",   :Int16),
+        (UInt16,  "i16",   :UInt16),
+        (Int32,   "i32",   :Int32),
+        (UInt32,  "i32",   :UInt32),
+        (Int64,   "i64",   :Int64),
+        (UInt64,  "i64",   :UInt64),
+        (Float16, "half",  :Float16),
+        (Float32, "float", :Float32),
+        (Float64, "double",:Float64)
+    ]
+    @eval begin
+        export sub_group_shuffle, sub_group_shuffle_xor
+        function sub_group_shuffle(x::$jltype, idx::Integer)
+            Base.llvmcall(
+                $("""
+                declare $llvmtype @__spirv_GroupNonUniformShuffle(i32, $llvmtype, i32)
+                define $llvmtype @entry($llvmtype %val, i32 %idx) #0 {
+                    %res = call $llvmtype @__spirv_GroupNonUniformShuffle(i32 3, $llvmtype %val, i32 %idx)
+                    ret $llvmtype %res
+                }
+                attributes #0 = { alwaysinline }
+                """, "entry"), $julia_type_str, Tuple{$julia_type_str, Int32}, x, idx % Int32 - 1i32)
+        end
+        function sub_group_shuffle_xor(x::$jltype, mask::Integer)
+            Base.llvmcall(
+                $("""
+                declare $llvmtype @__spirv_GroupNonUniformShuffleXor(i32, $llvmtype, i32)
+                define $llvmtype @entry($llvmtype %val, i32 %mask) #0 {
+                    %res = call $llvmtype @__spirv_GroupNonUniformShuffleXor(i32 3, $llvmtype %val, i32 %mask)
+                    ret $llvmtype %res
+                }
+                attributes #0 = { alwaysinline }
+                """, "entry"), $julia_type_str, Tuple{$julia_type_str, Int32}, x, mask % UInt32)
+        end
+    end
+end
+
 # 3D values
 for (julia_name, (spirv_name, offset)) in [
         # indices

src/mapreduce.jl

@@ -5,8 +5,78 @@
 # - group-stride loop to delay need for second kernel launch
 # - let the driver choose the local size
 
+function shuffle_expr(::Type{T}) where {T}
+    if T in SPIRVIntrinsics.generic_integer_types || T in SPIRVIntrinsics.generic_types
+        return :(sub_group_shuffle(val, i))
+    elseif Base.isstructtype(T)
+        ex = Expr(:new, T)
+        for f in fieldnames(T)
+            ex_f = shuffle_expr(fieldtype(T, f))
+            ex_f === nothing && return nothing
+            push!(ex.args, :(let val = getfield(val, $(QuoteNode(f)))
+                $ex_f
+            end))
+        end
+        return ex
+    else
+        return nothing
+    end
+end
+
+@inline @generated function reduce_group(op, val::T, neutral, ::Val{maxitems}) where {T, maxitems}
+    ex = shuffle_expr(T)
+    if ex === nothing
+        return :(reduce_group_fallback(op, val, neutral, Val(maxitems)))
+    end
+
+    quote
+        # Subgroup shuffle-based warp reduction
+        lane = get_sub_group_local_id()
+        width = get_sub_group_size()
+
+        offset = 1
+        while offset < width
+            if lane > offset
+                i = lane - offset
+                other = $ex
+                val = op(val, other)
+            end
+            offset <<= 1
+        end
+
+        items = get_num_sub_groups()
+        item = get_sub_group_id()
+
+        shared = CLLocalArray(T, (maxitems,))
+        if items > 1 && lane == 1
+            @inbounds shared[item] = val
+
+            d = 1
+            while d < items
+                work_group_barrier(LOCAL_MEM_FENCE)
+                index = 2 * d * (item-1) + 1
+                @inbounds if index <= items
+                    other_val = if index + d <= items
+                        shared[index+d]
+                    else
+                        neutral
+                    end
+                    shared[index] = op(shared[index], other_val)
+                end
+                d *= 2
+            end
+
+            if item == 1
+                val = @inbounds shared[item]
+            end
+        end
+
+        return val
+    end
+end
+
 # Reduce a value across a group, using local memory for communication
-@inline function reduce_group(op, val::T, neutral, ::Val{maxitems}) where {T, maxitems}
+@inline function reduce_group_fallback(op, val::T, neutral, ::Val{maxitems}) where {T, maxitems}
     items = get_local_size()
     item = get_local_id()