vulkan: add optimization for SSM scan

use subgroupAdd if available

Signed-off-by: Giuseppe Scrivano <[email protected]>

Author: giuseppe

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -3607,8 +3607,9 @@ static void ggml_vk_load_shaders(vk_device& device) {
 
     ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv7_f32, "rwkv_wkv7_f32", rwkv_wkv7_f32_len, rwkv_wkv7_f32_data, "main", 8, sizeof(vk_op_rwkv_wkv7_push_constants), {1, 1, 1}, {device->subgroup_size}, 1);
 
-    ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d128, "ssm_scan_f32", ssm_scan_f32_len, ssm_scan_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {128, device->subgroup_size, 16}, 1);
-    ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d256, "ssm_scan_f32", ssm_scan_f32_len, ssm_scan_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {256, device->subgroup_size, 16}, 1);
+    const uint32_t ssm_variant = device->subgroup_arithmetic ? 1 : 0;
+    ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d128, "ssm_scan_f32", arr_ssm_scan_f32_len[ssm_variant], arr_ssm_scan_f32_data[ssm_variant], "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {128, device->subgroup_size, 16}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d256, "ssm_scan_f32", arr_ssm_scan_f32_len[ssm_variant], arr_ssm_scan_f32_data[ssm_variant], "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {256, device->subgroup_size, 16}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_ssm_conv_f32, "ssm_conv_f32", ssm_conv_f32_len, ssm_conv_f32_data, "main", 3, sizeof(vk_op_ssm_conv_push_constants), {32, 1, 1}, {32}, 1);
 

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

@@ -1,6 +1,10 @@
 #version 450
 
 #extension GL_EXT_control_flow_attributes : require
+#if USE_SUBGROUP_ADD
+#extension GL_KHR_shader_subgroup_basic : require
+#extension GL_KHR_shader_subgroup_arithmetic : require
+#endif
 
 #include "types.glsl"
 
@@ -97,10 +101,18 @@ void main() {
         }
 
         [[unroll]] for (uint j = 0; j <= SPLIT_H / (D_STATE / SUBGROUP_SIZE); j++) {
+            const uint k = tid / SUBGROUP_SIZE + j * (D_STATE / SUBGROUP_SIZE);
             const uint idx = (tid % SUBGROUP_SIZE) +
                             D_STATE * (tid / SUBGROUP_SIZE) +
                             j * D_STATE * (D_STATE / SUBGROUP_SIZE);
-
+#if USE_SUBGROUP_ADD_
+            if (idx < SPLIT_H * D_STATE) {
+                float sum = subgroupAdd(stateC[idx]);
+                if (tid % SUBGROUP_SIZE == 0) {
+                    d[y_base_idx + i * stride_y + k] = sum;
+                }
+            }
+#else
             uint lane = tid % SUBGROUP_SIZE;
 
             [[unroll]] for (uint offset = SUBGROUP_SIZE / 2; offset > 0; offset >>= 1) {
@@ -111,9 +123,9 @@ void main() {
             }
 
             if (idx < SPLIT_H * D_STATE && tid % SUBGROUP_SIZE == 0) {
-                const uint k = tid / SUBGROUP_SIZE + j * (D_STATE / SUBGROUP_SIZE);
                 d[y_base_idx + i * stride_y + k] = stateC[idx];
             }
+#endif
         }
 
         barrier();

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -917,6 +917,7 @@ void process_shaders() {
     string_to_spv("multi_add_rms_f32", "multi_add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}, {"RTE16", "1"}, {"ADD_RMS" , "1"}});
 
     string_to_spv("ssm_scan_f32", "ssm_scan.comp", {{"A_TYPE", "float"}});
+    string_to_spv("ssm_scan_f32_subgroup", "ssm_scan.comp", {{"A_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}});
 
     string_to_spv("ssm_conv_f32", "ssm_conv.comp", {{"A_TYPE", "float"}});
 
@@ -1043,6 +1044,13 @@ void write_output_files() {
     }
     }
 
+    hdr << "extern const void * arr_ssm_scan_f32_data[2];\n";
+    hdr << "extern const uint64_t arr_ssm_scan_f32_len[2];\n";
+    if (basename(input_filepath) == "ssm_scan.comp") {
+        src << "const void * arr_ssm_scan_f32_data[2] = {ssm_scan_f32_data, ssm_scan_f32_subgroup_data};\n";
+        src << "const uint64_t arr_ssm_scan_f32_len[2] = {ssm_scan_f32_len, ssm_scan_f32_subgroup_len};\n";
+    }
+
     if (input_filepath == "") {
         write_file_if_changed(target_hpp, hdr.str());
     }