wehub-resource-sync
80 lines
2.3 KiB
Plaintext
80 lines
2.3 KiB
Plaintext
layout(set=0, binding=0) writeonly buffer s0{
|
|
FLOAT data[];
|
|
} uOutput;
|
|
|
|
layout(set=0, binding=1) readonly buffer s1{
|
|
FLOAT 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{
|
|
FLOAT data[];
|
|
} uGamma;
|
|
|
|
layout(set=0, binding=4) readonly buffer s3{
|
|
FLOAT data[];
|
|
} uBeta;
|
|
#endif
|
|
|
|
layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
|
|
#define UP_DIV(x, y) (((x)+(y)-1)/(y))
|
|
|
|
void main()
|
|
{
|
|
int y = int(gl_GlobalInvocationID.x);
|
|
if (y < uConstant.size.w)
|
|
{
|
|
float scale = 1.0 / float(uConstant.size.x);
|
|
if (USE_RMS != 0u) {
|
|
float squaremean = 0.0;
|
|
for (int j=0; j<uConstant.size.x; ++j)
|
|
{
|
|
float v = float(uInput.data[j + y * uConstant.size.x]);
|
|
squaremean = squaremean + (v * v * scale);
|
|
}
|
|
float invStd = 1.0 / sqrt(squaremean + uConstant.eps);
|
|
for (int j=0; j<uConstant.size.x; ++j)
|
|
{
|
|
float v = float(uInput.data[j + y * uConstant.size.x]);
|
|
#ifdef LAYERNORM_SCALE
|
|
float dst = (v * invStd) * float(uGamma.data[j]) + float(uBeta.data[j]);
|
|
#else
|
|
float dst = v * invStd;
|
|
#endif
|
|
uOutput.data[j + y * uConstant.size.x] = FLOAT(dst);
|
|
}
|
|
} else {
|
|
float mean = 0.0;
|
|
for (int j=0; j<uConstant.size.x; ++j)
|
|
{
|
|
mean = mean + (float(uInput.data[j + y * uConstant.size.x]) * scale);
|
|
}
|
|
float squaremean = 0.0;
|
|
for (int j=0; j<uConstant.size.x; ++j)
|
|
{
|
|
float v = float(uInput.data[j + y * uConstant.size.x]);
|
|
v = v - mean;
|
|
squaremean = squaremean + (v * v * scale);
|
|
}
|
|
float invStd = 1.0 / sqrt(squaremean + uConstant.eps);
|
|
for (int j=0; j<uConstant.size.x; ++j)
|
|
{
|
|
float v = float(uInput.data[j + y * uConstant.size.x]);
|
|
#ifdef LAYERNORM_SCALE
|
|
float dst = ((v - mean) * invStd) * float(uGamma.data[j]) + float(uBeta.data[j]);
|
|
#else
|
|
float dst = (v - mean) * invStd;
|
|
#endif
|
|
uOutput.data[j + y * uConstant.size.x] = FLOAT(dst);
|
|
}
|
|
}
|
|
}
|
|
}
|