fix_1

Author: Claude

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

@@ -40,26 +40,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();
-
-    int offset = WARP_SIZE / 2;
-    while (offset > 0) {
-        if (lane < offset) {
-            warp_sdata[warp_offset + lane] += warp_sdata[warp_offset + lane + offset];
-        }
-        barrier();
-        offset >>= 1;
-    }
-
-    return warp_sdata[warp_offset];
-}
 
 void main() {
     const int tid = int(gl_LocalInvocationID.x);
@@ -126,7 +106,22 @@ void main() {
             if (idx < SPLIT_H * D_STATE) {
                 y = stateC[idx];
             }
-            y = warp_reduce_sum(y);
+
+            int lane = tid % WARP_SIZE;
+            int warp_id = tid / WARP_SIZE;
+            int warp_offset = warp_id * WARP_SIZE;
+
+            warp_sdata[warp_offset + lane] = y;
+            barrier();
+
+            [[unroll]] for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
+                if (lane < offset) {
+                    warp_sdata[warp_offset + lane] += warp_sdata[warp_offset + lane + offset];
+                }
+                barrier();
+            }
+
+            y = warp_sdata[warp_offset];
 
             if (tid % WARP_SIZE == 0) {
                 const int k = tid / WARP_SIZE + j * (D_STATE / WARP_SIZE);