| {{ 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 %} |
| } |
| |
| {% if compensateHalfStats is defined and compensateHalfStats and xElement == "f16" %} |
| {% set halfOutputScalar = scalar %} |
| {% set halfScaleOffset = "scale_offset(base + i)" if scaleRank > 0 else "scale_offset()" %} |
| {% set halfOutputScalar = halfOutputScalar | default("f16") %} |
| {% if halfOutputScalar == "f16" %} |
| fn round_f16_bits_rte(value: f32) -> u32 { |
| let bits = bitcast<u32>(value); |
| let sign = (bits >> 16u) & 0x8000u; |
| let exponent_f32 = (bits >> 23u) & 0xffu; |
| let mantissa_f32 = bits & 0x7fffffu; |
| |
| if (exponent_f32 == 0xffu) { |
| if (mantissa_f32 != 0u) { |
| return 0x7e00u; |
| } |
| return sign | 0x7c00u; |
| } |
| |
| var exponent_f16 = i32(exponent_f32) - 127 + 15; |
| if (exponent_f16 >= 0x1f) { |
| return sign | 0x7c00u; |
| } |
| |
| if (exponent_f16 <= 0) { |
| if (exponent_f16 < -10) { |
| return sign; |
| } |
| let significand = mantissa_f32 | 0x800000u; |
| let shift = u32(14 - exponent_f16); |
| let halfway = 1u << (shift - 1u); |
| let discarded = significand & ((1u << shift) - 1u); |
| var fraction = significand >> shift; |
| if (discarded > halfway || (discarded == halfway && (fraction & 1u) == 1u)) { |
| fraction = fraction + 1u; |
| } |
| return sign | fraction; |
| } |
| |
| let halfway = 1u << 12u; |
| let discarded = mantissa_f32 & 0x1fffu; |
| var mantissa_f16 = mantissa_f32 >> 13u; |
| if (discarded > halfway || (discarded == halfway && (mantissa_f16 & 1u) == 1u)) { |
| mantissa_f16 = mantissa_f16 + 1u; |
| if (mantissa_f16 == 0x400u) { |
| mantissa_f16 = 0u; |
| exponent_f16 = exponent_f16 + 1; |
| } |
| } |
| if (exponent_f16 >= 0x1f) { |
| return sign | 0x7c00u; |
| } |
| return sign | (u32(exponent_f16) << 10u) | mantissa_f16; |
| } |
| |
| fn widen_f16_bits(value: u32) -> f32 { |
| return unpack2x16float(value & 0xffffu).x; |
| } |
| |
| // Typed ONNX edges must survive arithmetic fusion and narrow/wide casts. |
| // Integer rounding also fixes the ties-to-even rule independently of the |
| // implementation's floating-point conversion rounding mode. |
| fn half_stage(value: f32) -> f32 { |
| return widen_f16_bits(round_f16_bits_rte(value)); |
| } |
| {% else %} |
| fn round_f16_bits_rte(value: f32) -> u32 { |
| let bits = bitcast<u32>(value); |
| let sign = (bits >> 16u) & 0x8000u; |
| let exponent_f32 = (bits >> 23u) & 0xffu; |
| let mantissa_f32 = bits & 0x7fffffu; |
| |
| if (exponent_f32 == 0xffu) { |
| if (mantissa_f32 != 0u) { |
| return 0x7e00u; |
| } |
| return sign | 0x7c00u; |
| } |
| |
| var exponent_f16 = i32(exponent_f32) - 127 + 15; |
| if (exponent_f16 >= 0x1f) { |
| return sign | 0x7c00u; |
| } |
| |
| if (exponent_f16 <= 0) { |
| if (exponent_f16 < -10) { |
| return sign; |
| } |
| let significand = mantissa_f32 | 0x800000u; |
| let shift = u32(14 - exponent_f16); |
| let halfway = 1u << (shift - 1u); |
| let discarded = significand & ((1u << shift) - 1u); |
| var fraction = significand >> shift; |
| if (discarded > halfway || (discarded == halfway && (fraction & 1u) == 1u)) { |
| fraction = fraction + 1u; |
| } |
| return sign | fraction; |
| } |
| |
| let halfway = 1u << 12u; |
| let discarded = mantissa_f32 & 0x1fffu; |
| var mantissa_f16 = mantissa_f32 >> 13u; |
| if (discarded > halfway || (discarded == halfway && (mantissa_f16 & 1u) == 1u)) { |
| mantissa_f16 = mantissa_f16 + 1u; |
| if (mantissa_f16 == 0x400u) { |
| mantissa_f16 = 0u; |
| exponent_f16 = exponent_f16 + 1; |
| } |
| } |
| if (exponent_f16 >= 0x1f) { |
| return sign | 0x7c00u; |
| } |
| return sign | (u32(exponent_f16) << 10u) | mantissa_f16; |
| } |
| |
| fn widen_f16_bits(value: u32) -> f32 { |
| return unpack2x16float(value & 0xffffu).x; |
| } |
| {% endif %} |
| |
| // Half output magnifies statistics errors at rounding midpoints. Keep a low |
| // residual through the reduction and normalization, then round the typed |
| // float32-normalized/half-scale/half-bias edges explicitly. This is the standard ONNX half path; |
| // other normalization contracts retain their existing arithmetic. |
| |
| fn pair_add(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> { |
| let s = fma(a.x, 1.0, b.x); |
| // Materialize each rounded subtraction in the error-free transform. Plain |
| // cancellation expressions do not preserve the intended evaluation tree on |
| // every shader backend. |
| let bv = fma(-1.0, a.x, s); |
| let av = fma(-1.0, bv, s); |
| let a_error = fma(-1.0, av, a.x); |
| let b_error = fma(-1.0, bv, b.x); |
| let error = fma(a_error, 1.0, b_error); |
| let e = fma(fma(error, 1.0, a.y), 1.0, b.y); |
| let hi = fma(s, 1.0, e); |
| return vec2<f32>(hi, fma(-1.0, fma(-1.0, s, hi), e)); |
| } |
| |
| fn pair_mul(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> { |
| let p = fma(a.x, b.x, 0.0); |
| let error = fma(a.x, b.x, -p); |
| return pair_add(vec2<f32>(p, 0.0), vec2<f32>(error + (a.x * b.y + b.x * a.y), 0.0)); |
| } |
| |
| fn pair_div(a: vec2<f32>, b: f32) -> vec2<f32> { |
| let q = a.x / b; |
| let residual = pair_add(a, -pair_mul(vec2<f32>(q, 0.0), vec2<f32>(b, 0.0))); |
| return pair_add(vec2<f32>(q, 0.0), vec2<f32>((residual.x + residual.y) / b, 0.0)); |
| } |
| |
| fn pair_inverse_sqrt(a: vec2<f32>) -> vec2<f32> { |
| let r = inverseSqrt(a.x); |
| let rr = pair_mul(vec2<f32>(r, 0.0), vec2<f32>(r, 0.0)); |
| let residual = pair_add(vec2<f32>(1.0, 0.0), -pair_mul(a, rr)); |
| return pair_add(vec2<f32>(r, 0.0), vec2<f32>((0.5 * r) * (residual.x + residual.y), 0.0)); |
| } |
| |
| fn half_normalized(value: vec2<f32>) -> f32 { |
| // stash_type=1 materializes Normalized as float32 before its cast to half. |
| // Collapse the compensated residual at that typed edge; rounding the pair |
| // directly to half can choose a different result at a float32 midpoint. |
| return widen_f16_bits(round_f16_bits_rte(fma(value.x, 1.0, value.y))); |
| } |
| |
| var<workgroup> partial: array<vec2<f32>, WG>; |
| fn reduce_pair(value: vec2<f32>, tid: u32) -> vec2<f32> { |
| partial[tid] = value; |
| workgroupBarrier(); |
| for (var step = WG / 2u; step > 0u; step /= 2u) { |
| if (tid < step) { partial[tid] = pair_add(partial[tid], partial[tid + step]); } |
| workgroupBarrier(); |
| } |
| let total = partial[0]; |
| workgroupBarrier(); |
| return total; |
| } |
| {% set reduceArgs = "tid" %} |
| |
| fn load_value(index: u32) -> f32 { |
| return f32(x[index]); |
| } |
| |
| fn normalize_half_row(row: u32, tid: u32 |
| ) { |
| if (row >= params.rows) { return; } |
| let base = row * HIDDEN; |
| let mean = vec2<f32>(0.0); |
| var local_square = vec2<f32>(0.0); |
| for (var i = tid; i < HIDDEN; i += WG) { |
| let centered = pair_add(vec2<f32>(load_value(base + i), 0.0), -mean); |
| local_square = pair_add(local_square, pair_mul(centered, centered)); |
| } |
| let variance = pair_div(reduce_pair(local_square, {{ reduceArgs }}), f32(HIDDEN)); |
| let inv = pair_inverse_sqrt(pair_add(variance, vec2<f32>(EPSILON, 0.0))); |
| for (var i = tid; i < HIDDEN; i += WG) { |
| let normalized = half_normalized(pair_mul(pair_add(vec2<f32>(load_value(base + i), 0.0), -mean), inv)); |
| {% if halfOutputScalar == "f16" %} |
| var value = half_stage(normalized * f32(scale[{{ halfScaleOffset | default("i") }}])); |
| {% else %} |
| var value = normalized * f32(scale[{{ halfScaleOffset | default("i") }}]); |
| {% endif %} |
| y[base + i] = {{ halfOutputScalar }}(value); |
| } |
| } |
| @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; |
| normalize_half_row(row, lid.x); |
| } |
| {% else %} |
| 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); |
| for (var d = tid; d < HIDDEN; d = d + WG) { |
| let index = base + d; |
| // ONNX stage one ends by casting Normalized back to X's dtype; Scale is |
| // applied only after that rounding point. |
| let normalized = {{ xElement }}(f32(x[index]) * inv); |
| let value = f32(normalized) * f32(scale[scale_offset({% if scaleRank > 0 %}index{% endif %})]); |
| y[base + d] = {{ scalar }}(value); |
| } |
| } |
| {% endif %} |
| |
