// 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.
// Note: This file refers to modifiers in ETC1 spec as "selectors". The jargon
// was inherited from etcpak.
use std::simd::prelude::*;
use std::simd::{Mask, Simd};
use crate::{Reg, Reg32, UReg, SIMD_WIDTH};
// Selector tables from ETC1 spec. The negative part is omitted due to symmetry.
pub const TABLES: [[i16; 2]; 8] =
[[2, 8], [5, 17], [9, 29], [13, 42], [18, 60], [24, 80], [33, 106], [47, 183]];
/// Conditionally exchange the bottom left 2x2 block with top right 2x2 block,
/// if `flip` for that lane is true.
///
/// i.e. the goal is to flip from:
/// ```text
/// aeim
/// bfjn
/// cgko
/// dhlp
/// ```
/// to:
/// ```text
/// aecg
/// bfdh
/// imko
/// jnlp
/// ```
#[inline]
pub fn flip_pixels(d: &[[[Reg; 3]; 4]; 4], flip: Mask<i16, SIMD_WIDTH>) -> [[[Reg; 3]; 4]; 4] {
let mut o = [[[Reg::default(); 3]; 4]; 4];
for y0 in [0, 2] {
for x0 in [0, 2] {
for y1 in 0..2 {
for x1 in 0..2 {
for ch in 0..3 {
if y0 == x0 {
o[y0 + y1][x0 + x1][ch] = d[y0 + y1][x0 + x1][ch];
} else {
o[y0 + y1][x0 + x1][ch] =
flip.select(d[x0 + y1][y0 + x1][ch], d[y0 + y1][x0 + x1][ch]);
}
}
}
}
}
}
o
}
/// Flip the selector codeword if `flip` for that lane is true.
///
/// See [`flip_pixels`] for a description of the flip operation.
#[inline]
pub fn flip_selectors(x: UReg, flip: Mask<i16, SIMD_WIDTH>) -> UReg {
let keep = x & Simd::splat(0xCC33);
let bottom_left = x & Simd::splat(0x00CC);
let top_right = x & Simd::splat(0x3300);
let flipped = keep | (bottom_left << 6) | (top_right >> 6);
flip.select(flipped, x)
}
pub struct Fit {
pub err: Reg32,
pub table_idx: UReg,
pub selector_lo: UReg,
pub selector_hi: UReg,
}
/// Search for the optimal table and selectors for a subblock.
///
/// `data` should be in flipped layout, i.e. 4x2.
///
/// The error function used here is a bit quirky, see code comment for details.
#[inline]
pub fn search_table_and_selectors_subblock(data: &[[[Reg; 3]; 4]], base_color: [Reg; 3]) -> Fit {
assert_eq!(data.len(), 2);
// Use fold to compute minimum. Essentially a vector version of min_by_key.
TABLES
.iter()
.enumerate()
.fold(None, |best_fit, (table_idx, sel_table)| {
let mut outer_err = Reg32::splat(0);
let mut selector_lo = UReg::splat(0);
let mut selector_hi = UReg::splat(0);
for y in 0..2 {
for x in 0..4 {
// Below, we search for the optimal selector among [-lg, -sm, sm, lg] (sm
// and lg is from the selector table).
//
// We use the error metric:
// abs(gray(q + s - x))
// where q = quantized average, s = selector, x = pixel before compression
// gray(p) = 19*p.r + 38*p.g + 7*p.b (cf. rec601)
//
// Note that this is abs(gray(..)) not gray(abs(..)), i.e. the absolute
// is taken after computing to grayscale. This allows precomputing
// gray(q-x), then exploiting the fact that the selector is same for all
// three channels to calculate the final error with a single addition.
//
// We will first precompute gray(q - x).
let mut base_err = Reg::splat(0);
let rgb_weight = [19, 38, 7];
for ch in 0..3 {
base_err += (base_color[ch] - data[y][x][ch]) * Simd::splat(rgb_weight[ch]);
}
// Now, the sign of selector can be easily decided. To minimize the
// absolute value, the selector should be the opposite sign of
// gray(q - x).
let prefer_neg = base_err.simd_gt(Simd::splat(0));
// Finally, we compute the error metric for both sm and lg and decide the
// winner.
let base_err_abs = base_err.abs();
// Subtract in the direction that the final error metric is smaller.
// The selector is same for all three channels, so just multiply it by the
// total weight.
let weight_sum = 64;
let err_sm = (base_err_abs - Reg::splat(sel_table[0] * weight_sum)).abs();
let err_lg = (base_err_abs - Reg::splat(sel_table[1] * weight_sum)).abs();
let prefer_lg = err_lg.simd_lt(err_sm);
// The error can be fairly large (a crude upper bound is 255*64). To avoid
// overflow after squaring, we use widening multiply and accumulate. This
// is somewhat expensive.
let best_err = prefer_lg.select(err_lg, err_sm).cast::<i32>();
outer_err += best_err * best_err;
let pixel_idx = (y + x * 4) as u16;
selector_lo |= prefer_lg.select(UReg::splat(1 << pixel_idx), UReg::splat(0));
selector_hi |= prefer_neg.select(UReg::splat(1 << pixel_idx), UReg::splat(0));
}
}
let table_idx = UReg::splat(table_idx as u16);
match best_fit {
None => Some(Fit { err: outer_err, table_idx, selector_lo, selector_hi }),
Some(best) => {
let lt_32 = outer_err.simd_lt(best.err);
let lt = lt_32.cast::<i16>();
Some(Fit {
err: lt_32.select(outer_err, best.err),
table_idx: lt.select(table_idx, best.table_idx),
selector_lo: lt.select(selector_lo, best.selector_lo),
selector_hi: lt.select(selector_hi, best.selector_hi),
})
}
}
})
.unwrap()
}
/// Search through possible selector tables and selector values for each
/// subblock.
///
/// Returns: Four 16-bit codewords coding the optimal coefficients.
#[inline]
pub fn search_table_and_selectors(
mut hdr0: UReg,
hdr1: UReg,
data: &[[[Reg; 3]; 4]; 4],
base_color: [[Reg; 3]; 2],
) -> [UReg; 4] {
// We need to work on pixels in the first subblock, then the second. To allow
// uniform indices, the flip functions takes care of moving the first
// subblock to the top half and the second to bottom half. We will fix up
// the shuffled results in the end.
let flip = (hdr0 & (UReg::splat(1))).simd_ne(UReg::splat(0));
let permuted_data = flip_pixels(&data, !flip);
let mut selector_lo = UReg::splat(0);
let mut selector_hi = UReg::splat(0);
for subblock in 0..2 {
let best_fit = search_table_and_selectors_subblock(
&permuted_data[subblock * 2..subblock * 2 + 2],
base_color[subblock],
);
let subblock_bit = match subblock {
0 => 5,
1 => 2,
_ => unreachable!(),
};
hdr0 |= best_fit.table_idx << subblock_bit;
selector_lo |= best_fit.selector_lo << (subblock as u16 * 2);
selector_hi |= best_fit.selector_hi << (subblock as u16 * 2);
}
selector_lo = flip_selectors(selector_lo, !flip);
selector_hi = flip_selectors(selector_hi, !flip);
[selector_lo, selector_hi, hdr0, hdr1]
}