61 lines
1.7 KiB
Plaintext
61 lines
1.7 KiB
Plaintext
#version 450
|
|
|
|
#extension GL_EXT_control_flow_attributes : require
|
|
|
|
#include "types.glsl"
|
|
|
|
layout(constant_id = 0) const uint BLOCK_SIZE = 32;
|
|
layout(constant_id = 1) const uint TOKENS_PER_WG = 16;
|
|
layout(constant_id = 2) const bool APPLY_BIAS = false;
|
|
layout(constant_id = 3) const bool APPLY_SILU = false;
|
|
|
|
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z = 1) in;
|
|
|
|
layout(binding = 0) readonly buffer Src0 { float src0[]; };
|
|
layout(binding = 1) readonly buffer Src1 { float src1[]; };
|
|
layout(binding = 2) readonly buffer Bias { float bias[]; };
|
|
layout(binding = 3) buffer Dst { float dst[]; };
|
|
|
|
layout(push_constant) uniform PushConstants {
|
|
uint nb01; uint nb02;
|
|
uint nb11;
|
|
uint dst_nb0; uint dst_nb1; uint dst_nb2;
|
|
uint nc; uint ncs; uint nr; uint n_t; uint n_s;
|
|
};
|
|
|
|
void main() {
|
|
const uint i1 = gl_GlobalInvocationID.x;
|
|
const uint i2 = gl_WorkGroupID.y * TOKENS_PER_WG + gl_LocalInvocationID.y;
|
|
const uint i3 = gl_WorkGroupID.z;
|
|
|
|
if (i1 >= nr || i2 >= n_t || i3 >= n_s) {
|
|
return;
|
|
}
|
|
|
|
const uint src0_base = i3 * (nb02 / 4) + i2 + i1 * (nb01 / 4);
|
|
const uint src1_base = i1 * (nb11 / 4);
|
|
|
|
float sum = 0.0;
|
|
|
|
if (nc == 4) {
|
|
sum = dot(
|
|
vec4(src0[src0_base], src0[src0_base + 1], src0[src0_base + 2], src0[src0_base + 3]),
|
|
vec4(src1[src1_base], src1[src1_base + 1], src1[src1_base + 2], src1[src1_base + 3])
|
|
);
|
|
} else {
|
|
[[unroll]] for (uint i0 = 0; i0 < nc; i0++) {
|
|
sum += src0[src0_base + i0] * src1[src1_base + i0];
|
|
}
|
|
}
|
|
|
|
if (APPLY_BIAS) {
|
|
sum += bias[i1];
|
|
}
|
|
if (APPLY_SILU) {
|
|
sum = sum / (1.0f + exp(-sum));
|
|
}
|
|
|
|
const uint dst_idx = i3 * (dst_nb2 / 4) + i2 * (dst_nb1 / 4) + i1;
|
|
dst[dst_idx] = sum;
|
|
}
|