Commit 77c4540ce for llama.cpp
commit 77c4540ce4667f4577753fe034a3a8042f510b8f
Author: Jeff Bolz <jbolz@nvidia.com>
Date: Tue Oct 13 10:17:59 2026 -0500
vulkan: Optimize pool_2d when the pooling window is large (#30119)
When the window is larger than a subgroup, use a whole subgroup to
compute each output.
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
index 6ad178328..3d3b0ac00 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
@@ -958,6 +958,7 @@ struct vk_device_struct {
vk_pipeline pipeline_snake_bf16;
vk_pipeline pipeline_pool1d_f32;
vk_pipeline pipeline_pool2d_f32;
+ vk_pipeline pipeline_pool2d_subgroup_f32;
vk_pipeline pipeline_rwkv_wkv6_f32;
vk_pipeline pipeline_rwkv_wkv7_f32;
vk_pipeline pipeline_gated_linear_attn_f32;
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index e8589dc0e..1880098b9 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -3639,7 +3639,10 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
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_pool2d_f32, "pool2d_f32", pool2d_f32_len, pool2d_f32_data, "main", 2, sizeof(vk_op_pool2d_push_constants), {1, 1, 1}, {512}, 1);
+ if (device->subgroup_arithmetic && device->subgroup_require_full_support) {
+ ggml_vk_create_pipeline(device, device->pipeline_pool2d_subgroup_f32, "pool2d_subgroup_f32", pool2d_subgroup_f32_len, pool2d_subgroup_f32_data, "main", 2, sizeof(vk_op_pool2d_push_constants), {1, 1, 1}, {128}, 1, false, true, device->subgroup_size);
+ }
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);
@@ -9145,6 +9148,10 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
return nullptr;
case GGML_OP_POOL_2D:
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+ // Wide windows use neighboring subgroup lanes to read neighboring inputs.
+ if (ctx->device->pipeline_pool2d_subgroup_f32 && dst->op_params[1] >= int32_t(ctx->device->subgroup_size)) {
+ return ctx->device->pipeline_pool2d_subgroup_f32;
+ }
return ctx->device->pipeline_pool2d_f32;
}
return nullptr;
@@ -9587,11 +9594,12 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
} break;
case GGML_OP_POOL_2D:
{
- const uint32_t N = dst->ne[3];
- const uint32_t OC = dst->ne[2];
- const uint32_t OH = dst->ne[1];
- const uint32_t OW = dst->ne[0];
- elements = { N * OC * OH * OW, 1, 1};
+ const uint32_t outputs = ggml_nelements(dst);
+ const uint32_t outputs_per_workgroup = pipeline == ctx->device->pipeline_pool2d_subgroup_f32 ? 128 / ctx->device->subgroup_size : 512;
+ const uint32_t workgroups = CEIL_DIV(outputs, outputs_per_workgroup);
+ const uint32_t groups_y = CEIL_DIV(workgroups, ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
+ const uint32_t groups_x = CEIL_DIV(workgroups, groups_y);
+ elements = { groups_x, groups_y, 1 };
} break;
case GGML_OP_CONV_2D:
case GGML_OP_CONV_TRANSPOSE_2D:
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp b/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp
index d9d7166e3..10eedcfd2 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/pool2d.comp
@@ -1,5 +1,10 @@
#version 450
+#if defined(USE_SUBGROUP)
+#extension GL_KHR_shader_subgroup_basic : require
+#extension GL_KHR_shader_subgroup_arithmetic : require
+#endif
+
#include "types.glsl"
#extension GL_EXT_shader_16bit_storage : require
@@ -15,18 +20,27 @@ layout(push_constant) uniform parameter {
int p0; int p1;
} p;
-#define BLOCK_SIZE 512
+layout(constant_id = 0) const uint BLOCK_SIZE = 512;
#define FLT_MAX 3.402823466e+38F
#define OP_POOL_MAX 0u
#define OP_POOL_AVG 1u
-layout (local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
+layout (local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) readonly buffer X {A_TYPE data_a[];};
layout(binding = 1) writeonly buffer D {D_TYPE data_d[];};
void main() {
- const uint idx = gl_GlobalInvocationID.x;
+ const uint workgroup = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x;
+#if defined(USE_SUBGROUP)
+ const uint idx = workgroup * gl_NumSubgroups + gl_SubgroupID;
+ const int lane = int(gl_SubgroupInvocationID);
+ const int step = int(gl_SubgroupSize);
+#else
+ const uint idx = workgroup * BLOCK_SIZE + gl_LocalInvocationID.x;
+ const int lane = 0;
+ const int step = 1;
+#endif
if (idx >= p.pelements) {
return;
}
@@ -38,12 +52,12 @@ void main() {
const uint cur_ow = (idx % O_HW) % p.OW;
const int start_h = int(cur_oh) * p.s0 - p.p0;
- const uint bh = max(start_h, 0);
- const uint eh = min(start_h + p.k0, p.IH);
+ const int bh = max(start_h, 0);
+ const int eh = min(start_h + p.k0, int(p.IH));
const int start_w = int(cur_ow) * p.s1 - p.p1;
- const uint bw = max(start_w, 0);
- const uint ew = min(start_w + p.k1, p.IW);
+ const int bw = max(start_w, 0);
+ const int ew = min(start_w + p.k1, int(p.IW));
const float scale = 1.0 / float(p.k0 * p.k1);
float res;
@@ -57,18 +71,24 @@ void main() {
}
#pragma unroll
- for (uint i = bh; i < eh; i++) {
+ for (int i = bh; i < eh; i++) {
#pragma unroll
- for (uint j = bw; j < ew; j++) {
+ for (int j = bw + lane; j < ew; j += step) {
const float cur = D_TYPE(data_a[nc * p.IH * p.IW + i * p.IW + j]);
if (p.op == OP_POOL_AVG) {
- res += cur * scale;
+ res += cur;
} else if (p.op == OP_POOL_MAX) {
res = max(res, cur);
}
}
}
- data_d[nc * O_HW + cur_oh * p.OW + cur_ow] = res;
+#if defined(USE_SUBGROUP)
+ res = p.op == OP_POOL_AVG ? subgroupAdd(res) : subgroupMax(res);
+ if (!subgroupElect()) {
+ return;
+ }
+#endif
+ data_d[idx] = p.op == OP_POOL_AVG ? res * scale : res;
}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
index 94417f62b..9e8b89b1d 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
@@ -1141,6 +1141,7 @@ void process_shaders() {
string_to_spv("pool1d_f32", "pool1d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
string_to_spv("pool2d_f32", "pool2d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
+ string_to_spv("pool2d_subgroup_f32", "pool2d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_SUBGROUP", "1"}}));
string_to_spv("rwkv_wkv6_f32", "wkv6.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 0080c7e16..3ad705dee 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -9851,10 +9851,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
// MSA indexer block-max shape: wide 1-D window along ne0, no padding, non-divisible width
for (int64_t iw : {4096, 4160}) {
- for (int blk : {32, 64}) {
+ for (int blk : {32, 64, 128}) {
test_cases.emplace_back(new test_pool2d(pool_type, type_input, {iw, 2, 1, 1}, blk, 1, blk, 1, 0, 0));
}
}
+ test_cases.emplace_back(new test_pool2d(pool_type, type_input, {259, 5, 3, 2}, 128, 3, 64, 2, 64, 1));
+ test_cases.emplace_back(new test_pool2d(pool_type, type_input, {193, 2, 1, 1}, 128, 1, 128, 1, 128, 2));
}
}