{% if source.fromPartials %} // Parallel range finalization folds the guarded min/max partials and computes // the output scale and zero point. Each partial already includes zero in its range. {% else %} // The serial path computes the complete range, scale, zero point, and quantized // output in one invocation. {% endif %} {{ env.wgsl.resourceDeclarations }} // ONNX DynamicQuantizeLinear uses correctly rounded f32 division followed by // round-half-to-even. Every execution path uses these helpers so their numerical // behavior cannot drift apart. // WGSL permits f32 division to differ from the correctly-rounded result by // 2.5 ULP, and fma() inherits separate multiply/add accuracy rather than // promising a fused residual. Reconstruct the correctly-rounded normal result // with integer significand division when the quotient can affect an integer // rounding boundary. This is backend-independent and uses only exact u32 ops. fn dynamic_quantize_exact_div_normal(numerator: f32, denominator: f32) -> f32 { if (numerator == 0.0) { return numerator; } let numerator_bits = bitcast(numerator); let denominator_bits = bitcast(denominator); let sign_bits = (numerator_bits ^ denominator_bits) & 0x80000000u; let numerator_abs = numerator_bits & 0x7fffffffu; let denominator_abs = denominator_bits & 0x7fffffffu; if (denominator_abs == 0u || (numerator_abs & 0x7f800000u) == 0x7f800000u || (denominator_abs & 0x7f800000u) == 0x7f800000u) { return numerator / denominator; } var numerator_mantissa = numerator_abs & 0x007fffffu; var denominator_mantissa = denominator_abs & 0x007fffffu; let numerator_biased_exponent = (numerator_abs >> 23u) & 0xffu; let denominator_biased_exponent = (denominator_abs >> 23u) & 0xffu; var numerator_exponent: i32; var denominator_exponent: i32; if (numerator_biased_exponent == 0u) { numerator_exponent = -126; // Zero returned above. A non-zero subnormal reaches the implicit-bit // position in at most 23 exact shifts. while ((numerator_mantissa & 0x00800000u) == 0u) { numerator_mantissa = numerator_mantissa << 1u; numerator_exponent = numerator_exponent - 1; } } else { numerator_mantissa = numerator_mantissa | 0x00800000u; numerator_exponent = i32(numerator_biased_exponent) - 127; } if (denominator_biased_exponent == 0u) { denominator_exponent = -126; while ((denominator_mantissa & 0x00800000u) == 0u) { denominator_mantissa = denominator_mantissa << 1u; denominator_exponent = denominator_exponent - 1; } } else { denominator_mantissa = denominator_mantissa | 0x00800000u; denominator_exponent = i32(denominator_biased_exponent) - 127; } var quotient_exponent = numerator_exponent - denominator_exponent; var remainder = numerator_mantissa; if (remainder < denominator_mantissa) { remainder = remainder << 1u; quotient_exponent = quotient_exponent - 1; } // The normalized ratio is now in [1, 2). Emit its implicit bit followed by // all 23 stored significand bits using exact binary long division. var quotient_mantissa = 0x00800000u; remainder = remainder - denominator_mantissa; for (var digit = 0u; digit < 23u; digit = digit + 1u) { remainder = remainder << 1u; if (remainder >= denominator_mantissa) { remainder = remainder - denominator_mantissa; quotient_mantissa = quotient_mantissa | (1u << (22u - digit)); } } // Round the 24-bit significand to nearest, ties to even. remainder and its // doubled value are below 2^25, so no u32 overflow is possible. let twice_remainder = remainder << 1u; if (twice_remainder > denominator_mantissa || (twice_remainder == denominator_mantissa && (quotient_mantissa & 1u) != 0u)) { quotient_mantissa = quotient_mantissa + 1u; } if (quotient_mantissa == 0x01000000u) { quotient_mantissa = quotient_mantissa >> 1u; quotient_exponent = quotient_exponent + 1; } let biased_exponent = quotient_exponent + 127; if (biased_exponent <= 0 || biased_exponent >= 255) { // DynamicQuantizeLinear's GPU-supported ranges are normal and finite. // Preserve the implementation-defined subnormal/overflow behavior outside // that contract; the integer quantizer never takes this branch near a // half-integer boundary. return numerator / denominator; } let result_bits = sign_bits | (u32(biased_exponent) << 23u) | (quotient_mantissa & 0x007fffffu); return bitcast(result_bits); } fn dynamic_quantize_division_may_cross_half(estimate: f32) -> bool { let lower = floor(estimate); let fraction = estimate - lower; let magnitude = abs(estimate); let magnitude_bits = bitcast(magnitude); let adjacent = bitcast(magnitude_bits + 1u); let ulp = adjacent - magnitude; // Division is allowed 2.5 ULP error. Eight ULP also covers the factor-of-two // ULP change when an estimate straddles the 0.5 exponent boundary. return abs(fraction - 0.5) <= ulp * 8.0; } fn round_dynamic_half_to_even(value: f32, scale: f32) -> i32 { let estimate = value / scale; var scaled = estimate; // `select` evaluates both value operands in WGSL; use control flow so the // 23-bit software divide remains a rare boundary fallback, not O(23) work // for every quantized element. if (dynamic_quantize_division_may_cross_half(estimate)) { scaled = dynamic_quantize_exact_div_normal(value, scale); } let lower = floor(scaled); let fraction = scaled - lower; if (fraction < 0.5) { return i32(lower); } if (fraction > 0.5) { return i32(lower + 1.0); } let upper = lower + 1.0; let half_lower = floor(lower * 0.5); let lower_is_even = (lower - half_lower * 2.0) == 0.0; return i32(select(upper, lower, lower_is_even)); } @compute @workgroup_size(1) fn main(@builtin(global_invocation_id) gid: vec3) { if (gid.x != 0u) { return; } {% if source.fromPartials %} var min_value = partial_min[0]; var max_value = partial_max[0]; for (var i = 1u; i < params.numPartials; i = i + 1u) { min_value = min(min_value, partial_min[i]); max_value = max(max_value, partial_max[i]); } {% else %} var min_value = 0.0; var max_value = 0.0; for (var i = 0u; i < params.count; i = i + 1u) { min_value = min(min_value, x[i]); max_value = max(max_value, x[i]); } {% endif %} var scale = dynamic_quantize_exact_div_normal(max_value - min_value, 255.0); if (scale == 0.0) { scale = 1.0; } let zp_i32 = clamp(round_dynamic_half_to_even(-min_value, scale), 0, 255); y_scale[0] = scale; y_zero_point[0] = u32(zp_i32); {% if not source.fromPartials %} for (var i = 0u; i < params.count; i = i + 1u) { let q = clamp(round_dynamic_half_to_even(x[i], scale) + zp_i32, 0, 255); y[i] = u32(q); } {% endif %} }