wehub-resource-sync
86 lines
2.3 KiB
Plaintext
86 lines
2.3 KiB
Plaintext
layout(set=0, binding=0) writeonly buffer s0 {
|
|
FLOAT4 data[];
|
|
} uOutput;
|
|
|
|
layout(set=0, binding=1) readonly buffer s1 {
|
|
FLOAT4 data[];
|
|
} uInput;
|
|
|
|
layout(set=0, binding=2) readonly uniform constBuffer {
|
|
ivec4 size; // inside, outside, useRMSNorm, outside
|
|
float eps;
|
|
} uConstant;
|
|
|
|
layout(constant_id = 3) const uint USE_RMS = 0;
|
|
|
|
#ifdef LAYERNORM_SCALE
|
|
layout(set=0, binding=3) readonly buffer s2 {
|
|
FLOAT4 data[];
|
|
} uGamma;
|
|
|
|
layout(set=0, binding=4) readonly buffer s3 {
|
|
FLOAT4 data[];
|
|
} uBeta;
|
|
#endif
|
|
|
|
// 1 workgroup handles 1 "outside" row; parallel reduce across "inside".
|
|
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
|
|
|
|
shared float gSum[64];
|
|
shared float gSqSum[64];
|
|
|
|
void main() {
|
|
int inside = uConstant.size.x;
|
|
int outside = uConstant.size.w;
|
|
int y = int(gl_WorkGroupID.x);
|
|
if (y >= outside) {
|
|
return;
|
|
}
|
|
int inside4 = inside >> 2;
|
|
int base4 = y * inside4;
|
|
uint tid = gl_LocalInvocationID.x;
|
|
uint groupSize = gl_WorkGroupSize.x;
|
|
|
|
float sum = 0.0;
|
|
float sqSum = 0.0;
|
|
for (int j4 = int(tid); j4 < inside4; j4 += int(groupSize)) {
|
|
vec4 v4 = vec4(uInput.data[base4 + j4]);
|
|
if (USE_RMS == 0u) {
|
|
sum += v4.x + v4.y + v4.z + v4.w;
|
|
}
|
|
sqSum += dot(v4, v4);
|
|
}
|
|
|
|
if (USE_RMS == 0u) {
|
|
gSum[tid] = sum;
|
|
}
|
|
gSqSum[tid] = sqSum;
|
|
barrier();
|
|
|
|
for (uint stride = (groupSize >> 1); stride > 0; stride >>= 1) {
|
|
if (tid < stride) {
|
|
if (USE_RMS == 0u) {
|
|
gSum[tid] += gSum[tid + stride];
|
|
}
|
|
gSqSum[tid] += gSqSum[tid + stride];
|
|
}
|
|
barrier();
|
|
}
|
|
|
|
float invInside = 1.0 / float(inside);
|
|
float mean = (USE_RMS != 0u) ? 0.0 : (gSum[0] * invInside);
|
|
float squareMean = gSqSum[0] * invInside;
|
|
float var = (USE_RMS != 0u) ? squareMean : max(squareMean - mean * mean, 0.0);
|
|
float invStd = inversesqrt(var + uConstant.eps);
|
|
|
|
for (int j4 = int(tid); j4 < inside4; j4 += int(groupSize)) {
|
|
vec4 v4 = vec4(uInput.data[base4 + j4]);
|
|
vec4 nv = (USE_RMS != 0u) ? (v4 * invStd) : ((v4 - vec4(mean)) * invStd);
|
|
#ifdef LAYERNORM_SCALE
|
|
vec4 dst = nv * vec4(uGamma.data[j4]) + vec4(uBeta.data[j4]);
|
|
#else
|
|
vec4 dst = nv;
|
|
#endif
|
|
uOutput.data[base4 + j4] = FLOAT4(dst);
|
|
}
|
|
} |