611f419cff
* vulkan: optimize rms_norm, and allow the work to spread across multiple SMs There are really two parts to this change: (1) Some optimizations similar to what we have in soft_max, to unroll with different numbers of iterations. (2) A fusion optimization where we detect add followed by rms_norm, and make the add shader atomically accumulate the values^2 into memory. Then the rms_norm shader can just load that sum. This allows the rms_norm to be parallelized across multiple workgroups, it just becomes a simple per-element multiply. The fusion optimization is currently only applied when the rms_norm is on a single vector. This previously always ran on a single SM. It could apply more broadly, but when there are other dimensions the work can already spread across SMs, and there would be some complexity to tracking multiple atomic sums. * Change add+rms_norm optimization to write out an array of partial sums rather than using atomic add, to make it deterministic. The rms_norm shader fetches a subgroup's worth in parallel and uses subgroupAdd to add them up. * complete rebase against fused adds - multi_add shader can also compute partial sums * fix validation errors * disable add_rms_fusion for Intel due to possible driver bug * resolve against #15489, sync after clearing partial sums
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.comp"
|
|
#include "generic_binary_head.comp"
|
|
|
|
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
|
|
}
|