| {% if usesF16 %} |
| enable f16; |
| {% endif %} |
| {{ env.wgsl.resourceDeclarations }} |
| |
| // com.microsoft.EmbedLayerNormalization, normalization pass. |
| // output = (sum - mean) / sqrt(variance + epsilon) * gamma + beta |
| // One workgroup owns one (batch, position) row of the summed embedding the |
| // previous pass left in `output`, and normalizes it in place. The statistics |
| // accumulate in f32 over the stored tensor-type values. |
| const HIDDEN: u32 = {{ hidden }}u; |
| const EPSILON: f32 = {{ epsilon }}; |
| const WG: u32 = {{ workgroupSize }}u; |
| |
| var<workgroup> partial: array<f32, WG>; |
| |
| {% macro wgsl_tree_fold_stmt(a, op, idx, svar) %} |
| {% if op == "max" %} |
| {{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]); |
| {%- else %} |
| {{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}]; |
| {%- endif %} |
| {% endmacro %} |
| {% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %} |
| var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u; |
| loop { |
| {% if form == "head" %} |
| {% if breakInline %} |
| if ({{ svar }} == 0u) { break; } |
| {% else %} |
| if ({{ svar }} == 0u) { |
| break; |
| } |
| {% endif %} |
| {% endif %} |
| {% if bodyInline %} |
| if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} } |
| {% else %} |
| if ({{ idx }} < {{ svar }}) { |
| {% for a in arrays %} |
| {{ wgsl_tree_fold_stmt(a, op, idx, svar) }} |
| {% endfor %} |
| } |
| {% endif %} |
| {% if form == "head" %} |
| {% if barrierFirst %} |
| workgroupBarrier(); |
| {{ svar }} = {{ svar }} / 2u; |
| {% else %} |
| {{ svar }} = {{ svar }} / 2u; |
| workgroupBarrier(); |
| {% endif %} |
| {% else %} |
| workgroupBarrier(); |
| if ({{ svar }} == 1u) { |
| break; |
| } |
| {{ svar }} = {{ svar }} / 2u; |
| {% endif %} |
| } |
| {%- endmacro %} |
| |
| // Reusing partial after this reduction requires a barrier between the read of |
| // partial[0] and the next write, or the next round can race the prior readers. |
| {% set trailingBarrier = trailingBarrier is defined and trailingBarrier %} |
| fn reduce_sum(value: f32, tid: u32) -> f32 { |
| partial[tid] = value; |
| workgroupBarrier(); |
| {{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }} |
| {% if trailingBarrier %} |
| let total = partial[0]; |
| workgroupBarrier(); |
| return total; |
| {% else %} |
| return partial[0]; |
| {% endif %} |
| } |
| |
| |
| @compute @workgroup_size(WG, 1, 1) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, |
| @builtin(num_workgroups) nwg: vec3<u32>, |
| @builtin(local_invocation_id) lid: vec3<u32>) { |
| // 2D-folded row index: wg.y carries the high bits past the |
| // maxComputeWorkgroupsPerDimension dispatch limit. |
| let token = wg.x + wg.y * nwg.x; |
| if (token >= params.tokens) { |
| return; |
| } |
| let tid = lid.x; |
| let base = token * HIDDEN; |
| |
| var local_sum = 0.0; |
| for (var i = tid; i < HIDDEN; i = i + WG) { |
| local_sum = local_sum + f32(output[base + i]); |
| } |
| let mean = reduce_sum(local_sum, tid) / f32(HIDDEN); |
| // Separates the mean reduction's read of partial[0] from the variance |
| // reduction's writes to the same workgroup array. |
| workgroupBarrier(); |
| |
| var local_sq = 0.0; |
| for (var i = tid; i < HIDDEN; i = i + WG) { |
| let centred = f32(output[base + i]) - mean; |
| local_sq = local_sq + centred * centred; |
| } |
| let deviation = sqrt(reduce_sum(local_sq, tid) / f32(HIDDEN) + EPSILON); |
| workgroupBarrier(); |
| |
| for (var i = tid; i < HIDDEN; i = i + WG) { |
| let centred = f32(output[base + i]) - mean; |
| output[base + i] = {{ scalar }}(centred / deviation * f32(gamma[i]) + f32(beta[i])); |
| } |
| } |
| |
