Brunobkr/llama.cpp_AlgMor24_github

1#version 4502 3#include "types.glsl"4 5#extension GL_EXT_shader_16bit_storage : require6 7layout(push_constant) uniform parameter {8    uint IL;9    uint OL;10    uint OC;11    uint pelements;12    uint op;13    int k0;14    int s0;15    int p0;16} p;17 18#define BLOCK_SIZE 51219#define FLT_MAX 3.402823466e+38F20#define OP_POOL_MAX 0u21#define OP_POOL_AVG 1u22 23layout (local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;24 25layout(binding = 0) readonly buffer X {A_TYPE data_a[];};26layout(binding = 1) writeonly buffer D {D_TYPE data_d[];};27 28void main() {29    const uint idx = gl_GlobalInvocationID.x;30    if (idx >= p.pelements) {31        return;32    }33 34    const uint nc = idx / p.OL;35    const uint cur_ol = idx % p.OL;36 37    const int start = int(cur_ol) * p.s0 - p.p0;38    const int bl = max(start, 0);39    const int el = min(max(start + p.k0, 0), int(p.IL));40 41    const int window_size = el - bl;42    const float scale = window_size > 0 ? 1.0 / float(window_size) : 0.0;43    float res;44 45    if (p.op == OP_POOL_AVG) {46        res = 0.0;47    } else if (p.op == OP_POOL_MAX) {48        res = -FLT_MAX;49    } else {50        return;51    }52 53    #pragma unroll54    for (uint i = bl; i < el; i++) {55        const float cur = D_TYPE(data_a[nc * p.IL + i]);56 57        if (p.op == OP_POOL_AVG) {58            res += cur * scale;59        } else if (p.op == OP_POOL_MAX) {60            res = max(res, cur);61        }62    }63 64    data_d[nc * p.OL + cur_ol] = res;65}66