70 lines
2.0 KiB
Plaintext
70 lines
2.0 KiB
Plaintext
#version 450
|
|
|
|
#extension GL_EXT_shader_16bit_storage : require
|
|
#if ADD_RMS
|
|
#extension GL_KHR_shader_subgroup_arithmetic : enable
|
|
#extension GL_KHR_shader_subgroup_basic : enable
|
|
#endif
|
|
|
|
#include "types.glsl"
|
|
#include "generic_binary_head.glsl"
|
|
|
|
const uint num_threads = 256;
|
|
|
|
layout (binding = 3, std430) buffer PartialBuf {float partial_sums[];};
|
|
|
|
layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in;
|
|
|
|
#if ADD_RMS
|
|
// XXX TODO this could be sized based on number of subgroups, but that't not considered a constant
|
|
shared FLOAT_TYPE sumsh[num_threads];
|
|
#endif
|
|
|
|
void main() {
|
|
uint idx = get_idx();
|
|
uint orig_idx = idx;
|
|
|
|
// num_threads * num_iter must equal 512, to match the wg_denoms and get_idx calculation
|
|
const uint num_iter = 2;
|
|
|
|
FLOAT_TYPE sum_sq = 0;
|
|
|
|
[[unroll]] for (uint i = 0; i < num_iter; ++i) {
|
|
if (idx >= p.ne) {
|
|
continue;
|
|
}
|
|
uint i00, i01, i02, i03;
|
|
get_indices(idx, i00, i01, i02, i03);
|
|
|
|
FLOAT_TYPE sum = FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]) + FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)]);
|
|
sum_sq += sum*sum;
|
|
|
|
data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] = D_TYPE(sum);
|
|
|
|
idx += num_threads;
|
|
}
|
|
|
|
#if ADD_RMS
|
|
if (p.param3 != 0) {
|
|
// reduce the sum within each subgroup, then across subgroups
|
|
const uint NumSubgroups = num_threads / gl_SubgroupSize;
|
|
sum_sq = subgroupAdd(sum_sq);
|
|
if (gl_SubgroupInvocationID == 0) {
|
|
sumsh[gl_SubgroupID] = sum_sq;
|
|
}
|
|
barrier();
|
|
[[unroll]] for (uint s = NumSubgroups / 2; s > 0; s >>= 1) {
|
|
if (gl_SubgroupID < s && gl_SubgroupInvocationID == 0) {
|
|
sum_sq += sumsh[gl_SubgroupID + s];
|
|
sumsh[gl_SubgroupID] = sum_sq;
|
|
}
|
|
barrier();
|
|
}
|
|
|
|
if (gl_SubgroupID == 0 && gl_SubgroupInvocationID == 0) {
|
|
partial_sums[orig_idx / (num_iter * num_threads)] = sum_sq;
|
|
}
|
|
}
|
|
#endif
|
|
}
|