From be5295bd51edf5c32fa98ee45d47f511e8b3072f Mon Sep 17 00:00:00 2001
From: Eve <139727413+netrunnereve@users.noreply.github.com>
Date: Thu, 21 Nov 2024 11:27:18 -0500
Subject: [PATCH] test

---
 .../vulkan-shaders/mul_mat_vec.comp           | 21 +++++++------------
 1 file changed, 7 insertions(+), 14 deletions(-)

diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp
index d3ccba7fc..e8209c511 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp
@@ -1,5 +1,7 @@
 #version 450
 
+#extension GL_KHR_shader_subgroup_arithmetic: enable
+
 #ifdef FLOAT16
 #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
 #endif
@@ -10,8 +12,6 @@ layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
 layout (constant_id = 0) const uint BLOCK_SIZE = 32;
 
-shared FLOAT_TYPE tmp[BLOCK_SIZE];
-
 void main() {
     const uint row = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z;
     const uint tid = gl_LocalInvocationID.x;
@@ -28,7 +28,7 @@ void main() {
 
     const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
 
-    tmp[tid] = FLOAT_TYPE(0.0f);
+    FLOAT_TYPE tmp = FLOAT_TYPE(0.0f);
 
     [[unroll]] for (uint i = 0; i < p.ncols/block_size; i += 2) {
         const uint col = i*block_size + 2*tid;
@@ -39,18 +39,11 @@ void main() {
         vec2 v = dequantize(ib, iqs, a_offset / QUANT_K);
 
         // matrix multiplication
-        tmp[tid] = fma(FLOAT_TYPE(v.x), FLOAT_TYPE(data_b[b_offset + iybs + iqs]), fma(FLOAT_TYPE(v.y), FLOAT_TYPE(data_b[b_offset + iybs + iqs + y_offset]), tmp[tid]));
+        tmp = fma(FLOAT_TYPE(v.x), FLOAT_TYPE(data_b[b_offset + iybs + iqs]), fma(FLOAT_TYPE(v.y), FLOAT_TYPE(data_b[b_offset + iybs + iqs + y_offset]), tmp));
     }
 
     // sum up partial sums and write back result
-    barrier();
-    [[unroll]] for (uint s = block_size/2; s > 0; s >>= 1) {
-        if (tid < s) {
-            tmp[tid] += tmp[tid + s];
-        }
-        barrier();
-    }
-    if (tid == 0) {
-        data_d[d_offset + row] = D_TYPE(tmp[0]);
-    }
+    tmp = subgroupAdd(tmp);
+    if (tid == 0)
+        data_d[d_offset + row] = D_TYPE(tmp);
 }