File size: 3,582 Bytes
5b47652 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | {% 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]));
}
}
|