KBaba7/llama.cpp

1#version 4502 3#extension GL_EXT_shader_16bit_storage : require4#extension GL_EXT_spirv_intrinsics: enable5#extension GL_EXT_control_flow_attributes : require6 7#if RTE168spirv_execution_mode(capabilities = [4467], 4462, 16); // RoundingModeRTE, 16 bits9#endif10 11layout (push_constant) uniform parameter12{13    uint batch_offset; uint offset_delta;14    uint IC;15    uint IW; uint IH;16    uint OW; uint OH;17    uint KW; uint KH;18    uint pelements;19    uint CHW;20    int s0; int s1;21    int p0; int p1;22    int d0; int d1;23} p;24 25#include "types.comp"26 27layout(constant_id = 0) const uint BLOCK_SIZE = 32;28 29const uint NUM_ITER = 512 / BLOCK_SIZE;30 31layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;32 33layout (binding = 0) readonly buffer X {A_TYPE data_a[];};34layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};35 36void main() {37    const uint gidx = gl_GlobalInvocationID.x;38 39    const uint oh = gl_GlobalInvocationID.y;40    const uint batch = gl_GlobalInvocationID.z / p.IC;41    const uint ic = gl_GlobalInvocationID.z % p.IC;42 43    A_TYPE values[NUM_ITER];44    uint offset_dst[NUM_ITER];45    [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) {46        values[idx] = A_TYPE(0);47    }48 49    [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) {50 51        const uint i = gidx * NUM_ITER + idx;52 53        const uint ksize = p.OW * (p.KH > 1 ? p.KW : 1);54        const uint kx = i / ksize;55        const uint kd = kx * ksize;56        const uint ky = (i - kd) / p.OW;57        const uint ix = i % p.OW;58 59        const uint iiw = ix * p.s0 + kx * p.d0 - p.p0;60        const uint iih = oh * p.s1 + ky * p.d1 - p.p1;61 62        offset_dst[idx] =63            ((batch * p.OH + oh) * p.OW + ix) * p.CHW +64            (ic * (p.KW * p.KH) + ky * p.KW + kx);65 66        if (i >= p.pelements) {67            continue;68        }69 70        if (iih < p.IH && iiw < p.IW) {71            const uint offset_src = ic * p.offset_delta + batch * p.batch_offset;72            values[idx] = data_a[offset_src + iih * p.IW + iiw];73        }74    }75 76    [[unroll]] for (uint idx = 0; idx < NUM_ITER; ++idx) {77 78        const uint i = gidx * NUM_ITER + idx;79 80        if (i >= p.pelements) {81            continue;82        }83 84        data_d[offset_dst[idx]] = D_TYPE(values[idx]);85    }86 87}88