use core::num::NonZeroUsize;
use core::time::Duration;
use std::sync::OnceLock;
pub const PAYS: f64 = 3.0;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DispatchKind {
Team,
InPool,
Outside,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WidthCosts {
widths: [u16; CAPACITY],
ns: [u64; CAPACITY],
len: u8,
}
const CAPACITY: usize = 12;
const fn to_ns(d: Duration) -> u64 {
let n = d.as_nanos();
if n > u64::MAX as u128 {
u64::MAX
} else {
n as u64
}
}
#[inline]
fn interpolate(c0: u64, c1: u64, w0: u16, w1: u16, w: usize) -> u64 {
let dw = (w - usize::from(w0)) as i64;
if let (Ok(a), Ok(b)) = (i64::try_from(c0), i64::try_from(c1)) {
if let Some(t) = (b - a).checked_mul(dw) {
let v = a + t / i64::from(w1 - w0);
return if v < 0 { 0 } else { v as u64 };
}
}
let (c0, c1) = (i128::from(c0), i128::from(c1));
let v = c0 + (c1 - c0) * (w - usize::from(w0)) as i128 / i128::from(w1 - w0);
if v < 0 {
0
} else {
v as u64
}
}
impl WidthCosts {
pub const MAX_ANCHORS: usize = CAPACITY;
pub const fn new(at_2_4_6: [Duration; 3]) -> Self {
Self::store(&[(2, at_2_4_6[0]), (4, at_2_4_6[1]), (6, at_2_4_6[2])])
}
pub const fn flat(cost: Duration) -> Self {
WidthCosts::new([cost; 3])
}
pub const fn from_anchors(anchors: &[(usize, Duration)]) -> Option<Self> {
let n = anchors.len();
if n == 0 || n > CAPACITY {
return None;
}
let mut k = 0;
while k < n {
let w = anchors[k].0;
if w < 2 || w > u16::MAX as usize {
return None;
}
if k > 0 && w <= anchors[k - 1].0 {
return None;
}
k += 1;
}
Some(Self::store(anchors))
}
const fn store(anchors: &[(usize, Duration)]) -> Self {
let mut widths = [0u16; CAPACITY];
let mut ns = [0u64; CAPACITY];
let mut k = 0;
while k < CAPACITY {
if k < anchors.len() {
widths[k] = anchors[k].0 as u16;
ns[k] = to_ns(anchors[k].1);
}
k += 1;
}
WidthCosts {
widths,
ns,
len: anchors.len() as u8,
}
}
#[inline]
pub(crate) fn ns_at(&self, width: NonZeroUsize) -> u64 {
let w = width.get();
let len = usize::from(self.len);
if w <= usize::from(self.widths[0]) {
return self.ns[0];
}
let mut i = 0;
while i + 1 < len {
let w1 = self.widths[i + 1];
if w <= usize::from(w1) {
return interpolate(self.ns[i], self.ns[i + 1], self.widths[i], w1, w);
}
i += 1;
}
let (last, w_m) = (self.ns[len - 1], u64::from(self.widths[len - 1]));
if let Some(p) = last.checked_mul(w as u64) {
return p / w_m;
}
let v = u128::from(last) * w as u128 / u128::from(w_m);
if v > u64::MAX as u128 {
u64::MAX
} else {
v as u64
}
}
#[inline]
pub fn at(&self, width: NonZeroUsize) -> Duration {
Duration::from_nanos(self.ns_at(width))
}
#[inline]
pub fn min(&self) -> Duration {
let len = usize::from(self.len);
let mut m = self.ns[0];
let mut k = 1;
while k < len {
if self.ns[k] < m {
m = self.ns[k];
}
k += 1;
}
Duration::from_nanos(m)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum HostClass {
X86_64Avx512,
X86_64Avx2,
X86_64Scalar,
Aarch64Apple,
Aarch64,
Generic,
}
impl HostClass {
pub fn detect() -> Self {
static CLASS: OnceLock<HostClass> = OnceLock::new();
*CLASS.get_or_init(Self::detect_uncached)
}
#[cfg(target_arch = "x86_64")]
fn detect_uncached() -> Self {
if std::arch::is_x86_feature_detected!("avx512f") {
HostClass::X86_64Avx512
} else if std::arch::is_x86_feature_detected!("avx2")
&& std::arch::is_x86_feature_detected!("fma")
{
HostClass::X86_64Avx2
} else {
HostClass::X86_64Scalar
}
}
#[cfg(all(target_arch = "aarch64", target_os = "macos"))]
fn detect_uncached() -> Self {
HostClass::Aarch64Apple
}
#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
fn detect_uncached() -> Self {
HostClass::Aarch64
}
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
fn detect_uncached() -> Self {
HostClass::Generic
}
pub const fn name(self) -> &'static str {
match self {
HostClass::X86_64Avx512 => "x86_64-avx512",
HostClass::X86_64Avx2 => "x86_64-avx2",
HostClass::X86_64Scalar => "x86_64-scalar",
HostClass::Aarch64Apple => "aarch64-apple",
HostClass::Aarch64 => "aarch64",
HostClass::Generic => "generic",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DispatchCosts {
team: WidthCosts,
in_pool: WidthCosts,
outside: WidthCosts,
}
const fn us_f(x: f64) -> Duration {
Duration::from_nanos((x * 1000.0) as u64)
}
const fn anchored(anchors: &[(usize, Duration)]) -> WidthCosts {
match WidthCosts::from_anchors(anchors) {
Some(costs) => costs,
None => panic!("invalid measured anchors"),
}
}
const AVX512_TEAM: WidthCosts = anchored(&[
(2, us_f(0.6)),
(4, us_f(0.7)),
(6, us_f(0.9)),
(8, us_f(0.98)),
]);
const AVX2_TEAM: WidthCosts = anchored(&[
(2, Duration::from_nanos(1423)),
(4, Duration::from_nanos(1992)),
(6, Duration::from_nanos(2329)),
(8, Duration::from_nanos(2738)),
(12, Duration::from_nanos(3558)),
(16, Duration::from_nanos(4076)),
(24, Duration::from_nanos(4292)),
(32, Duration::from_nanos(4838)),
(48, Duration::from_nanos(6402)),
(64, Duration::from_nanos(9333)),
]);
impl DispatchCosts {
pub fn new(team: WidthCosts, in_pool: WidthCosts, outside: WidthCosts) -> Self {
DispatchCosts {
team,
in_pool,
outside,
}
}
pub fn for_class(class: HostClass) -> Self {
match class {
HostClass::X86_64Avx512 => DispatchCosts {
team: AVX512_TEAM,
in_pool: WidthCosts::new([us_f(3.0), us_f(3.6), us_f(3.5)]),
outside: WidthCosts::new([us_f(9.7), us_f(10.1), us_f(8.7)]),
},
HostClass::Aarch64Apple => DispatchCosts {
team: WidthCosts::new([us_f(3.0), us_f(2.3), us_f(1.4)]),
in_pool: WidthCosts::new([us_f(6.0), us_f(11.0), us_f(5.5)]),
outside: WidthCosts::new([us_f(14.0), us_f(18.0), us_f(19.0)]),
},
HostClass::X86_64Avx2 => DispatchCosts {
team: AVX2_TEAM,
in_pool: WidthCosts::flat(us_f(15.0)),
outside: WidthCosts::flat(us_f(40.0)),
},
HostClass::X86_64Scalar | HostClass::Aarch64 | HostClass::Generic => DispatchCosts {
team: WidthCosts::flat(us_f(3.0)),
in_pool: WidthCosts::flat(us_f(15.0)),
outside: WidthCosts::flat(us_f(40.0)),
},
}
}
pub fn host_default() -> Self {
Self::for_class(HostClass::detect())
}
#[must_use]
pub fn with(mut self, kind: DispatchKind, costs: WidthCosts) -> Self {
match kind {
DispatchKind::Team => self.team = costs,
DispatchKind::InPool => self.in_pool = costs,
DispatchKind::Outside => self.outside = costs,
}
self
}
pub fn min_cost(&self, kind: DispatchKind) -> Duration {
match kind {
DispatchKind::Team => self.team.min(),
DispatchKind::InPool => self.in_pool.min(),
DispatchKind::Outside => self.outside.min(),
}
}
#[inline]
pub fn cost(&self, kind: DispatchKind, width: NonZeroUsize) -> Duration {
match kind {
DispatchKind::Team => self.team.at(width),
DispatchKind::InPool => self.in_pool.at(width),
DispatchKind::Outside => self.outside.at(width),
}
}
}
impl Default for DispatchCosts {
fn default() -> Self {
DispatchCosts::host_default()
}
}
const PAYS_INT: u128 = 3;
#[inline]
fn pays_ns(est_ns: u64, sections: u32, cost_ns: u64) -> bool {
u128::from(est_ns) >= PAYS_INT * u128::from(sections) * u128::from(cost_ns)
}
pub fn pays_at(estimate: Duration, sections: u32, dispatch_cost: Duration) -> bool {
if dispatch_cost == Duration::MAX {
return false;
}
if sections == 0 {
return true;
}
pays_ns(to_ns(estimate), sections, to_ns(dispatch_cost))
}
pub fn width_for_by(
estimate: Duration,
sections: u32,
budget: NonZeroUsize,
cost: impl Fn(NonZeroUsize) -> Duration,
) -> NonZeroUsize {
let est_ns = to_ns(estimate);
let est = est_ns as f64;
let mut best = NonZeroUsize::MIN;
let mut best_obj = est;
for k in 2..=budget.get() {
let Some(width) = NonZeroUsize::new(k) else {
continue;
};
let d = cost(width);
if d == Duration::MAX {
continue;
}
let d_ns = to_ns(d);
if sections != 0 && !pays_ns(est_ns, sections, d_ns) {
continue;
}
let obj = f64::from(sections) * d_ns as f64 + est / k as f64;
if obj < best_obj {
best = width;
best_obj = obj;
}
}
best
}
#[cfg(test)]
mod tests {
use super::*;
fn nz(n: usize) -> NonZeroUsize {
NonZeroUsize::new(n).unwrap()
}
fn us(n: u64) -> Duration {
Duration::from_micros(n)
}
#[test]
fn anchors_are_returned_exactly() {
let c = DispatchCosts::for_class(HostClass::Aarch64Apple);
assert_eq!(c.cost(DispatchKind::Team, nz(2)), us_f(3.0));
assert_eq!(c.cost(DispatchKind::Team, nz(4)), us_f(2.3));
assert_eq!(c.cost(DispatchKind::Team, nz(6)), us_f(1.4));
assert_eq!(c.cost(DispatchKind::Outside, nz(6)), us_f(19.0));
}
#[test]
fn interpolation_and_extrapolation() {
let c = WidthCosts::new([us(3), us(2), us(1)]);
assert_eq!(c.at(nz(1)), us(3));
assert_eq!(c.at(nz(3)), Duration::from_nanos(2500));
assert_eq!(c.at(nz(5)), Duration::from_nanos(1500));
assert_eq!(c.at(nz(8)), Duration::from_nanos(1333));
assert_eq!(c.at(nz(12)), us(2));
let n = WidthCosts::new([us(1), us(5), us(0)]);
assert_eq!(n.at(nz(5)), Duration::from_nanos(2500));
assert_eq!(n.at(nz(12)), Duration::ZERO);
}
#[test]
fn a_flat_row_extrapolates_in_proportion_to_the_width() {
let c = WidthCosts::flat(us(1));
assert_eq!(c.at(nz(6)), us(1));
assert_eq!(c.at(nz(12)), us(2));
assert_eq!(c.at(nz(24)), us(4));
assert_eq!(c.at(nz(48)), us(8));
assert_eq!(c.at(nz(64)), Duration::from_nanos(10_666));
assert_eq!(c.min(), us(1));
}
#[test]
fn above_the_last_anchor_is_non_decreasing_and_at_least_linear() {
let rows: [[Duration; 3]; 5] = [
[us_f(0.6), us_f(0.7), us_f(0.9)],
[us_f(3.0), us_f(3.6), us_f(3.5)],
[us_f(9.7), us_f(10.1), us_f(8.7)],
[us(3), us(2), us(1)],
[us(1), us(5), us(0)],
];
for row in rows {
let c = WidthCosts::new(row);
let cm = u128::from(to_ns(c.at(nz(6))));
let mut prev = c.at(nz(6));
for w in 7..=64u64 {
let v = c.at(nz(w as usize));
assert!(v >= prev, "{row:?}: width {w} fell");
assert!(
u128::from(to_ns(v)) >= cm * u128::from(w) / 6,
"{row:?}: width {w} below the linear bound"
);
prev = v;
}
}
}
#[test]
fn anchors_beyond_six_interpolate_and_extrapolate() {
let c = WidthCosts::from_anchors(&[
(2, Duration::from_nanos(8000)),
(4, Duration::from_nanos(4000)),
(8, Duration::from_nanos(2000)),
(16, Duration::from_nanos(1000)),
(32, Duration::from_nanos(500)),
(64, Duration::from_nanos(500)),
])
.unwrap();
let ns = |w: usize| c.at(nz(w)).as_nanos() as u64;
assert_eq!(ns(2), 8000);
assert_eq!(ns(3), 6000); assert_eq!(ns(6), 3000); assert_eq!(ns(12), 1500); assert_eq!(ns(24), 750); assert_eq!(ns(48), 500); assert_eq!(ns(64), 500); assert_eq!(ns(96), 750); assert_eq!(ns(128), 1000);
assert_eq!(c.min(), Duration::from_nanos(500));
}
#[test]
fn from_anchors_rejects_invalid_rows() {
assert!(WidthCosts::from_anchors(&[]).is_none());
assert!(WidthCosts::from_anchors(&[(0, us(1))]).is_none());
assert!(WidthCosts::from_anchors(&[(1, us(1)), (2, us(1))]).is_none());
assert!(WidthCosts::from_anchors(&[(2, us(1)), (2, us(1))]).is_none());
assert!(WidthCosts::from_anchors(&[(4, us(1)), (2, us(1))]).is_none());
let too_wide = [(usize::from(u16::MAX) + 1, us(1))];
assert!(WidthCosts::from_anchors(&too_wide).is_none());
assert!(WidthCosts::from_anchors(&[(usize::from(u16::MAX), us(1))]).is_some());
let mut many = [(0usize, Duration::ZERO); WidthCosts::MAX_ANCHORS + 1];
for (k, a) in many.iter_mut().enumerate() {
*a = (k + 2, us(1));
}
assert!(WidthCosts::from_anchors(&many[..WidthCosts::MAX_ANCHORS]).is_some());
assert!(WidthCosts::from_anchors(&many).is_none());
}
#[test]
fn from_anchors_is_const_and_capacity_is_invisible() {
const ROW: WidthCosts = match WidthCosts::from_anchors(&[
(2, Duration::from_micros(3)),
(4, Duration::from_micros(2)),
(8, Duration::from_micros(1)),
]) {
Some(c) => c,
None => panic!("valid anchors"),
};
assert_eq!(WidthCosts::MAX_ANCHORS, 12);
assert_eq!(ROW.at(nz(6)), Duration::from_nanos(1500));
assert_eq!(ROW.at(nz(16)), us(2));
assert_eq!(ROW.min(), us(1));
let three = WidthCosts::from_anchors(&[(2, us(3)), (4, us(2)), (6, us(1))]).unwrap();
assert_eq!(three, WidthCosts::new([us(3), us(2), us(1)]));
let four =
WidthCosts::from_anchors(&[(2, us(3)), (4, us(2)), (6, us(1)), (8, us(1))]).unwrap();
assert_ne!(three, four);
}
#[test]
fn a_duration_max_row_never_pays() {
let cap = Duration::from_nanos(u64::MAX);
let c = WidthCosts::flat(Duration::MAX);
assert_eq!(c.at(nz(6)), cap);
assert_eq!(c.at(nz(64)), cap); assert_eq!(c.min(), cap);
assert!(!pays_at(cap, 1, c.at(nz(64))));
assert!(!pays_at(cap, 2, c.at(nz(64))));
let mixed = WidthCosts::from_anchors(&[(2, Duration::ZERO), (4, Duration::MAX)]).unwrap();
assert_eq!(mixed.at(nz(64)), cap);
assert_eq!(mixed.min(), Duration::ZERO);
}
#[test]
fn a_wide_region_selects_an_interior_width_above_six() {
let row = WidthCosts::flat(us(1));
let d = |k: NonZeroUsize| row.at(k);
let budget = nz(64);
assert_eq!(width_for_by(us(24), 1, budget, d).get(), 12);
assert_eq!(width_for_by(us(1_000_000), 1, budget, d).get(), 64);
let below = Duration::from_nanos(3 * 1000 - 1);
assert!(!pays_at(below, 1, row.at(nz(6))));
assert!(!pays_at(below, 1, row.at(nz(64))));
assert_eq!(width_for_by(below, 1, budget, d).get(), 1);
}
fn old_at(ns: [u64; 3], w: u64) -> u64 {
if w <= 2 {
ns[0]
} else if w >= 6 {
ns[2]
} else {
let i = usize::from(w > 4);
let (w0, w1) = ([2u64, 4, 6][i], [2u64, 4, 6][i + 1]);
let (c0, c1) = (i128::from(ns[i]), i128::from(ns[i + 1]));
(c0 + (c1 - c0) * i128::from(w - w0) / i128::from(w1 - w0)).max(0) as u64
}
}
#[test]
fn rows_keep_the_old_values_inside_the_anchors() {
let tables = [
[3000, 2300, 1400],
[2800, 2200, 1700],
[14000, 18000, 19000],
[9500, 10100, 8800],
[3100, 3600, 3500],
[1000, 5000, 0],
[0, 0, 0],
[7, 3, 11],
];
for ns in tables {
let c = WidthCosts::new(ns.map(Duration::from_nanos));
for w in 1..=6u64 {
assert_eq!(
c.at(nz(w as usize)),
Duration::from_nanos(old_at(ns, w)),
"{ns:?} width {w}"
);
}
}
for d in [
DispatchCosts::for_class(HostClass::Aarch64Apple),
DispatchCosts::for_class(HostClass::X86_64Avx512),
DispatchCosts::for_class(HostClass::Generic),
] {
for kind in [
DispatchKind::Team,
DispatchKind::InPool,
DispatchKind::Outside,
] {
let m = (1..=64).map(|w| d.cost(kind, nz(w))).min().unwrap();
assert_eq!(d.min_cost(kind), m);
}
}
}
#[test]
fn measured_rows_are_pinned_at_widths_one_to_six() {
let avx512 = DispatchCosts::for_class(HostClass::X86_64Avx512);
let apple = DispatchCosts::for_class(HostClass::Aarch64Apple);
let generic = DispatchCosts::for_class(HostClass::Generic);
for (costs, kind, expected) in [
(
avx512,
DispatchKind::Team,
[600u64, 600, 650, 700, 800, 900],
),
(
avx512,
DispatchKind::InPool,
[3000, 3000, 3300, 3600, 3550, 3500],
),
(
avx512,
DispatchKind::Outside,
[9700, 9700, 9900, 10100, 9400, 8700],
),
(
apple,
DispatchKind::Team,
[3000, 3000, 2650, 2300, 1850, 1400],
),
(
apple,
DispatchKind::InPool,
[6000, 6000, 8500, 11000, 8250, 5500],
),
(
apple,
DispatchKind::Outside,
[14000, 14000, 16000, 18000, 18500, 19000],
),
(generic, DispatchKind::Team, [3000; 6]),
(generic, DispatchKind::InPool, [15000; 6]),
(generic, DispatchKind::Outside, [40000; 6]),
] {
for (k, want) in expected.into_iter().enumerate() {
assert_eq!(
costs.cost(kind, nz(k + 1)).as_nanos() as u64,
want,
"{kind:?} width {}",
k + 1
);
}
}
assert_eq!(avx512.cost(DispatchKind::Team, nz(8)), us_f(0.98));
assert_eq!(
avx512.cost(DispatchKind::Team, nz(7)),
Duration::from_nanos(940)
);
assert_eq!(
avx512.cost(DispatchKind::Team, nz(64)),
Duration::from_nanos(7_840)
);
assert_eq!(
avx512.cost(DispatchKind::InPool, nz(8)),
Duration::from_nanos(4666)
);
assert_eq!(
avx512.cost(DispatchKind::Outside, nz(8)),
Duration::from_nanos(11_600)
);
}
#[test]
fn min_is_right_for_a_rising_table_and_a_falling_one() {
let rising = WidthCosts::new([us(2), us(5), us(9)]);
assert_eq!(rising.min(), us(2));
let falling = WidthCosts::new([us(9), us(5), us(2)]);
assert_eq!(falling.min(), us(2));
let dip = WidthCosts::new([us(9), us(1), us(9)]);
assert_eq!(dip.min(), us(1));
assert_eq!(WidthCosts::flat(us(4)).min(), us(4));
}
#[test]
fn pays_matches_the_float_rule() {
assert_eq!(PAYS, PAYS_INT as f64);
for cost in [1u64, 700, 1700, 2800, 19_000] {
for sections in [1u32, 2, 5] {
let edge = 3 * u64::from(sections) * cost;
for est in [edge - 1, edge, edge + 1] {
let f = est as f64 >= PAYS * f64::from(sections) * cost as f64;
assert_eq!(
pays_at(
Duration::from_nanos(est),
sections,
Duration::from_nanos(cost)
),
f
);
}
}
}
}
#[test]
fn pays_is_the_multiple_of_the_cost() {
let d = us(2);
assert!(pays_at(us(6), 1, d));
assert!(!pays_at(Duration::from_nanos(5999), 1, d));
assert!(pays_at(us(12), 2, d));
assert!(!pays_at(us(11), 2, d));
assert!(!pays_at(us(1_000_000), 1, Duration::MAX));
assert!(pays_at(Duration::ZERO, 0, d));
}
#[test]
fn width_for_minimises_the_model() {
let budget = nz(6);
let d = |k: NonZeroUsize| us(2) * k.get() as u32;
assert_eq!(width_for_by(us(200), 1, budget, d).get(), 6);
assert_eq!(width_for_by(us(50), 1, budget, d).get(), 5);
assert_eq!(width_for_by(us(5), 1, budget, d).get(), 1);
assert!(width_for_by(us(200), 4, budget, d).get() < 6);
assert_eq!(
width_for_by(us(1_000_000), 1, budget, |_| Duration::MAX).get(),
1
);
assert_eq!(width_for_by(us(1_000_000), 1, nz(1), d).get(), 1);
}
#[test]
fn wide_values_take_the_128_bit_path_with_the_same_result() {
let big = Duration::from_nanos(u64::MAX / 3);
let huge = Duration::from_nanos(u64::MAX / 2);
let c = WidthCosts::from_anchors(&[(2, big), (64, huge)]).unwrap();
let (c0, c1) = (i128::from(u64::MAX / 3), i128::from(u64::MAX / 2));
for w in [3usize, 17, 33, 63] {
let want = c0 + (c1 - c0) * (w as i128 - 2) / 62;
assert_eq!(u128::from(c.ns_at(nz(w))), want as u128, "width {w}");
}
assert_eq!(c.ns_at(nz(200)), u64::MAX);
let want = u128::from(u64::MAX / 2) * 100 / 64;
assert_eq!(u128::from(c.ns_at(nz(100))), want.min(u128::from(u64::MAX)));
}
#[test]
fn falling_interpolation_truncates_toward_zero() {
let ns = Duration::from_nanos;
let c = WidthCosts::from_anchors(&[(2, ns(4)), (5, ns(0))]).unwrap();
assert_eq!(c.at(nz(3)), ns(3));
assert_eq!(c.at(nz(4)), ns(2));
}
}