KBaba7/llama.cpp

1#version 4502 3layout (push_constant) uniform parameter4{5    uint ne; uint a_offset; uint d_offset;6    uint nb00; uint nb01; uint nb02; uint nb03;7    uint ne10; uint ne11; uint ne12; uint ne13;8    float sf0; float sf1; float sf2; float sf3;9} p;10 11#include "types.comp"12 13layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;14 15layout (binding = 0) readonly buffer A {A_TYPE data_a[];};16layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};17 18void main() {19    const uint idx = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x;20 21    if (idx >= p.ne) {22        return;23    }24 25    const uint i10 = idx % p.ne10;26    const uint i11 = (idx / p.ne10) % p.ne11;27    const uint i12 = (idx / (p.ne10 * p.ne11)) % p.ne12;28    const uint i13 = (idx / (p.ne10 * p.ne11 * p.ne12)) % p.ne13;29 30    const uint i00 = uint(i10 / p.sf0);31    const uint i01 = uint(i11 / p.sf1);32    const uint i02 = uint(i12 / p.sf2);33    const uint i03 = uint(i13 / p.sf3);34 35    data_d[p.d_offset + idx] = D_TYPE(data_a[p.a_offset + i03 * p.nb03 + i02 * p.nb02 + i01 * p.nb01 + i00 * p.nb00]);36}37