1

Author: Claude

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

@@ -49,16 +49,14 @@ float warp_reduce_sum(float val) {
     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();
+    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];
 }