Vulkan: Improve Q8 OUT_PROD performance

Increase OUT_PROD Q8 performance through improving memory locality.

Author: Italo Nicola

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

@@ -4,7 +4,8 @@
 #include "generic_binary_head.comp"
 #include "dequant_funcs.comp"
 
-const uint num_threads = 256;
+const uint quant_group_sz = 2;
+const uint num_threads = 512 / quant_group_sz;
 layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in;
 
 void get_dst_indices(uint idx, out uint i20, out uint i21, out uint i22, out uint i23) {
@@ -17,38 +18,47 @@ void get_dst_indices(uint idx, out uint i20, out uint i21, out uint i22, out uin
 }
 
 void main() {
-    // num_threads * num_iter must equal 512 to match the wg_denoms and get_idx
-    const uint num_iter = 2;
-
     const uint broadcast2 = uint(p.param2);
     const uint broadcast3 = p.param3;
 
-    uint idx = get_idx();
+    uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x * quant_group_sz;
+
+    uint aoffset = get_aoffset();
+    uint boffset = get_boffset();
+    uint doffset = get_doffset();
 
-    [[unroll]] for (uint it = 0; it < num_iter; ++it) {
-        if (idx < p.ne) {
-            uint i0, i1, i2, i3;
-            get_dst_indices(idx, i0, i1, i2, i3);
+    if (idx < p.ne) {
+        uint i0, i1, i2, i3;
+        get_dst_indices(idx, i0, i1, i2, i3);
 
-            float acc = 0.0f;
+        vec2 acc = vec2(0.0);
+
+        for (uint k = 0; k < p.ne01; k++) {
+            if (i0 + 1 >= p.ne20) { // XXX
+                continue;
+            }
 
-            for (uint k = 0; k < p.ne01; k += 1) {
-                const uint a_block_base = get_aoffset() + (i3 / broadcast3) * p.nb03 + (i2 / broadcast2) * p.nb02 + k * p.nb01;
-                const uint ib = a_block_base + (i0 / QUANT_K) * p.nb00;
-                const uint iqs = (i0 % QUANT_K) / QUANT_R;
+            const uint a_block_base = aoffset + (i3 / broadcast3) * p.nb03 + (i2 / broadcast2) * p.nb02 + k * p.nb01;
+            const uint ib = a_block_base + ((i0) / QUANT_K) * p.nb00;
+            const uint iqs = ((i0) % QUANT_K) / QUANT_R;
 
-                const vec2 v = dequantize(ib, iqs, 0);
-                const vec2 dm = get_dm(ib, 0);
-                const float a_val = v.x * dm.x + dm.y;
+            const vec2 v = dequantize(ib, iqs, 0);
+            const vec2 dm = get_dm(ib, 0);
+            const vec2 a_vals = v * dm.x + dm.y;
+
+            const uint b_idx = src1_idx(i1, k, i2, i3);
+            const float b = data_b[boffset + b_idx];
+
+            acc += a_vals * b;
+        }
 
-                const uint b_idx = src1_idx(i1, k, i2, i3);
-                const float b = data_b[get_boffset() + b_idx];
-                acc += a_val * b;
+        uint d_idx = dst_idx(i0, i1, i2, i3);
+        for (uint q = 0; q < quant_group_sz; q++) {
+            if (d_idx + q + 0 >= p.ne) { // XXX
+                continue;
             }
 
-            uint d_idx = dst_idx(i0, i1, i2, i3);
-            data_d[get_doffset() + d_idx] = acc;
+            data_d[doffset + d_idx + q] = acc[q];
         }
-        idx += num_threads;
     }
 }