File size: 3,431 Bytes
48220c7 8e0c6a5 48220c7 8e0c6a5 48220c7 9088f09 48220c7 9088f09 48220c7 9088f09 48220c7 8e0c6a5 48220c7 8e0c6a5 48220c7 | 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 | {{ env.wgsl.resourceDeclarations }}
const HIDDEN: u32 = {{ hiddenSize }}u;
const EPSILON: f32 = {{ epsilon }};
const WG: u32 = {{ workgroupSize }}u;
{% if scaleRank > 0 %}
const X_RANK: u32 = {{ xRank }}u;
const SCALE_RANK: u32 = {{ scaleRank }}u;
const X_SHAPE: array<u32, {{ xRank }}> = array<u32, {{ xRank }}>({% for d in xShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
const SCALE_SHAPE: array<u32, {{ scaleRank }}> = array<u32, {{ scaleRank }}>({% for d in scaleShape %}{{ d }}u{% if not loop.last %}, {% endif %}{% endfor %});
fn x_stride(axis: u32) -> u32 {
var stride = 1u;
for (var i = axis + 1u; i < X_RANK; i += 1u) {
stride *= X_SHAPE[i];
}
return stride;
}
fn scale_stride(axis: u32) -> u32 {
var stride = 1u;
for (var i = axis + 1u; i < SCALE_RANK; i += 1u) {
stride *= SCALE_SHAPE[i];
}
return stride;
}
{% endif %}
fn scale_offset({% if scaleRank > 0 %}out_index: u32{% endif %}) -> u32 {
{% if scaleRank == 0 %}
return 0u;
{% else %}
var rem = out_index;
var offset = 0u;
for (var axis = 0u; axis < X_RANK; axis += 1u) {
let stride = x_stride(axis);
let coord = rem / stride;
rem %= stride;
let scale_axis = i32(axis) - i32(X_RANK - SCALE_RANK);
if (scale_axis >= 0) {
let s_axis = u32(scale_axis);
if (SCALE_SHAPE[s_axis] != 1u) {
offset += coord * scale_stride(s_axis);
}
}
}
return offset;
{% endif %}
}
var<workgroup> partial: array<f32, WG>;
{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
{% if op == "max" or op == "min" %}
{{ a }}[{{ idx }}] = {{ op }}({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);{% else %}
{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] {{ "*" if op == "prod" else "+" }} {{ 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, reuse=false) %}
var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
loop {
if ({{ svar }} == 0u) {
break;
}
if ({{ idx }} < {{ svar }}) {
{% for a in arrays %}
{{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
{% endfor %}
}
{{ svar }} = {{ svar }} / 2u;
workgroupBarrier();
}{% 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.
fn reduce_sum(value: f32, tid: u32) -> f32 {
partial[tid] = value;
workgroupBarrier();
{{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }}
return partial[0];
}
@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let row = wg.x + wg.y * params.rowStride;
if (row >= params.rows) {
return;
}
let tid = lid.x;
let base = row * HIDDEN;
var local_sq = 0.0;
for (var d = tid; d < HIDDEN; d = d + WG) {
let value = f32(x[base + d]);
local_sq = local_sq + value * value;
}
let inv = inverseSqrt(reduce_sum(local_sq, tid) / f32(HIDDEN) + EPSILON);
{% if writeStats %}
if (tid == 0u) {
inv_std_out[row] = inv;
}
{% endif %}
for (var d = tid; d < HIDDEN; d = d + WG) {
let index = base + d;
let value = f32(x[index]) * inv * f32(scale[scale_offset({% if scaleRank > 0 %}index{% endif %})]);
y[base + d] = {{ scalar }}(value);
}
}
|