| // Copyright 2025 The Chromium Authors |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| use std::simd::prelude::*; |
| use std::simd::Simd; |
| |
| use crate::{Reg, Reg32, UReg}; |
| |
| /// Subblock sums and averages, used in eval_quant_err. |
| #[derive(Default, Copy, Clone)] |
| pub struct SubblockStats { |
| pub sum: [Reg; 3], |
| pub avg: [Reg; 3], |
| } |
| |
| #[inline] |
| pub fn fast_div_255_round(x: Reg) -> Reg { |
| let r = x + Simd::splat(128); |
| (r + ((r + Simd::splat(257)) >> 8)) >> 8 |
| } |
| |
| /// Compute subblock (2x4 or 4x2) sums and averages. |
| /// |
| /// Returns: subblock averages in order of top, bottom, left, right. |
| #[inline] |
| pub fn prepare_averages(data: &[[[Reg; 3]; 4]; 4]) -> [SubblockStats; 4] { |
| let mut sum_2x2 = [[Reg::default(); 3]; 4]; |
| for y in 0..2 { |
| for x in 0..2 { |
| for ch in 0..3 { |
| sum_2x2[y * 2 + x][ch] = data[y * 2][x * 2][ch] |
| + data[y * 2 + 1][x * 2][ch] |
| + data[y * 2][x * 2 + 1][ch] |
| + data[y * 2 + 1][x * 2 + 1][ch]; |
| } |
| } |
| } |
| let mut out = [SubblockStats::default(); 4]; |
| for (i, (s0, s1)) in [(0, 1), (2, 3), (0, 2), (1, 3)].into_iter().enumerate() { |
| for ch in 0..3 { |
| out[i].sum[ch] = sum_2x2[s0][ch] + sum_2x2[s1][ch]; |
| out[i].avg[ch] = (out[i].sum[ch] + Simd::splat(4)) >> 3; |
| } |
| } |
| out |
| } |
| |
| struct QuantResultWithErr { |
| /// Bit 47..32 of ETC1 codeword |
| lo: UReg, |
| /// Bit 63..48 of ETC1 codeword |
| hi: UReg, |
| /// Value of endpoint 0, scaled to `0..=255`` |
| scaled0: [Reg; 3], |
| /// Value of endpoint 1, scaled to `0..=255`` |
| scaled1: [Reg; 3], |
| /// Error metric, see [`eval_quant_err`] |
| err: Reg32, |
| } |
| |
| pub struct QuantResult { |
| /// Bit 47..32 of ETC1 codeword |
| pub lo: UReg, |
| /// Bit 63..48 of ETC1 codeword |
| pub hi: UReg, |
| /// Value of endpoint 0, scaled to `0..=255`` |
| pub scaled0: [Reg; 3], |
| /// Value of endpoint 1, scaled to `0..=255`` |
| pub scaled1: [Reg; 3], |
| } |
| |
| #[inline] |
| fn quant_444( |
| avg0: [Reg; 3], |
| avg1: [Reg; 3], |
| sum0: [Reg; 3], |
| sum1: [Reg; 3], |
| flip: bool, |
| ) -> QuantResultWithErr { |
| #[inline] |
| fn quant(avg: [Reg; 3]) -> [Reg; 3] { |
| avg.map(|x| fast_div_255_round(x * Simd::splat(15))) |
| } |
| |
| #[inline] |
| fn scale(q: [Reg; 3]) -> [Reg; 3] { |
| q.map(|x| (x << 4) | x) |
| } |
| |
| #[inline] |
| fn encode(q0: [Reg; 3], q1: [Reg; 3], flip: bool) -> (UReg, UReg) { |
| let flip = if flip { UReg::splat(1) } else { UReg::splat(0) }; |
| let diff = UReg::splat(0); |
| let base1_b = q1[2].cast::<u16>() << 8; |
| let base0_b = q0[2].cast::<u16>() << 12; |
| let lo = flip | diff | base1_b | base0_b; |
| |
| let base1_g = q1[1].cast::<u16>(); |
| let base0_g = q0[1].cast::<u16>() << 4; |
| let base1_r = q1[0].cast::<u16>() << 8; |
| let base0_r = q0[0].cast::<u16>() << 12; |
| let hi = base1_g | base0_g | base1_r | base0_r; |
| |
| (lo, hi) |
| } |
| |
| let q0 = quant(avg0); |
| let q1 = quant(avg1); |
| let scaled0 = scale(q0); |
| let scaled1 = scale(q1); |
| let err = eval_quant_err(scaled0, scaled1, sum0, sum1); |
| let (lo, hi) = encode(q0, q1, flip); |
| QuantResultWithErr { lo, hi, scaled0, scaled1, err } |
| } |
| |
| #[inline] |
| fn quant_555( |
| avg0: [Reg; 3], |
| avg1: [Reg; 3], |
| sum0: [Reg; 3], |
| sum1: [Reg; 3], |
| flip: bool, |
| ) -> QuantResultWithErr { |
| #[inline] |
| fn quant(avg: [Reg; 3]) -> [Reg; 3] { |
| avg.map(|x| fast_div_255_round(x * Simd::splat(31))) |
| } |
| #[inline] |
| fn scale(q: [Reg; 3]) -> [Reg; 3] { |
| // Per ETC1 spec, the "five-bit codewords are extended to eight bits by |
| // replicating the top three highest-order bits to the three lowest |
| // order bits". |
| q.map(|x| (x << 3) | (x >> 2)) |
| } |
| #[inline] |
| fn encode(q0: [Reg; 3], delta: [Reg; 3], flip: bool) -> (UReg, UReg) { |
| #[inline] |
| fn encode_delta(d: Reg) -> UReg { |
| d.cast::<u16>() & Simd::splat(0b111) |
| } |
| |
| let flip = if flip { UReg::splat(1) } else { UReg::splat(0) }; |
| let diff = UReg::splat(1 << 1); |
| let delta_b = encode_delta(delta[2]) << 8; |
| let base_b = q0[2].cast::<u16>() << 11; |
| let lo = flip | diff | delta_b | base_b; |
| |
| let delta_g = encode_delta(delta[1]); |
| let base_g = q0[1].cast::<u16>() << 3; |
| let delta_r = encode_delta(delta[0]) << 8; |
| let base_r = q0[0].cast::<u16>() << 11; |
| let hi = delta_g | base_g | delta_r | base_r; |
| |
| (lo, hi) |
| } |
| |
| let q0 = quant(avg0); |
| let q1 = quant(avg1); |
| let delta = [0, 1, 2].map(|i| (q1[i] - q0[i]).simd_clamp(Simd::splat(-4), Simd::splat(3))); |
| let q1 = [0, 1, 2].map(|i| q0[i] + delta[i]); |
| |
| let scaled0 = scale(q0); |
| let scaled1 = scale(q1); |
| let err = eval_quant_err(scaled0, scaled1, sum0, sum1); |
| let (lo, hi) = encode(q0, delta, flip); |
| QuantResultWithErr { lo, hi, scaled0, scaled1, err } |
| } |
| |
| #[inline] |
| fn eval_quant_err(q0: [Reg; 3], q1: [Reg; 3], sum0: [Reg; 3], sum1: [Reg; 3]) -> Reg32 { |
| // Target error metric: |
| // sum((x - q) ** 2) (for each pixel) |
| // where x is the original pixel value, and |
| // q is the quantized average of the block |
| // This can be rewritten as: |
| // sum(x ** 2) - 2 * sum(x * q) + sum(q ** 2) |
| // For relative comparisons, sum(x ** 2) is constant and can be omitted. |
| // With this and more simplification: |
| // q * sum(q - 2 * x) |
| // Assuming that we are computing the sum for 8 pixels within a subblock: |
| // q * (8 * q - 2 * sum(x)) |
| // Dividing by 2: |
| // q * ((q << 2) - sum(x)) |
| (0..3).fold(Reg32::splat(0), |mut acc, i| { |
| let q0 = q0[i].cast::<i32>(); |
| let q1 = q1[i].cast::<i32>(); |
| let sum0 = sum0[i].cast::<i32>(); |
| let sum1 = sum1[i].cast::<i32>(); |
| acc += q0 * ((q0 << 2) - sum0); |
| acc += q1 * ((q1 << 2) - sum1); |
| acc |
| }) |
| } |
| |
| #[inline] |
| fn quantize_endpoint_pairs( |
| avg0: [Reg; 3], |
| avg1: [Reg; 3], |
| sum0: [Reg; 3], |
| sum1: [Reg; 3], |
| flip: bool, |
| ) -> QuantResultWithErr { |
| let q444 = quant_444(avg0, avg1, sum0, sum1, flip); |
| let q555 = quant_555(avg0, avg1, sum0, sum1, flip); |
| |
| let prefer555_32 = q555.err.simd_lt(q444.err); |
| let prefer555 = prefer555_32.cast::<i16>(); |
| QuantResultWithErr { |
| lo: prefer555.select(q555.lo, q444.lo), |
| hi: prefer555.select(q555.hi, q444.hi), |
| scaled0: [0, 1, 2].map(|i| prefer555.select(q555.scaled0[i], q444.scaled0[i])), |
| scaled1: [0, 1, 2].map(|i| prefer555.select(q555.scaled1[i], q444.scaled1[i])), |
| err: prefer555_32.select(q555.err, q444.err), |
| } |
| } |
| |
| #[inline] |
| /// Search through flip / no-flip and individual / differential modes, and |
| /// return the result with the least MSE from original pixels. |
| pub fn quantize_averages(data: &[[[Reg; 3]; 4]; 4]) -> QuantResult { |
| let stats = prepare_averages(&data); |
| |
| let flip = |
| quantize_endpoint_pairs(stats[0].avg, stats[1].avg, stats[0].sum, stats[1].sum, true); |
| let no_flip = |
| quantize_endpoint_pairs(stats[2].avg, stats[3].avg, stats[2].sum, stats[3].sum, false); |
| |
| let prefer_flip = flip.err.simd_lt(no_flip.err).cast::<i16>(); |
| QuantResult { |
| lo: prefer_flip.select(flip.lo, no_flip.lo), |
| hi: prefer_flip.select(flip.hi, no_flip.hi), |
| scaled0: [0, 1, 2].map(|i| prefer_flip.select(flip.scaled0[i], no_flip.scaled0[i])), |
| scaled1: [0, 1, 2].map(|i| prefer_flip.select(flip.scaled1[i], no_flip.scaled1[i])), |
| } |
| } |
