| {% set yUnsigned = yUnsigned is defined and yUnsigned %} |
| {% set qMin = qMin | default(0 if yUnsigned else 0 - 128) %} |
| {% set qMax = qMax | default(255 if yUnsigned else 127) %} |
| // int8/uint8 tensors are stored widened (one u32/i32 per element), so the vec4 |
| // binding gives 128-bit loads/stores of four elements. Per-component arithmetic |
| // remains identical to the scalar quantize/dequantize paths. |
| {{ env.wgsl.resourceDeclarations }} |
| |
| {% set divisionF16 = divisionF16 is defined and divisionF16 %} |
| // Exact ONNX QuantizeLinear round-to-nearest-even, with identical handling of |
| // infinities, saturation, and halfway values across every kernel route. |
| {% if divisionF16 %} |
| fn round_quotient_half_to_even(quotient: f32) -> i32 { |
| let v = clamp(quotient, -2.0e9, 2.0e9); |
| let fl = floor(v); |
| let hi = fl + 1.0; |
| let fraction = v - fl; |
| if (fraction < 0.5) { return i32(fl); } |
| if (fraction > 0.5) { return i32(hi); } |
| let half = floor(fl * 0.5); |
| let is_even = (fl - half * 2.0) == 0.0; |
| return i32(select(hi, fl, is_even)); |
| } |
| |
| fn round_scaled_half_to_even(value: f32, scale: f32) -> i32 { |
| // ONNX precision=FLOAT16 (and an omitted precision with f16 y_scale) requires |
| // the division itself—not merely its operands—to round in f16. |
| return round_quotient_half_to_even(f32(f16(value) / f16(scale))); |
| } |
| {% else %} |
| fn round_scaled_half_to_even(value: f32, scale: f32) -> i32 { |
| // Clamp before the i32 cast so infinite and huge finite inputs saturate |
| // instead of invoking undefined conversion behavior. Compare distances in |
| // the input domain: doing the comparison on value / scale can move a value |
| // across a halfway boundary because WGSL division is not required to be |
| // correctly rounded. |
| let v = clamp(value / scale, -2.0e9, 2.0e9); |
| let fl = floor(v); |
| let lo = fl; |
| let hi = fl + 1.0; |
| let lo_dist = abs(value - lo * scale); |
| let hi_dist = abs(hi * scale - value); |
| if (lo_dist < hi_dist) { |
| return i32(lo); |
| } |
| if (hi_dist < lo_dist) { |
| return i32(hi); |
| } |
| let half = floor(fl * 0.5); |
| let is_even = (fl - half * 2.0) == 0.0; |
| return i32(select(hi, lo, is_even)); |
| } |
| {% endif %} |
| |
| |
| {% if vectorParams is defined and vectorParams %} |
| fn read_zero4({% if hasZero %}index: u32{% endif %}) -> vec4<i32> { |
| {% if hasZero %} |
| {% if yUnsigned %} |
| return vec4<i32>(y_zero_point[index]); |
| {% else %} |
| return y_zero_point[index]; |
| {% endif %} |
| {% else %} |
| return vec4<i32>(0); |
| {% endif %} |
| } |
| {% else %} |
| fn read_zero({% if hasZero %}index: u32{% endif %}) -> i32 { |
| {% if hasZero %} |
| {% if yUnsigned %} |
| return i32(y_zero_point[index]); |
| {% else %} |
| return y_zero_point[index]; |
| {% endif %} |
| {% else %} |
| return 0; |
| {% endif %} |
| } |
| {% endif %} |
| |
| @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }}) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| // The flat dispatch is folded across x/y at a fixed per-axis workgroup |
| // width; gid.y carries the high portion of the vector index. |
| let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u; |
| if (i >= params.count4) { |
| return; |
| } |
| {% if vectorParams is defined and vectorParams %} |
| let scale_index = i % (params.scaleSize / 4u); |
| {% elif crossingParams is defined and crossingParams %} |
| let base_index = i * 4u; |
| let scale_index0 = (base_index / params.inner) % params.scaleSize; |
| let scale_index1 = ((base_index + 1u) / params.inner) % params.scaleSize; |
| let scale_index2 = ((base_index + 2u) / params.inner) % params.scaleSize; |
| let scale_index3 = ((base_index + 3u) / params.inner) % params.scaleSize; |
| {% elif perAxis %} |
| // inner % 4 == 0, so all 4 lanes share one scale index. |
| let scale_index = ((i * 4u) / params.inner) % params.scaleSize; |
| {% else %} |
| let scale_index = 0u; |
| {% endif %} |
| let xv = x[i]; |
| {% if vectorParams is defined and vectorParams %} |
| let zp4 = read_zero4({% if hasZero %}scale_index{% endif %}); |
| {% elif crossingParams is defined and crossingParams %} |
| let zp0 = read_zero(scale_index0); |
| let zp1 = read_zero(scale_index1); |
| let zp2 = read_zero(scale_index2); |
| let zp3 = read_zero(scale_index3); |
| {% else %} |
| let zp = read_zero({% if hasZero %}scale_index{% endif %}); |
| {% endif %} |
| {% if vectorParams is defined and vectorParams %} |
| let scale = vec4<f32>(y_scale[scale_index]); |
| let q0 = clamp(round_scaled_half_to_even(f32(xv.x), scale.x) + zp4.x, {{ qMin }}, {{ qMax }}); |
| let q1 = clamp(round_scaled_half_to_even(f32(xv.y), scale.y) + zp4.y, {{ qMin }}, {{ qMax }}); |
| let q2 = clamp(round_scaled_half_to_even(f32(xv.z), scale.z) + zp4.z, {{ qMin }}, {{ qMax }}); |
| let q3 = clamp(round_scaled_half_to_even(f32(xv.w), scale.w) + zp4.w, {{ qMin }}, {{ qMax }}); |
| {% elif crossingParams is defined and crossingParams %} |
| let q0 = clamp(round_scaled_half_to_even(f32(xv.x), f32(y_scale[scale_index0])) + zp0, {{ qMin }}, {{ qMax }}); |
| let q1 = clamp(round_scaled_half_to_even(f32(xv.y), f32(y_scale[scale_index1])) + zp1, {{ qMin }}, {{ qMax }}); |
| let q2 = clamp(round_scaled_half_to_even(f32(xv.z), f32(y_scale[scale_index2])) + zp2, {{ qMin }}, {{ qMax }}); |
| let q3 = clamp(round_scaled_half_to_even(f32(xv.w), f32(y_scale[scale_index3])) + zp3, {{ qMin }}, {{ qMax }}); |
| {% else %} |
| let scale = f32(y_scale[scale_index]); |
| let q0 = clamp(round_scaled_half_to_even(f32(xv.x), scale) + zp, {{ qMin }}, {{ qMax }}); |
| let q1 = clamp(round_scaled_half_to_even(f32(xv.y), scale) + zp, {{ qMin }}, {{ qMax }}); |
| let q2 = clamp(round_scaled_half_to_even(f32(xv.z), scale) + zp, {{ qMin }}, {{ qMax }}); |
| let q3 = clamp(round_scaled_half_to_even(f32(xv.w), scale) + zp, {{ qMin }}, {{ qMax }}); |
| {% endif %} |
| {% if yUnsigned %} |
| y[i] = vec4<u32>(u32(q0), u32(q1), u32(q2), u32(q3)); |
| {% else %} |
| y[i] = vec4<i32>(q0, q1, q2, q3); |
| {% endif %} |
| } |
| |
