Repository: ggml-org/llama.cpp
Pull request: vulkan: Use coopmat2 for conv2d

Stacked on https://github.com/ggml-org/llama.cpp/pull/14933, Draft until that's merged.

I haven't done any perf tuning on this yet, there may still be more perf to get.

Directed perf tests:

```
5090 before:

  CONV_2D(ne_input=[19,19,256,16],ne_kernel=[4,4,256,4096],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                    307 runs -  3262.35 us/run - 137.42 GFLOP/run -  42.12 TFLOPS
  CONV_2D(ne_input=[19,19,8,16],ne_kernel=[4,4,8,128],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              137632 runs -     7.28 us/run - 133.69 MFLOP/run -  18.38 TFLOPS
  CONV_2D(ne_input=[19,19,8,16],ne_kernel=[4,4,8,130],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               99495 runs -    10.11 us/run - 135.78 MFLOP/run -  13.43 TFLOPS
  CONV_2D(ne_input=[19,19,4,16],ne_kernel=[2,2,4,4],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                483328 runs -     2.08 us/run - 642.82 kFLOP/run - 308.66 GFLOPS
  CONV_2D(ne_input=[224,224,3,1],ne_kernel=[3,3,3,8],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               181868 runs -     5.63 us/run -  20.90 MFLOP/run -   3.71 TFLOPS
  CONV_2D(ne_input=[224,224,1,1],ne_kernel=[2,2,1,8],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               188416 runs -     5.51 us/run -   2.78 MFLOP/run - 505.39 GFLOPS
  CONV_2D(ne_input=[224,224,1,8],ne_kernel=[2,2,1,8],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                58357 runs -    18.41 us/run -  22.28 MFLOP/run -   1.21 TFLOPS
  CONV_2D(ne_input=[58,58,32,1],ne_kernel=[3,3,32,64],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              100572 runs -     9.95 us/run - 115.40 MFLOP/run -  11.60 TFLOPS
  CONV_2D(ne_input=[58,58,32,8],ne_kernel=[3,3,32,64],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               26269 runs -    38.22 us/run - 923.24 MFLOP/run -  24.16 TFLOPS
  CONV_2D(ne_input=[16,16,128,8],ne_kernel=[3,3,128,512],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                    10670 runs -    93.73 us/run -   1.85 GFLOP/run -  19.73 TFLOPS
  CONV_2D(ne_input=[19,19,256,16],ne_kernel=[4,4,256,4096],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                    316 runs -  3174.30 us/run - 137.42 GFLOP/run -  43.29 TFLOPS
  CONV_2D(ne_input=[19,19,8,16],ne_kernel=[4,4,8,128],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              135388 runs -     7.41 us/run - 133.69 MFLOP/run -  18.04 TFLOPS
  CONV_2D(ne_input=[19,19,8,16],ne_kernel=[4,4,8,130],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               98021 runs -    10.24 us/run - 135.78 MFLOP/run -  13.26 TFLOPS
  CONV_2D(ne_input=[19,19,4,16],ne_kernel=[2,2,4,4],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                507904 runs -     1.98 us/run - 642.82 kFLOP/run - 325.29 GFLOPS
  CONV_2D(ne_input=[224,224,3,1],ne_kernel=[3,3,3,8],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               177082 runs -     5.67 us/run -  20.90 MFLOP/run -   3.68 TFLOPS
  CONV_2D(ne_input=[224,224,1,1],ne_kernel=[2,2,1,8],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               180224 runs -     5.57 us/run -   2.78 MFLOP/run - 499.80 GFLOPS
  CONV_2D(ne_input=[224,224,1,8],ne_kernel=[2,2,1,8],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                58357 runs -    18.41 us/run -  22.28 MFLOP/run -   1.21 TFLOPS
  CONV_2D(ne_input=[58,58,32,1],ne_kernel=[3,3,32,64],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              102306 runs -     9.79 us/run - 115.40 MFLOP/run -  11.78 TFLOPS
  CONV_2D(ne_input=[58,58,32,8],ne_kernel=[3,3,32,64],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               27250 runs -    36.84 us/run - 923.24 MFLOP/run -  25.06 TFLOPS
  CONV_2D(ne_input=[16,16,128,8],ne_kernel=[3,3,128,512],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                    11605 runs -    86.26 us/run -   1.85 GFLOP/run -  21.43 TFLOPS

5090 after:

  CONV_2D(ne_input=[19,19,256,16],ne_kernel=[4,4,256,4096],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                    952 runs -  1051.23 us/run - 137.42 GFLOP/run - 130.72 TFLOPS
  CONV_2D(ne_input=[19,19,8,16],ne_kernel=[4,4,8,128],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              203456 runs -     4.93 us/run - 133.69 MFLOP/run -  27.14 TFLOPS
  CONV_2D(ne_input=[19,19,8,16],ne_kernel=[4,4,8,130],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              152559 runs -     6.57 us/run - 135.78 MFLOP/run -  20.67 TFLOPS
  CONV_2D(ne_input=[19,19,4,16],ne_kernel=[2,2,4,4],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                466944 runs -     2.15 us/run - 642.82 kFLOP/run - 299.08 GFLOPS
  CONV_2D(ne_input=[224,224,3,1],ne_kernel=[3,3,3,8],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               196226 runs -     5.20 us/run -  20.90 MFLOP/run -   4.02 TFLOPS
  CONV_2D(ne_input=[224,224,1,1],ne_kernel=[2,2,1,8],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               196608 runs -     5.18 us/run -   2.78 MFLOP/run - 537.48 GFLOPS
  CONV_2D(ne_input=[224,224,1,8],ne_kernel=[2,2,1,8],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                62846 runs -    15.92 us/run -  22.28 MFLOP/run -   1.40 TFLOPS
  CONV_2D(ne_input=[58,58,32,1],ne_kernel=[3,3,32,64],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              140454 runs -     7.14 us/run - 115.40 MFLOP/run -  16.17 TFLOPS
  CONV_2D(ne_input=[58,58,32,8],ne_kernel=[3,3,32,64],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               51666 runs -    19.40 us/run - 923.24 MFLOP/run -  47.59 TFLOPS
  CONV_2D(ne_input=[16,16,128,8],ne_kernel=[3,3,128,512],type_kernel=f32,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                    25025 runs -    40.00 us/run -   1.85 GFLOP/run -  46.22 TFLOPS
  CONV_2D(ne_input=[19,19,256,16],ne_kernel=[4,4,256,4096],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                    961 runs -  1040.79 us/run - 137.42 GFLOP/run - 132.04 TFLOPS
  CONV_2D(ne_input=[19,19,8,16],ne_kernel=[4,4,8,128],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              201212 runs -     4.98 us/run - 133.69 MFLOP/run -  26.85 TFLOPS
  CONV_2D(ne_input=[19,19,8,16],ne_kernel=[4,4,8,130],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              154770 runs -     6.47 us/run - 135.78 MFLOP/run -  21.00 TFLOPS
  CONV_2D(ne_input=[19,19,4,16],ne_kernel=[2,2,4,4],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                499712 runs -     2.00 us/run - 642.82 kFLOP/run - 320.75 GFLOPS
  CONV_2D(ne_input=[224,224,3,1],ne_kernel=[3,3,3,8],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               196226 runs -     5.20 us/run -  20.90 MFLOP/run -   4.02 TFLOPS
  CONV_2D(ne_input=[224,224,1,1],ne_kernel=[2,2,1,8],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               204800 runs -     5.03 us/run -   2.78 MFLOP/run - 553.85 GFLOPS
  CONV_2D(ne_input=[224,224,1,8],ne_kernel=[2,2,1,8],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                62846 runs -    16.02 us/run -  22.28 MFLOP/run -   1.39 TFLOPS
  CONV_2D(ne_input=[58,58,32,1],ne_kernel=[3,3,32,64],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):              148257 runs -     6.75 us/run - 115.40 MFLOP/run -  17.11 TFLOPS
  CONV_2D(ne_input=[58,58,32,8],ne_kernel=[3,3,32,64],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):               52865 runs -    18.95 us/run - 923.24 MFLOP/run -  48.72 TFLOPS
  CONV_2D(ne_input=[16,16,128,8],ne_kernel=[3,3,128,512],type_kernel=f16,stride0=1,stride1=1,padding0=0,padding1=0,dilation0=1,dilation1=1,cwhn=0):                    23155 runs -    43.27 us/run -   1.85 GFLOP/run -  42.73 TFLOPS
```

stable-diffusion:

```
5090 before:

Vulkan Timings:
ADD: 85 x 68.008 us
CONT: 3 x 23.626 us
CONV_2D M=Cout=128, K=Cin*KW*KH=1152, N=N*OW*OH=262144: 5 x 1785.04 us (43290.9 GFLOPS/s)
CONV_2D M=Cout=128, K=Cin*KW*KH=2304, N=N*OW*OH=262144: 1 x 3306.5 us (46752 GFLOPS/s)
CONV_2D M=Cout=128, K=Cin*KW*KH=256, N=N*OW*OH=262144: 1 x 653.312 us (26245.2 GFLOPS/s)
CONV_2D M=Cout=256, K=Cin*KW*KH=2304, N=N*OW*OH=262144: 1 x 6177.79 us (50045.5 GFLOPS/s)
CONV_2D M=Cout=256, K=Cin*KW*KH=2304, N=N*OW*OH=65536: 5 x 1771.92 us (43620.9 GFLOPS/s)
CONV_2D M=Cout=256, K=Cin*KW*KH=4608, N=N*OW*OH=65536: 1 x 3477.5 us (44457.8 GFLOPS/s)
CONV_2D M=Cout=256, K=Cin*KW*KH=512, N=N*OW*OH=65536: 1 x 557.056 us (30810.4 GFLOPS/s)
CONV_2D M=Cout=3, K=Cin*KW*KH=1152, N=N*OW*OH=262144: 1 x 483.328 us (3747.25 GFLOPS/s)
CONV_2D M=Cout=4, K=Cin*KW*KH=4, N=N*OW*OH=4096: 1 x 4.352 us (26.3529 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=36, N=N*OW*OH=4096: 1 x 12.672 us (11750.1 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=4608, N=N*OW*OH=16384: 7 x 1892.63 us (40843.2 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=4608, N=N*OW*OH=4096: 10 x 692.732 us (27897.1 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=4608, N=N*OW*OH=65536: 1 x 6333.79 us (48818.2 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=512, N=N*OW*OH=4096: 4 x 90.832 us (23619.3 GFLOPS/s)
GROUP_NORM: 30 x 842.821 us
MUL: 30 x 59.145 us
MUL_MAT m=4096 n=4096 k=512: 1 x 223.072 us (76939.7 GFLOPS/s)
MUL_MAT m=512 n=4096 k=4096: 1 x 393.632 us (43639.2 GFLOPS/s)
SCALE: 1 x 26.592 us
SILU: 29 x 64.174 us
SOFT_MAX: 1 x 45.056 us
UPSCALE: 3 x 159.05 us
Total time: 95267.3 us.

5090 after:

Vulkan Timings:
ADD: 85 x 64.818 us
CONT: 3 x 23.946 us
CONV_2D M=Cout=128, K=Cin*KW*KH=1152, N=N*OW*OH=262144: 5 x 549.702 us (140578 GFLOPS/s)
CONV_2D M=Cout=128, K=Cin*KW*KH=2304, N=N*OW*OH=262144: 1 x 1081.38 us (142952 GFLOPS/s)
CONV_2D M=Cout=128, K=Cin*KW*KH=256, N=N*OW*OH=262144: 1 x 296.928 us (57745.7 GFLOPS/s)
CONV_2D M=Cout=256, K=Cin*KW*KH=2304, N=N*OW*OH=262144: 1 x 1934.34 us (159833 GFLOPS/s)
CONV_2D M=Cout=256, K=Cin*KW*KH=2304, N=N*OW*OH=65536: 5 x 565.254 us (136740 GFLOPS/s)
CONV_2D M=Cout=256, K=Cin*KW*KH=4608, N=N*OW*OH=65536: 1 x 1142.78 us (135285 GFLOPS/s)
CONV_2D M=Cout=256, K=Cin*KW*KH=512, N=N*OW*OH=65536: 1 x 215.072 us (79801.6 GFLOPS/s)
CONV_2D M=Cout=3, K=Cin*KW*KH=1152, N=N*OW*OH=262144: 1 x 335.872 us (5392.39 GFLOPS/s)
CONV_2D M=Cout=4, K=Cin*KW*KH=4, N=N*OW*OH=4096: 1 x 5.6 us (20.48 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=36, N=N*OW*OH=4096: 1 x 10.656 us (13973.1 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=4608, N=N*OW*OH=16384: 7 x 626.902 us (123306 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=4608, N=N*OW*OH=4096: 10 x 328.428 us (58841.5 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=4608, N=N*OW*OH=65536: 1 x 1936.38 us (159681 GFLOPS/s)
CONV_2D M=Cout=512, K=Cin*KW*KH=512, N=N*OW*OH=4096: 4 x 37.312 us (57498.6 GFLOPS/s)
GROUP_NORM: 30 x 828.746 us
MUL: 30 x 59.382 us
MUL_MAT m=4096 n=4096 k=512: 1 x 93.44 us (183680 GFLOPS/s)
MUL_MAT m=512 n=4096 k=4096: 1 x 100.704 us (170577 GFLOPS/s)
SCALE: 1 x 26.656 us
SILU: 29 x 63.516 us
SOFT_MAX: 1 x 40.928 us
UPSCALE: 3 x 165.888 us
Total time: 55182.3 us.
```

Changed files:
filename: ggml/src/ggml-vulkan/ggml-vulkan.cpp
patch: @@ -3096,6 +3096,10 @@ static void ggml_vk_load_shaders(vk_device& device) {
         uint32_t conv2d_SHMEM_PAD = 4;
         bool conv2d_UNROLL = true;
 
+        if (device->coopmat2) {
+            conv2d_SHMEM_PAD = 8; // 8 float16_t
+        }
+
         if (device->vendor_id == VK_VENDOR_ID_INTEL) {
             conv2d_SHMEM_PAD = 0;
             conv2d_UNROLL = false;
@@ -3154,7 +3158,14 @@ static void ggml_vk_load_shaders(vk_device& device) {
         std::array<uint32_t, 3> wg_denoms = { conv2d_BS_K, conv2d_BS_NPQ, 1 };
         std::vector<uint32_t> spec_constants = { conv2d_WG_SIZE, conv2d_BS_K, conv2d_BS_CRS, conv2d_BS_NPQ, conv2d_TS_K, use_collectives, conv2d_SHMEM_PAD };
 
-        if (conv2d_UNROLL) {
+        if (device->coopmat2) {
+            ggml_vk_create_pipeline(
+                device, device->pipeline_conv2d_f32[s], "conv2d_f32", conv2d_f32_cm2_len, conv2d_f32_cm2_data, "main", 3,
+                sizeof(vk_op_conv2d_push_constants), wg_denoms, spec_constants, 1, true, use_collectives);
+            ggml_vk_create_pipeline(
+                device, device->pipeline_conv2d_f16_f32[s], "conv2d_f16_f32", conv2d_f16_f32_cm2_len, conv2d_f16_f32_cm2_data, "main", 3,
+                sizeof(vk_op_conv2d_push_constants), wg_denoms, spec_constants, 1, true, use_collectives);
+        } else if (conv2d_UNROLL) {
             ggml_vk_create_pipeline(
                 device, device->pipeline_conv2d_f32[s], "conv2d_f32", conv2d_f32_unroll_len, conv2d_f32_unroll_data, "main", 3,
                 sizeof(vk_op_conv2d_push_constants), wg_denoms, spec_constants, 1, true, use_collectives);

filename: ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp
patch: @@ -1,6 +1,11 @@
 #version 450
 
 #extension GL_EXT_control_flow_attributes : enable
+#ifdef COOPMAT2
+#extension GL_NV_cooperative_matrix2 : enable
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#extension GL_KHR_memory_scope_semantics : enable
+#endif
 
 #ifdef USE_COLLECTIVES
 #    extension GL_KHR_shader_subgroup_shuffle : enable
@@ -91,6 +96,12 @@ uint32_t n_elems_out = K * NPQ;
 // Number of blocktiles per input
 uint32_t NB_CRS = splitWork(CRS, BS_CRS);
 
+#ifdef COOPMAT2
+#define SHMEM_TYPE float16_t
+#else
+#define SHMEM_TYPE float
+#endif
+
 const uint32_t Ash_stride = BS_CRS + SHMEM_PAD;
 const uint32_t Bsh_stride = BS_NPQ + SHMEM_PAD;
 
@@ -100,8 +111,8 @@ const uint32_t Bsh_numel = BS_CRS * BS_NPQ;
 const uint32_t Ash_len = BS_K * Ash_stride;
 const uint32_t Bsh_len = BS_CRS * Bsh_stride;
 
-shared float Ash[Ash_len];  // K x CRS
-shared float Bsh[Bsh_len];  // CRS x NPQ
+shared SHMEM_TYPE Ash[Ash_len];  // K x CRS
+shared SHMEM_TYPE Bsh[Bsh_len];  // CRS x NPQ
 
 // Threadtile sizes
 const uint32_t TS_NPQ = BS_K * BS_NPQ / WG_SIZE / TS_K;
@@ -110,10 +121,6 @@ const uint32_t TS_NPQ = BS_K * BS_NPQ / WG_SIZE / TS_K;
 const uint32_t NT_K   = BS_K / TS_K;
 const uint32_t NT_NPQ = BS_NPQ / TS_NPQ;
 
-float regA[TS_K];
-float regB[TS_NPQ];
-float regC[TS_K][TS_NPQ];
-
 /*
 Compute
 KxCRS @ CRSxNPQ = K x NPQ
@@ -145,12 +152,36 @@ uint fastdiv(uint n, uint mp, uint L) {
     return (msbs + n) >> L;
 }
 
+#ifdef COOPMAT2
+#define ACC_TYPE float16_t
+
+ACC_TYPE perElemOpStore(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem)
+{
+    uint32_t K_idx   = B_idx_K * BS_K + r;
+    uint32_t NPQ_idx = B_idx_NPQ * BS_NPQ + c;
+    uint32_t N_idx   = fastdiv(NPQ_idx, p.OWOHmp, p.OWOHL); // divide by p.OH * p.OW;
+    uint32_t OH_idx  = fastdiv(NPQ_idx - N_idx * p.OH * p.OW, p.OWmp, p.OWL); // divide by p.OW;
+    uint32_t OW_idx  = NPQ_idx - N_idx * p.OH * p.OW - OH_idx * p.OW;
+    uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + K_idx * p.nb2 + N_idx * p.nb3;
+    if (K_idx < K && NPQ_idx < NPQ) {
+        dst_data[dst_idx] = D_TYPE(elem);
+    }
+    return elem;
+}
+#endif
+
 void main() {
+#ifdef COOPMAT2
+    coopmat<ACC_TYPE, gl_ScopeWorkgroup, BS_K, BS_NPQ, gl_MatrixUseAccumulator> matC;
+    matC = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BS_K, BS_NPQ, gl_MatrixUseAccumulator>(0.0);
+#else
+    float regC[TS_K][TS_NPQ];
     for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
         for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
             regC[T_ly][T_lx] = 0.0;
         }
     }
+#endif
     /* Advance block in CRS dim */
     for (uint32_t B_idx_CRS = 0; B_idx_CRS < NB_CRS; B_idx_CRS++) {
         uint32_t CRS_idx_a;
@@ -199,7 +230,7 @@ void main() {
             if (K_idx >= K || CRS_idx_a >= CRS) {
                 val = 0.0;
             }
-            Ash[B_ly * Ash_stride + B_lx] = val;
+            Ash[B_ly * Ash_stride + B_lx] = SHMEM_TYPE(val);
         }
         /* Load input to B_block: (BS_CRS x BS_NPQ) */
         UNROLL for (uint32_t r_offset = 0; r_offset < BS_CRS; r_offset += BrpWg) {
@@ -244,11 +275,21 @@ void main() {
             if (CRS_idx_b >= CRS || NPQ_idx >= NPQ || H_idx < 0 || H_idx >= p.H || W_idx < 0 || W_idx >= p.W) {
                 val = 0.0;
             }
-            Bsh[B_ly * Bsh_stride + B_lx] = val;
+            Bsh[B_ly * Bsh_stride + B_lx] = SHMEM_TYPE(val);
         }
         barrier();
+#ifdef COOPMAT2
+        coopmat<float16_t, gl_ScopeWorkgroup, BS_K, BS_CRS, gl_MatrixUseA> matA;
+        coopmat<float16_t, gl_ScopeWorkgroup, BS_CRS, BS_NPQ, gl_MatrixUseB> matB;
+
+        coopMatLoad(matA, Ash, 0, Ash_stride, gl_CooperativeMatrixLayoutRowMajor);
+        coopMatLoad(matB, Bsh, 0, Bsh_stride, gl_CooperativeMatrixLayoutRowMajor);
+        matC = coopMatMulAdd(matA, matB, matC);
+#else
         if (T_y * TS_K < K) {
             UNROLL for (uint32_t CRS_lidx = 0; CRS_lidx < BS_CRS; CRS_lidx++) {
+                float regA[TS_K];
+                float regB[TS_NPQ];
                 for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
                     regA[T_ly] = Ash[(T_y * TS_K + T_ly) * Ash_stride + CRS_lidx];
                 }
@@ -262,9 +303,13 @@ void main() {
                 }
             }
         }
+#endif
         barrier();
     }
     /* Save C* */
+#ifdef COOPMAT2
+    coopMatPerElementNV(matC, matC, perElemOpStore);
+#else
     if (T_y * TS_K < K) {
         for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
             for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
@@ -280,4 +325,5 @@ void main() {
             }
         }
     }
+#endif
 }

filename: ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
patch: @@ -661,6 +661,9 @@ void process_shaders() {
     string_to_spv("conv2d_f32", "conv2d_mm.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_COLLECTIVES", "1"}, {"UNROLL", ""}});
     string_to_spv("conv2d_f16_f32", "conv2d_mm.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_COLLECTIVES", "1"}, {"UNROLL", ""}});
 
+    string_to_spv("conv2d_f32", "conv2d_mm.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_COLLECTIVES", "1"}, {"UNROLL", "[[unroll]]"}, {"COOPMAT2", "1"}}, true, false, true);
+    string_to_spv("conv2d_f16_f32", "conv2d_mm.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_COLLECTIVES", "1"}, {"UNROLL", "[[unroll]]"}, {"COOPMAT2", "1"}}, true, false, true);
+
     string_to_spv("conv2d_dw_whcn_f32", "conv2d_dw.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"WHCN", "1"}}));
     string_to_spv("conv2d_dw_cwhn_f32", "conv2d_dw.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"CWHN", "1"}}));