Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

vulkan: add POOL_1D op (#25431)

Browse Source 
* vulkan : add pool1d push constants and pipeline field

Declared data structures needed for POOL1D OP, which are the vk_op_pool1d_push_constants struct and pipeline_pool1d_f32 field.

* vulkan : add pool1d compute shader

Added pool1d.comp for Vulkan backend mirroring the existing pool2d shader.

* vulkan : add full GGML_OP_POOL_1D support

Added pipeline creation and op dispatch for 1D pooling in the Vulkan backend.

* vulkan : fix pool1d shader logic

Registered pool1d_f32 in vulkan-shaders-gen.cpp and fixed tensor dimension indices and avg pool scale.

* vulkan : fix pool1d end boundary crash and expand test coverage

Fixed an issue where the shader crashed when the end boundary was negative when k0 < p0. Also, added more test cases related to this fix.
This commit is contained in:
Anand Patil
2026-07-31 16:48:58 +02:00
committed by GitHub
parent eb41d503ba
commit 876a432116
4 changed files with 130 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+61
View File
@@ -998,6 +998,7 @@ struct vk_device_struct {
vk_pipeline pipeline_snake_f32;
vk_pipeline pipeline_snake_f16;
vk_pipeline pipeline_snake_bf16;
vk_pipeline pipeline_pool1d_f32;
vk_pipeline pipeline_pool2d_f32;
vk_pipeline pipeline_rwkv_wkv6_f32;
vk_pipeline pipeline_rwkv_wkv7_f32;
@@ -1687,6 +1688,17 @@ struct vk_op_snake_push_constants {
uint32_t ne1;
};
struct vk_op_pool1d_push_constants {
uint32_t IL;
uint32_t OL;
uint32_t OC;
uint32_t pelements;
uint32_t op;
int32_t k0;
int32_t s0;
int32_t p0;
};
struct vk_op_pool2d_push_constants {
uint32_t IW; uint32_t IH;
uint32_t OW; uint32_t OH;
@@ -5622,6 +5634,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_snake_f16, "snake_f16", snake_f16_len, snake_f16_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_snake_bf16, "snake_bf16", snake_bf16_len, snake_bf16_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_pool1d_f32, "pool1d_f32", pool1d_f32_len, pool1d_f32_data, "main", 2, sizeof(vk_op_pool1d_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_pool2d_f32, "pool2d_f32", pool2d_f32_len, pool2d_f32_data, "main", 2, sizeof(vk_op_pool2d_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_rwkv_wkv6_f32, "rwkv_wkv6_f32", rwkv_wkv6_f32_len, rwkv_wkv6_f32_data, "main", 7, sizeof(vk_op_rwkv_wkv6_push_constants), {1, 1, 1}, {device->subgroup_size}, 1);
@@ -11335,6 +11348,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
case GGML_TYPE_BF16: return ctx->device->pipeline_col2im_1d_bf16;
default: return nullptr;
}
case GGML_OP_POOL_1D:
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
return ctx->device->pipeline_pool1d_f32;
}
return nullptr;
case GGML_OP_POOL_2D:
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
return ctx->device->pipeline_pool2d_f32;
@@ -11857,6 +11875,13 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
{
elements = { uint32_t(dst->ne[0]), uint32_t(dst->ne[1]), 1 };
} break;
case GGML_OP_POOL_1D:
{
const uint32_t N = dst->ne[3] * dst->ne[2];
const uint32_t OC = dst->ne[1];
const uint32_t OL = dst->ne[0];
elements = { N * OC * OL, 1, 1};
} break;
case GGML_OP_POOL_2D:
{
const uint32_t N = dst->ne[3];
@@ -13776,6 +13801,29 @@ static void ggml_vk_snake_dispatch_fused(ggml_backend_vk_context * ctx, vk_conte
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, a_buf, inv_b_buf, dst_buf }, pc, elements);
}
static void ggml_vk_pool_1d(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
uint32_t op = static_cast<uint32_t>(dst->op_params[0]);
const int32_t k0 = dst->op_params[1];
const int32_t s0 = dst->op_params[2];
const int32_t p0 = dst->op_params[3];
const uint32_t IL = src0->ne[0];
const uint32_t N = dst->ne[3] * dst->ne[2];
const uint32_t OC = dst->ne[1];
const uint32_t OL = dst->ne[0];
const uint32_t parallel_elements = N * OC * OL;
ggml_vk_op_f32<vk_op_pool1d_push_constants>(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_POOL_1D, {
IL, OL, OC,
parallel_elements,
op,
k0, s0, p0,
});
}
static void ggml_vk_pool_2d(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
uint32_t op = static_cast<uint32_t>(dst->op_params[0]);
const int32_t k1 = dst->op_params[1];
@@ -15287,6 +15335,10 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
case GGML_OP_CONV_TRANSPOSE_1D:
ggml_vk_conv_transpose_1d(ctx, compute_ctx, src0, src1, node);
break;
case GGML_OP_POOL_1D:
ggml_vk_pool_1d(ctx, compute_ctx, src0, node);
break;
case GGML_OP_POOL_2D:
ggml_vk_pool_2d(ctx, compute_ctx, src0, node);
@@ -18004,6 +18056,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
case GGML_OP_CONV_2D_DW:
return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16)
&& op->src[1]->type == GGML_TYPE_F32;
case GGML_OP_POOL_1D:
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
case GGML_OP_POOL_2D:
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
case GGML_OP_RWKV_WKV6:
@@ -18939,6 +18993,13 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
const int32_t oc = tensor->op_params[1];
const int32_t p0 = tensor->op_params[2];
tensor_clone = ggml_col2im_1d(ggml_ctx, src_clone[0], stride, oc, p0);
} else if (tensor->op == GGML_OP_POOL_1D) {
enum ggml_op_pool op = static_cast<ggml_op_pool>(tensor->op_params[0]);
const int32_t k0 = tensor->op_params[1];
const int32_t s0 = tensor->op_params[2];
const int32_t p0 = tensor->op_params[3];
tensor_clone = ggml_pool_1d(ggml_ctx, src_clone[0], op, k0, s0, p0);
} else if (tensor->op == GGML_OP_POOL_2D) {
enum ggml_op_pool op = static_cast<ggml_op_pool>(tensor->op_params[0]);
const int32_t k0 = tensor->op_params[1];