fixup scan

Author: Claude

ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp

@@ -1,6 +1,8 @@
 #version 450
 
 #extension GL_EXT_control_flow_attributes : require
+#extension GL_KHR_shader_subgroup_basic : require
+#extension GL_KHR_shader_subgroup_arithmetic : require
 
 #include "types.glsl"
 
@@ -39,29 +41,6 @@ float softplus(float x) {
 }
 
 shared float stateC[SPLIT_H * D_STATE];
-shared float warp_sdata[D_STATE];
-float warp_reduce_sum(float val) {
-    const int tid = int(gl_LocalInvocationID.x);
-    int lane = tid % WARP_SIZE;
-    int warp_id = tid / WARP_SIZE;
-    int warp_offset = warp_id * WARP_SIZE;
-
-    warp_sdata[warp_offset + lane] = val;
-    barrier();
-
-    if (lane < 16) warp_sdata[warp_offset + lane] += warp_sdata[warp_offset + lane + 16];
-    barrier();
-    if (lane < 8) warp_sdata[warp_offset + lane] += warp_sdata[warp_offset + lane + 8];
-    barrier();
-    if (lane < 4) warp_sdata[warp_offset + lane] += warp_sdata[warp_offset + lane + 4];
-    barrier();
-    if (lane < 2) warp_sdata[warp_offset + lane] += warp_sdata[warp_offset + lane + 2];
-    barrier();
-    if (lane < 1) warp_sdata[warp_offset + lane] += warp_sdata[warp_offset + lane + 1];
-    barrier();
-
-    return warp_sdata[warp_offset];
-}
 
 void main() {
     const int tid = int(gl_LocalInvocationID.x);
@@ -128,9 +107,10 @@ void main() {
             if (idx < SPLIT_H * D_STATE) {
                 y = stateC[idx];
             }
-            y = warp_reduce_sum(y);
 
-            if (tid % WARP_SIZE == 0) {
+            y = subgroupAdd(y);
+
+            if (gl_SubgroupInvocationID == 0) {
                 const int k = tid / WARP_SIZE + j * (D_STATE / WARP_SIZE);
                 d[y_base_idx + uint(i) * uint(stride_y) + uint(k)] = y;
             }