Repository: ggml-org/llama.cpp
Pull request: vulkan: Add ACC_TYPE_VEC2 implementation

This PR adds the implementation for `ACC_TYPE_VEC2`. This change, with non-`coopmat` shaders, using `ACC_TYPE_VEC2` improves caching behavior, as accessing 32-bit values is generally more efficient than accessing 16-bit values.

<details>
<summary><strong>Performance Comparison (Without <code>coopmat</code> and <code>coopmat2</code>) NVIDIA GeForce RTX 4060 Ti </strong></summary>

| Name | Before (us/run) | After (us/run) | Δ% (Improvement) |
|-------------|------------------|----------------|------------------|
| MUL_MAT(type_a=f32,type_b=f32,m=4096,n=512,k=14336) | 5767.64 | 5479.83 | +4.99% |
| MUL_MAT(type_a=f16,type_b=f32,m=4096,n=512,k=14336) | 5421.40 | 5047.91 | +6.88% |
| MUL_MAT(type_a=bf16,type_b=f32,m=4096,n=512,k=14336) | 5281.02 | 6002.14 | −13.66% |
| MUL_MAT(type_a=q4_0,type_b=f32,m=4096,n=512,k=14336) | 2741.43 | 2748.71 | −0.27% |
| MUL_MAT(type_a=q4_1,type_b=f32,m=4096,n=512,k=14336) | 2766.60 | 2764.23 | +0.09% |
| MUL_MAT(type_a=q5_0,type_b=f32,m=4096,n=512,k=14336) | 2877.49 | 2875.25 | +0.08% |
| MUL_MAT(type_a=q5_1,type_b=f32,m=4096,n=512,k=14336) | 2869.17 | 2867.33 | +0.06% |
| MUL_MAT(type_a=q8_0,type_b=f32,m=4096,n=512,k=14336) | 2887.17 | 2890.27 | −0.11% |
| MUL_MAT(type_a=mxfp4,type_b=f32,m=4096,n=512,k=14336) | 4976.57 | 4043.75 | +18.75% |
| MUL_MAT(type_a=q2_K,type_b=f32,m=4096,n=512,k=14336) | 4938.25 | 4120.32 | +16.56% |
| MUL_MAT(type_a=q3_K,type_b=f32,m=4096,n=512,k=14336) | 5287.85 | 4548.30 | +13.99% |
| MUL_MAT(type_a=q4_K,type_b=f32,m=4096,n=512,k=14336) | 5373.34 | 4566.63 | +15.01% |
| MUL_MAT(type_a=q5_K,type_b=f32,m=4096,n=512,k=14336) | 5769.13 | 4907.47 | +14.94% |
| MUL_MAT(type_a=q6_K,type_b=f32,m=4096,n=512,k=14336) | 5507.98 | 4524.96 | +17.85% |
| MUL_MAT(type_a=iq2_xxs,type_b=f32,m=4096,n=512,k=14336) | 4877.02 | 4043.75 | +17.07% |
| MUL_MAT(type_a=iq2_xs,type_b=f32,m=4096,n=512,k=14336) | 5010.98 | 4112.35 | +17.94% |
| MUL_MAT(type_a=iq2_s,type_b=f32,m=4096,n=512,k=14336) | 4863.99 | 4065.67 | +16.41% |
| MUL_MAT(type_a=iq3_xxs,type_b=f32,m=4096,n=512,k=14336) | 4957.83 | 4129.54 | +16.70% |
| MUL_MAT(type_a=iq1_s,type_b=f32,m=4096,n=512,k=14336) | 4583.30 | 3788.42 | +17.34% |
| MUL_MAT(type_a=iq1_m,type_b=f32,m=4096,n=512,k=14336) | 5128.29 | 4280.64 | +16.52% |
| MUL_MAT(type_a=iq4_nl,type_b=f32,m=4096,n=512,k=14336) | 4885.91 | 3992.67 | +18.27% |
| MUL_MAT(type_a=iq3_s,type_b=f32,m=4096,n=512,k=14336) | 4933.56 | 4084.30 | +17.22% |
| MUL_MAT(type_a=iq4_xs,type_b=f32,m=4096,n=512,k=14336) | 5389.60 | 4489.23 | +16.67% |

</details>

<details>
<summary><strong>Performance before(Without <code>coopmat</code> and <code>coopmat2</code>) NVIDIA GeForce RTX 4060 Ti </strong></summary>

```text
MUL_MAT(type_a=f32,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                   174 runs -  5767.64 us/run -  60.13 GFLOP/run -  10.43 TFLOPS
MUL_MAT(type_a=f16,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                   186 runs -  5421.40 us/run -  60.13 GFLOP/run -  11.09 TFLOPS
MUL_MAT(type_a=bf16,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  190 runs -  5281.02 us/run -  60.13 GFLOP/run -  11.39 TFLOPS
MUL_MAT(type_a=q4_0,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  366 runs -  2741.43 us/run -  60.13 GFLOP/run -  21.93 TFLOPS
MUL_MAT(type_a=q4_1,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  362 runs -  2766.60 us/run -  60.13 GFLOP/run -  21.73 TFLOPS
MUL_MAT(type_a=q5_0,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  348 runs -  2877.49 us/run -  60.13 GFLOP/run -  20.90 TFLOPS
MUL_MAT(type_a=q5_1,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  350 runs -  2869.17 us/run -  60.13 GFLOP/run -  20.96 TFLOPS
MUL_MAT(type_a=q8_0,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  348 runs -  2887.17 us/run -  60.13 GFLOP/run -  20.83 TFLOPS
MUL_MAT(type_a=mxfp4,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 202 runs -  4976.57 us/run -  60.13 GFLOP/run -  12.08 TFLOPS
MUL_MAT(type_a=q2_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  204 runs -  4938.25 us/run -  60.13 GFLOP/run -  12.18 TFLOPS
MUL_MAT(type_a=q3_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  190 runs -  5287.85 us/run -  60.13 GFLOP/run -  11.37 TFLOPS
MUL_MAT(type_a=q4_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  188 runs -  5373.34 us/run -  60.13 GFLOP/run -  11.19 TFLOPS
MUL_MAT(type_a=q5_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  174 runs -  5769.13 us/run -  60.13 GFLOP/run -  10.42 TFLOPS
MUL_MAT(type_a=q6_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  182 runs -  5507.98 us/run -  60.13 GFLOP/run -  10.92 TFLOPS
MUL_MAT(type_a=iq2_xxs,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):               206 runs -  4877.02 us/run -  60.13 GFLOP/run -  12.33 TFLOPS
MUL_MAT(type_a=iq2_xs,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                200 runs -  5010.98 us/run -  60.13 GFLOP/run -  12.00 TFLOPS
MUL_MAT(type_a=iq2_s,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 206 runs -  4863.99 us/run -  60.13 GFLOP/run -  12.36 TFLOPS
MUL_MAT(type_a=iq3_xxs,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):               202 runs -  4957.83 us/run -  60.13 GFLOP/run -  12.13 TFLOPS
MUL_MAT(type_a=iq1_s,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 220 runs -  4583.30 us/run -  60.13 GFLOP/run -  13.12 TFLOPS
MUL_MAT(type_a=iq1_m,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 196 runs -  5128.29 us/run -  60.13 GFLOP/run -  11.73 TFLOPS
MUL_MAT(type_a=iq4_nl,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                206 runs -  4885.91 us/run -  60.13 GFLOP/run -  12.31 TFLOPS
MUL_MAT(type_a=iq3_s,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 204 runs -  4933.56 us/run -  60.13 GFLOP/run -  12.19 TFLOPS
MUL_MAT(type_a=iq4_xs,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                186 runs -  5389.60 us/run -  60.13 GFLOP/run -  11.16 TFLOPS

```

</details>

<details>
<summary><strong>Performance after(Without <code>coopmat</code> and <code>coopmat2</code>) NVIDIA GeForce RTX 4060 Ti </strong></summary>

```text
MUL_MAT(type_a=f32,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                   184 runs -  5479.83 us/run -  60.13 GFLOP/run -  10.97 TFLOPS
MUL_MAT(type_a=f16,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                   200 runs -  5047.91 us/run -  60.13 GFLOP/run -  11.91 TFLOPS
MUL_MAT(type_a=bf16,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  168 runs -  6002.14 us/run -  60.13 GFLOP/run -  10.02 TFLOPS
MUL_MAT(type_a=q4_0,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  364 runs -  2748.71 us/run -  60.13 GFLOP/run -  21.88 TFLOPS
MUL_MAT(type_a=q4_1,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  362 runs -  2764.23 us/run -  60.13 GFLOP/run -  21.75 TFLOPS
MUL_MAT(type_a=q5_0,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  348 runs -  2875.25 us/run -  60.13 GFLOP/run -  20.91 TFLOPS
MUL_MAT(type_a=q5_1,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  350 runs -  2867.33 us/run -  60.13 GFLOP/run -  20.97 TFLOPS
MUL_MAT(type_a=q8_0,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  348 runs -  2890.27 us/run -  60.13 GFLOP/run -  20.80 TFLOPS
MUL_MAT(type_a=mxfp4,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 248 runs -  4043.75 us/run -  60.13 GFLOP/run -  14.87 TFLOPS
MUL_MAT(type_a=q2_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  244 runs -  4120.32 us/run -  60.13 GFLOP/run -  14.59 TFLOPS
MUL_MAT(type_a=q3_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  220 runs -  4548.30 us/run -  60.13 GFLOP/run -  13.22 TFLOPS
MUL_MAT(type_a=q4_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  220 runs -  4566.63 us/run -  60.13 GFLOP/run -  13.17 TFLOPS
MUL_MAT(type_a=q5_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  204 runs -  4907.47 us/run -  60.13 GFLOP/run -  12.25 TFLOPS
MUL_MAT(type_a=q6_K,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                  222 runs -  4524.96 us/run -  60.13 GFLOP/run -  13.29 TFLOPS
MUL_MAT(type_a=iq2_xxs,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):               248 runs -  4043.75 us/run -  60.13 GFLOP/run -  14.87 TFLOPS
MUL_MAT(type_a=iq2_xs,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                244 runs -  4112.35 us/run -  60.13 GFLOP/run -  14.62 TFLOPS
MUL_MAT(type_a=iq2_s,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 246 runs -  4065.67 us/run -  60.13 GFLOP/run -  14.79 TFLOPS
MUL_MAT(type_a=iq3_xxs,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):               244 runs -  4129.54 us/run -  60.13 GFLOP/run -  14.56 TFLOPS
MUL_MAT(type_a=iq1_s,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 264 runs -  3788.42 us/run -  60.13 GFLOP/run -  15.87 TFLOPS
MUL_MAT(type_a=iq1_m,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 234 runs -  4280.64 us/run -  60.13 GFLOP/run -  14.05 TFLOPS
MUL_MAT(type_a=iq4_nl,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                252 runs -  3992.67 us/run -  60.13 GFLOP/run -  15.06 TFLOPS
MUL_MAT(type_a=iq3_s,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                 246 runs -  4084.30 us/run -  60.13 GFLOP/run -  14.72 TFLOPS
MUL_MAT(type_a=iq4_xs,type_b=f32,m=4096,n=512,k=14336,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1):                224 runs -  4489.23 us/run -  60.13 GFLOP/run -  13.39 TFLOPS

```

</details>

Changed files:
filename: ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp
patch: @@ -313,12 +313,12 @@ void main() {
         sums[i] = coopmat<ACC_TYPE, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
     }
 #else
-    ACC_TYPE sums[WMITER * TM * WNITER * TN];
+    ACC_TYPE_VEC2 sums[WMITER * TM * WNITER * TN/2];
     FLOAT_TYPE_VEC2 cache_a[WMITER * TM];
-    FLOAT_TYPE_VEC2 cache_b[TN];
+    FLOAT_TYPE_VEC2 cache_b;
 
-    [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) {
-        sums[i] = ACC_TYPE(0.0f);
+    [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN/2; i++) {
+        sums[i] = ACC_TYPE_VEC2(0.0f, 0.0f);
     }
 #endif
 
@@ -360,20 +360,22 @@ void main() {
                     cache_a[wsir * TM + j] = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * SHMEM_STRIDE + i];
                 }
             }
-            [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) {
-                [[unroll]] for (uint j = 0; j < TN; j++) {
-                    cache_b[j] = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + j) * SHMEM_STRIDE + i];
-                }
 
-                [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
-                    [[unroll]] for (uint cc = 0; cc < TN; cc++) {
-                        [[unroll]] for (uint cr = 0; cr < TM; cr++) {
-                            const uint sums_idx = (wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr;
-                            sums[sums_idx] = fma(ACC_TYPE(cache_a[wsir * TM + cr].x), ACC_TYPE(cache_b[cc].x), fma(ACC_TYPE(cache_a[wsir * TM + cr].y), ACC_TYPE(cache_b[cc].y), sums[sums_idx]));
+            [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) {
+                [[unroll]] for (uint cc = 0; cc < TN; cc++) {
+                    cache_b = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + cc) * SHMEM_STRIDE + i];
+
+                    [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
+                        [[unroll]] for (uint cr = 0; cr < TM / 2; cr++) {
+                            // [WNITER][TN][WMITER][TM / 2] -> [wsic][cc][wsir][cr]
+                            const uint sums_idx = (wsic * TN + cc) * WMITER * (TM / 2) + wsir * (TM / 2) + cr;
+                            sums[sums_idx].x = fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr    ].x), ACC_TYPE(cache_b.x), fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr    ].y), ACC_TYPE(cache_b.y), sums[sums_idx].x));
+                            sums[sums_idx].y = fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr + 1].x), ACC_TYPE(cache_b.x), fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr + 1].y), ACC_TYPE(cache_b.y), sums[sums_idx].y));
                         }
                     }
                 }
             }
+
         }
 #endif
 
@@ -388,8 +390,9 @@ void main() {
         }
     }
 #else
-    [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) {
-        sums[i] = clamp(sums[i], -ACC_TYPE_MAX, ACC_TYPE_MAX);
+    [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN/2; i++) {
+        sums[i].x = clamp(sums[i].x, -ACC_TYPE_MAX, ACC_TYPE_MAX);
+        sums[i].y = clamp(sums[i].y, -ACC_TYPE_MAX, ACC_TYPE_MAX);
     }
 #endif
 #endif
@@ -463,14 +466,21 @@ void main() {
 
                 const u16vec2 row_idx = row_ids[row_i - ic * BN];
 #endif // MUL_MAT_ID
-                [[unroll]] for (uint cr = 0; cr < TM; cr++) {
+                [[unroll]] for (uint cr = 0; cr < TM / 2; cr++) {
+                    const uint sums_idx = (wsic * TN + cc) * WMITER * (TM / 2) + wsir * (TM / 2) + cr;
 #ifdef MUL_MAT_ID
-                    if (dr_warp + cr < p.M) {
-                        data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]);
+                    if (dr_warp + 2 * cr < p.M) {
+                        data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr] = D_TYPE(sums[sums_idx].x);
+                    }
+                    if (dr_warp + 2 * cr + 1 < p.M) {
+                        data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr + 1] = D_TYPE(sums[sums_idx].y);
                     }
 #else
-                    if (dr_warp + cr < p.M && dc_warp + cc < p.N) {
-                        data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]);
+                    if (dr_warp + 2 * cr < p.M && dc_warp + cc < p.N) {
+                        data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + 2 * cr] = D_TYPE(sums[sums_idx].x);
+                    }
+                    if (dr_warp + 2 * cr + 1 < p.M && dc_warp + cc < p.N) {
+                        data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + 2 * cr + 1] = D_TYPE(sums[sums_idx].y);
                     }
 #endif // MUL_MAT_ID
                 }