vulkan: implement SSM conv operation

Add State Space Model conv operation to the Vulkan backend.

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

Author: giuseppe

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

@@ -585,6 +585,7 @@ struct vk_device_struct {
     vk_pipeline pipeline_ssm_scan_f32_d16;
     vk_pipeline pipeline_ssm_scan_f32_d128;
     vk_pipeline pipeline_ssm_scan_f32_d256;
+    vk_pipeline pipeline_ssm_conv_f32;
     vk_pipeline pipeline_opt_step_adamw_f32;
     vk_pipeline pipeline_opt_step_sgd_f32;
     vk_pipeline pipeline_conv2d_f32[CONV_SHAPE_COUNT];
@@ -1097,6 +1098,12 @@ struct vk_op_ssm_scan_push_constants {
     uint32_t s_off;
     uint32_t n_head, d_head, n_group, n_tok;
 };
+struct vk_op_ssm_conv_push_constants {
+    uint32_t src0_nb1, src0_nb2;
+    uint32_t src1_nb1;
+    uint32_t dst_nb0, dst_nb1, dst_nb2;
+    uint32_t nc, ncs, nr, n_t, n_s;
+};
 
 struct vk_op_conv2d_push_constants {
     uint32_t Cout;
@@ -3602,6 +3609,8 @@ static void ggml_vk_load_shaders(vk_device& device) {
     ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d128, "ssm_scan_f32_d128", ssm_scan_f32_d128_len, ssm_scan_f32_d128_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {128, 32, 16}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d256, "ssm_scan_f32_d256", ssm_scan_f32_d256_len, ssm_scan_f32_d256_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {256, 32, 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), {1, 1, 1}, {32}, 1);
+
     ggml_vk_create_pipeline(device, device->pipeline_opt_step_adamw_f32, "opt_step_adamw_f32", opt_step_adamw_f32_len, opt_step_adamw_f32_data, "main", 5, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_opt_step_sgd_f32, "opt_step_sgd_f32", opt_step_sgd_f32_len, opt_step_sgd_f32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
@@ -8113,6 +8122,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
             }
         }
         return nullptr;
+    case GGML_OP_SSM_CONV:
+        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return ctx->device->pipeline_ssm_conv_f32;
+        }
+        return nullptr;
     case GGML_OP_OPT_STEP_ADAMW:
         if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
             return ctx->device->pipeline_opt_step_adamw_f32;
@@ -9153,6 +9167,77 @@ static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx,
     }, pc, elements);
 }
 
+static void ggml_vk_ssm_conv(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, bool dryrun = false) {
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+
+    GGML_ASSERT(!ggml_is_quantized(src0->type));
+    GGML_ASSERT(!ggml_is_quantized(src1->type));
+    GGML_ASSERT(dst->buffer != nullptr);
+
+    const uint32_t nc  = src1->ne[0];
+    const uint32_t ncs = src0->ne[0];
+    const uint32_t nr  = src0->ne[1];
+    const uint32_t n_t = dst->ne[1];
+    const uint32_t n_s = dst->ne[2];
+
+    vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, nullptr, dst, dst->op);
+    GGML_ASSERT(pipeline != nullptr);
+
+    if (dryrun) {
+        ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+        return;
+    }
+
+    const vk_op_ssm_conv_push_constants pc = {
+        (uint32_t)src0->nb[1], (uint32_t)src0->nb[2],
+        (uint32_t)src1->nb[1],
+        (uint32_t)dst->nb[0], (uint32_t)dst->nb[1], (uint32_t)dst->nb[2],
+        nc, ncs, nr, n_t, n_s
+    };
+
+    ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context;
+    ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context;
+    ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context;
+
+    vk_buffer d_D = nullptr, d_src0 = nullptr, d_src1 = nullptr;
+    size_t dst_offset = 0, src0_offset = 0, src1_offset = 0;
+    bool dst_uma = false, src0_uma = false, src1_uma = false;
+
+    if (ctx->device->uma) {
+        ggml_vk_host_get(ctx->device, src0->data, d_src0, src0_offset);
+        src0_uma = d_src0 != nullptr;
+        ggml_vk_host_get(ctx->device, src1->data, d_src1, src1_offset);
+        src1_uma = d_src1 != nullptr;
+        ggml_vk_host_get(ctx->device, dst->data, d_D, dst_offset);
+        dst_uma = d_D != nullptr;
+    }
+
+    if (!dst_uma) {
+        d_D = dst_buf_ctx->dev_buffer;
+        dst_offset = vk_tensor_offset(dst) + dst->view_offs;
+    }
+    if (!src0_uma) {
+        d_src0 = src0_buf_ctx->dev_buffer;
+        src0_offset = vk_tensor_offset(src0) + src0->view_offs;
+    }
+    if (!src1_uma) {
+        d_src1 = src1_buf_ctx->dev_buffer;
+        src1_offset = vk_tensor_offset(src1) + src1->view_offs;
+    }
+
+    size_t dst_size = ggml_nbytes(dst);
+    size_t src0_size = ggml_nbytes(src0);
+    size_t src1_size = ggml_nbytes(src1);
+
+    std::array<uint32_t, 3> elements = { CEIL_DIV(nr, 32), n_t, n_s };
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, {
+        vk_subbuffer{ d_src0, src0_offset, src0_size },
+        vk_subbuffer{ d_src1, src1_offset, src1_size },
+        vk_subbuffer{ d_D, dst_offset, dst_size }
+    }, pc, elements);
+}
+
 static void ggml_vk_op_f32_opt_step_adamw(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, const vk_op_push_constants&& pc, bool dryrun = false) {
     const ggml_tensor * x = dst->src[0];
     const ggml_tensor * g = dst->src[1];
@@ -10870,6 +10955,7 @@ static bool ggml_vk_compute_forward(ggml_backend_vk_context* ctx, ggml_cgraph *
 // If submit is true the current all operations queued so far are being submitted to Vulkan to overlap cmdlist creation and GPU execution.
 static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int node_idx, ggml_tensor *node_begin, int node_idx_begin, bool dryrun, bool last_node, bool almost_ready, bool submit){
     ggml_tensor * node = cgraph->nodes[node_idx];
+
     if (ggml_is_empty(node) || !node->buffer) {
         return false;
     }
@@ -10986,6 +11072,7 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
     case GGML_OP_RWKV_WKV6:
     case GGML_OP_RWKV_WKV7:
     case GGML_OP_SSM_SCAN:
+    case GGML_OP_SSM_CONV:
     case GGML_OP_LEAKY_RELU:
     case GGML_OP_FLASH_ATTN_EXT:
     case GGML_OP_OPT_STEP_ADAMW:
@@ -11408,6 +11495,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
 
         break;
 
+    case GGML_OP_SSM_CONV:
+        ggml_vk_ssm_conv(ctx, compute_ctx, node, dryrun);
+
+        break;
+
     case GGML_OP_OPT_STEP_ADAMW:
         ggml_vk_opt_step_adamw(ctx, compute_ctx, node, dryrun);
 
@@ -11520,6 +11612,7 @@ static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_cgraph *
     case GGML_OP_RWKV_WKV6:
     case GGML_OP_RWKV_WKV7:
     case GGML_OP_SSM_SCAN:
+    case GGML_OP_SSM_CONV:
     case GGML_OP_LEAKY_RELU:
     case GGML_OP_REPEAT:
     case GGML_OP_REPEAT_BACK:
@@ -13053,6 +13146,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
 
                 return true;
             }
+        case GGML_OP_SSM_CONV:
+            return true;
         case GGML_OP_CONV_TRANSPOSE_1D:
             return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
         case GGML_OP_CONV_2D:
@@ -13714,6 +13809,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
     } else if (tensor->op == GGML_OP_SSM_SCAN) {
         tensor_clone = ggml_ssm_scan(ggml_ctx, src_clone[0], src_clone[1], src_clone[2],
                                      src_clone[3], src_clone[4], src_clone[5], src_clone[6]);
+    } else if (tensor->op == GGML_OP_SSM_CONV) {
+        tensor_clone = ggml_ssm_conv(ggml_ctx, src_clone[0], src_clone[1]);
     }
     else {
         std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl;

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

@@ -0,0 +1,44 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+
+#include "types.glsl"
+
+layout(constant_id = 0) const int BLOCK_SIZE = 32;
+
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout(binding = 0) readonly buffer Src0 { float src0[]; };
+layout(binding = 1) readonly buffer Src1 { float src1[]; };
+layout(binding = 2) buffer Dst { float dst[]; };
+
+layout(push_constant) uniform PushConstants {
+    uint src0_nb1; uint src0_nb2;
+    uint src1_nb1;
+    uint dst_nb0; uint dst_nb1; uint dst_nb2;
+    uint nc; uint ncs; uint nr; uint n_t; uint n_s;
+};
+
+void main() {
+    const uint global_thread_id = gl_WorkGroupID.x * gl_WorkGroupSize.x + gl_LocalInvocationID.x;
+    const uint i2 = gl_WorkGroupID.y;
+    const uint i3 = gl_WorkGroupID.z;
+
+    if (global_thread_id >= nr || i2 >= n_t || i3 >= n_s) {
+        return;
+    }
+
+    const uint i1 = global_thread_id;
+    const uint src0_base = i3 * (src0_nb2 / 4) + i2 + i1 * (src0_nb1 / 4);
+    const uint src1_base = i1 * (src1_nb1 / 4);
+    const uint dst_idx = i3 * (dst_nb2 / 4) + i2 * (dst_nb1 / 4) + i1;
+
+    float sum = 0.0;
+    [[unroll]] for (uint i0 = 0; i0 < nc; i0++) {
+        const uint src0_idx = src0_base + i0;
+        const uint src1_idx = src1_base + i0;
+        sum += src0[src0_idx] * src1[src1_idx];
+    }
+
+    dst[dst_idx] = sum;
+}

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

@@ -923,6 +923,8 @@ void process_shaders() {
     string_to_spv("ssm_scan_f32_d128", "ssm_scan.comp", {{"A_TYPE", "float"}});
     string_to_spv("ssm_scan_f32_d256", "ssm_scan.comp", {{"A_TYPE", "float"}});
 
+    string_to_spv("ssm_conv_f32", "ssm_conv.comp", {{"A_TYPE", "float"}});
+
     for (auto &c : compiles) {
         c.wait();
     }