vulkan: add POOL_1D op (#25431)
* 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:
@@ -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;
|
||||
}
|
||||
@@ -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"}}));
|
||||
|
||||
Reference in New Issue
Block a user