diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 9c01dabe6..ff72fa60a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -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(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(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(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(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(tensor->op_params[0]); const int32_t k0 = tensor->op_params[1]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/pool1d.comp b/ggml/src/ggml-vulkan/vulkan-shaders/pool1d.comp new file mode 100644 index 000000000..bb87631ce --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/pool1d.comp @@ -0,0 +1,65 @@ +#version 450 + +#include "types.glsl" + +#extension GL_EXT_shader_16bit_storage : require + +layout(push_constant) uniform parameter { + uint IL; + uint OL; + uint OC; + uint pelements; + uint op; + int k0; + int s0; + int p0; +} p; + +#define 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(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; + if (idx >= p.pelements) { + return; + } + + const uint nc = idx / p.OL; + const uint cur_ol = idx % p.OL; + + const int start = int(cur_ol) * p.s0 - p.p0; + const int bl = max(start, 0); + const int el = min(max(start + p.k0, 0), int(p.IL)); + + const int window_size = el - bl; + const float scale = window_size > 0 ? 1.0 / float(window_size) : 0.0; + float res; + + if (p.op == OP_POOL_AVG) { + res = 0.0; + } else if (p.op == OP_POOL_MAX) { + res = -FLT_MAX; + } else { + return; + } + + #pragma unroll + for (uint i = bl; i < el; i++) { + const float cur = D_TYPE(data_a[nc * p.IL + i]); + + if (p.op == OP_POOL_AVG) { + res += cur * scale; + } else if (p.op == OP_POOL_MAX) { + res = max(res, cur); + } + } + + data_d[nc * p.OL + cur_ol] = 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 3e6b39576..0223d2e01 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -1052,6 +1052,7 @@ void process_shaders() { string_to_spv("snake_f16", "snake.comp", {{"DATA_A_F16", "1"}, {"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); string_to_spv("snake_bf16", "snake.comp", {{"DATA_A_BF16", "1"}, {"DATA_D_BF16", "1"}, {"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); + 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("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 e101ca390..4098acaaf 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8192,9 +8192,9 @@ static std::vector> make_test_cases_eval() { for (ggml_type type_input : {GGML_TYPE_F32}) { for (ggml_op_pool pool_type : {GGML_OP_POOL_AVG, GGML_OP_POOL_MAX}) { - for (int k0 : {1, 3}) { - for (int s0 : {1, 2}) { - for (int p0 : {0, 1}) { + for (int k0 : {1, 2, 3}) { + for (int s0 : {1, 2, 3}) { + for (int p0 : {0, 1, 2, 3}) { test_cases.emplace_back(new test_pool1d(pool_type, type_input, { 10, 3, 2, 1 }, k0, s0, p0)); test_cases.emplace_back(new test_pool1d(pool_type, type_input, { 11, 1, 3, 2 }, k0, s0, p0)); test_cases.emplace_back(new test_pool1d(pool_type, type_input, { 128, 2, 1, 3 }, k0, s0, p0));