use crate::test_support::KtstrTestEntry;
use crate::vmm::topology::Topology;
pub(crate) struct TestCandidate {
pub name: String,
pub features: u64,
pub estimated_secs: f64,
}
const SCHED_SHIFT: u32 = 0;
const REQ_FLAGS_SHIFT: u32 = 4;
const EXCL_FLAGS_SHIFT: u32 = 10;
const PROFILE_FLAGS_SHIFT: u32 = 16;
const CPU_BUCKET_SHIFT: u32 = 22;
const LLC_BUCKET_SHIFT: u32 = 27;
const SMT_SHIFT: u32 = 32;
const PERF_MODE_SHIFT: u32 = 33;
const HOST_ONLY_SHIFT: u32 = 34;
const EXPECT_ERR_SHIFT: u32 = 35;
const DURATION_SHIFT: u32 = 36;
const WORKERS_SHIFT: u32 = 39;
const GAUNTLET_SHIFT: u32 = 40;
const NAME_HASH_SHIFT: u32 = 41;
fn flag_bit(name: &str) -> Option<u32> {
crate::scenario::flags::ALL
.iter()
.position(|&n| n == name)
.map(|i| i as u32)
}
fn encode_flags(flags: &[&str]) -> u64 {
let mut mask = 0u64;
for &f in flags {
if let Some(bit) = flag_bit(f) {
mask |= 1 << bit;
}
}
mask
}
fn djb2_hash(name: &str) -> u32 {
let mut h: u32 = 5381;
for b in name.bytes() {
h = h.wrapping_mul(33).wrapping_add(b as u32);
}
h
}
fn cpu_bucket(total_cpus: u32) -> u64 {
match total_cpus {
0..=8 => 1 << 0,
9..=16 => 1 << 1,
17..=64 => 1 << 2,
65..=128 => 1 << 3,
_ => 1 << 4,
}
}
fn llc_bucket(num_llcs: u32) -> u64 {
match num_llcs {
0..=1 => 1 << 0,
2 => 1 << 1,
3..=4 => 1 << 2,
5..=8 => 1 << 3,
_ => 1 << 4,
}
}
fn duration_bucket(duration_secs: u64) -> u64 {
match duration_secs {
0..=2 => 1 << 0,
3..=10 => 1 << 1,
_ => 1 << 2,
}
}
fn workers_bucket(workers: u32) -> u64 {
if workers <= 2 { 0 } else { 1 }
}
pub(crate) fn extract_features(
entry: &KtstrTestEntry,
topo: &Topology,
active_flags: &[&str],
is_gauntlet: bool,
test_name: &str,
) -> u64 {
let mut bits = 0u64;
bits |= (1u64 << (djb2_hash(entry.scheduler.name) % 4)) << SCHED_SHIFT;
bits |= encode_flags(entry.required_flags) << REQ_FLAGS_SHIFT;
bits |= encode_flags(entry.excluded_flags) << EXCL_FLAGS_SHIFT;
bits |= encode_flags(active_flags) << PROFILE_FLAGS_SHIFT;
bits |= cpu_bucket(topo.total_cpus()) << CPU_BUCKET_SHIFT;
bits |= llc_bucket(topo.num_llcs()) << LLC_BUCKET_SHIFT;
if topo.threads_per_core > 1 {
bits |= 1 << SMT_SHIFT;
}
if entry.performance_mode {
bits |= 1 << PERF_MODE_SHIFT;
}
if entry.host_only {
bits |= 1 << HOST_ONLY_SHIFT;
}
if entry.expect_err {
bits |= 1 << EXPECT_ERR_SHIFT;
}
bits |= duration_bucket(entry.duration.as_secs()) << DURATION_SHIFT;
bits |= workers_bucket(entry.workers_per_cgroup) << WORKERS_SHIFT;
if is_gauntlet {
bits |= 1 << GAUNTLET_SHIFT;
}
bits |= (1u64 << (djb2_hash(test_name) % 4)) << NAME_HASH_SHIFT;
bits
}
pub(crate) fn estimate_duration(entry: &KtstrTestEntry, topo: &Topology) -> f64 {
let duration_secs = entry.duration.as_secs();
if entry.host_only {
return (duration_secs + 2) as f64;
}
let cpus = topo.total_cpus() as u64;
let boot_overhead = 10 + cpus.saturating_sub(16) / 10;
let perf_overhead: u64 = if entry.performance_mode { 3 } else { 0 };
(boot_overhead + duration_secs + 2 + perf_overhead) as f64
}
pub(crate) fn select(candidates: &[TestCandidate], budget_secs: f64) -> Vec<usize> {
if candidates.is_empty() || budget_secs <= 0.0 {
return Vec::new();
}
let n = candidates.len();
let mut selected = Vec::new();
let mut used = vec![false; n];
let mut covered: u64 = 0;
let mut remaining_budget = budget_secs;
loop {
let mut best_idx: Option<usize> = None;
let mut best_ratio: f64 = 0.0;
let mut best_name: &str = "";
for (i, c) in candidates.iter().enumerate() {
if used[i] || c.estimated_secs > remaining_budget {
continue;
}
let marginal = (c.features & !covered).count_ones() as f64;
if marginal == 0.0 {
continue;
}
let ratio = if c.estimated_secs > 0.0 {
marginal / c.estimated_secs
} else {
marginal * 1e6 };
let better = ratio > best_ratio || (ratio == best_ratio && c.name.as_str() < best_name);
if better {
best_ratio = ratio;
best_idx = Some(i);
best_name = &c.name;
}
}
match best_idx {
Some(i) => {
selected.push(i);
used[i] = true;
covered |= candidates[i].features;
remaining_budget -= candidates[i].estimated_secs;
}
None => break,
}
}
selected.sort_unstable();
selected
}
pub(crate) struct SelectionStats {
pub selected: usize,
pub total: usize,
pub budget_used: f64,
pub budget_total: f64,
pub bits_covered: u32,
pub bits_possible: u32,
}
pub(crate) fn selection_stats(
candidates: &[TestCandidate],
selected: &[usize],
budget_secs: f64,
) -> SelectionStats {
let mut covered: u64 = 0;
let mut budget_used = 0.0;
for &i in selected {
covered |= candidates[i].features;
budget_used += candidates[i].estimated_secs;
}
let mut all_features: u64 = 0;
for c in candidates {
all_features |= c.features;
}
SelectionStats {
selected: selected.len(),
total: candidates.len(),
budget_used,
budget_total: budget_secs,
bits_covered: covered.count_ones(),
bits_possible: all_features.count_ones(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn flag_bit_known() {
assert_eq!(flag_bit("llc"), Some(0));
assert_eq!(flag_bit("borrow"), Some(1));
assert_eq!(flag_bit("steal"), Some(2));
assert_eq!(flag_bit("rebal"), Some(3));
assert_eq!(flag_bit("reject-pin"), Some(4));
assert_eq!(flag_bit("no-ctrl"), Some(5));
}
#[test]
fn flag_bit_unknown() {
assert_eq!(flag_bit("nonexistent"), None);
}
#[test]
fn encode_flags_empty() {
assert_eq!(encode_flags(&[]), 0);
}
#[test]
fn encode_flags_single() {
assert_eq!(encode_flags(&["borrow"]), 0b10);
}
#[test]
fn encode_flags_multiple() {
let mask = encode_flags(&["llc", "steal"]);
assert_eq!(mask, 0b101); }
#[test]
fn cpu_bucket_one_hot() {
assert_eq!(cpu_bucket(1), 1 << 0);
assert_eq!(cpu_bucket(8), 1 << 0);
assert_eq!(cpu_bucket(9), 1 << 1);
assert_eq!(cpu_bucket(16), 1 << 1);
assert_eq!(cpu_bucket(17), 1 << 2);
assert_eq!(cpu_bucket(64), 1 << 2);
assert_eq!(cpu_bucket(65), 1 << 3);
assert_eq!(cpu_bucket(128), 1 << 3);
assert_eq!(cpu_bucket(129), 1 << 4);
assert_eq!(cpu_bucket(252), 1 << 4);
}
#[test]
fn cpu_bucket_no_shared_bits() {
let vals = [
cpu_bucket(4),
cpu_bucket(12),
cpu_bucket(32),
cpu_bucket(96),
cpu_bucket(200),
];
for (i, &a) in vals.iter().enumerate() {
assert_eq!(a.count_ones(), 1, "bucket {i} not one-hot: {a:#b}");
for &b in &vals[i + 1..] {
assert_eq!(a & b, 0, "buckets share bits: {a:#b} & {b:#b}");
}
}
}
#[test]
fn llc_bucket_one_hot() {
assert_eq!(llc_bucket(1), 1 << 0);
assert_eq!(llc_bucket(2), 1 << 1);
assert_eq!(llc_bucket(3), 1 << 2);
assert_eq!(llc_bucket(4), 1 << 2);
assert_eq!(llc_bucket(5), 1 << 3);
assert_eq!(llc_bucket(8), 1 << 3);
assert_eq!(llc_bucket(9), 1 << 4);
assert_eq!(llc_bucket(15), 1 << 4);
}
#[test]
fn llc_bucket_no_shared_bits() {
let vals = [
llc_bucket(1),
llc_bucket(2),
llc_bucket(3),
llc_bucket(7),
llc_bucket(10),
];
for (i, &a) in vals.iter().enumerate() {
assert_eq!(a.count_ones(), 1);
for &b in &vals[i + 1..] {
assert_eq!(a & b, 0);
}
}
}
#[test]
fn duration_bucket_one_hot() {
assert_eq!(duration_bucket(0), 1 << 0);
assert_eq!(duration_bucket(2), 1 << 0);
assert_eq!(duration_bucket(3), 1 << 1);
assert_eq!(duration_bucket(10), 1 << 1);
assert_eq!(duration_bucket(11), 1 << 2);
}
#[test]
fn duration_bucket_no_shared_bits() {
let vals = [duration_bucket(1), duration_bucket(5), duration_bucket(20)];
for (i, &a) in vals.iter().enumerate() {
assert_eq!(a.count_ones(), 1);
for &b in &vals[i + 1..] {
assert_eq!(a & b, 0);
}
}
}
#[test]
fn workers_bucket_values() {
assert_eq!(workers_bucket(1), 0);
assert_eq!(workers_bucket(2), 0);
assert_eq!(workers_bucket(3), 1);
assert_eq!(workers_bucket(32), 1);
}
#[test]
fn extract_features_base_test() {
use crate::test_support::KtstrTestEntry;
let entry = KtstrTestEntry::DEFAULT;
let topo = Topology {
sockets: 1,
cores_per_socket: 2,
threads_per_core: 1,
};
let features = extract_features(&entry, &topo, &[], false, "basic_test");
assert_eq!(features & (1 << GAUNTLET_SHIFT), 0);
assert_eq!(features & (1 << SMT_SHIFT), 0);
assert_eq!(features & (1 << PERF_MODE_SHIFT), 0);
}
#[test]
fn extract_features_smt_set() {
let entry = KtstrTestEntry::DEFAULT;
let topo = Topology {
sockets: 2,
cores_per_socket: 2,
threads_per_core: 2,
};
let features = extract_features(&entry, &topo, &[], false, "smt_test");
assert_ne!(features & (1 << SMT_SHIFT), 0);
}
#[test]
fn extract_features_gauntlet() {
let entry = KtstrTestEntry::DEFAULT;
let topo = Topology {
sockets: 4,
cores_per_socket: 4,
threads_per_core: 2,
};
let features = extract_features(&entry, &topo, &["llc", "borrow"], true, "gauntlet_test");
assert_ne!(features & (1 << GAUNTLET_SHIFT), 0);
let profile_mask = (features >> PROFILE_FLAGS_SHIFT) & 0x3F;
assert_ne!(profile_mask & 0b01, 0); assert_ne!(profile_mask & 0b10, 0); }
#[test]
fn extract_features_required_flags() {
let entry = KtstrTestEntry {
required_flags: &["borrow", "rebal"],
..KtstrTestEntry::DEFAULT
};
let topo = Topology {
sockets: 1,
cores_per_socket: 2,
threads_per_core: 1,
};
let features = extract_features(&entry, &topo, &[], false, "flag_test");
let req_mask = (features >> REQ_FLAGS_SHIFT) & 0x3F;
assert_ne!(req_mask & (1 << 1), 0); assert_ne!(req_mask & (1 << 3), 0); }
#[test]
fn estimate_duration_small_topo() {
let entry = KtstrTestEntry {
duration: std::time::Duration::from_secs(2),
..KtstrTestEntry::DEFAULT
};
let topo = Topology {
sockets: 1,
cores_per_socket: 2,
threads_per_core: 1,
};
assert_eq!(estimate_duration(&entry, &topo), 14.0);
}
#[test]
fn estimate_duration_large_topo() {
let entry = KtstrTestEntry {
duration: std::time::Duration::from_secs(5),
..KtstrTestEntry::DEFAULT
};
let topo = Topology {
sockets: 14,
cores_per_socket: 9,
threads_per_core: 2,
};
assert_eq!(estimate_duration(&entry, &topo), 40.0);
}
#[test]
fn estimate_duration_performance_mode() {
let entry = KtstrTestEntry {
duration: std::time::Duration::from_secs(2),
performance_mode: true,
..KtstrTestEntry::DEFAULT
};
let topo = Topology {
sockets: 1,
cores_per_socket: 2,
threads_per_core: 1,
};
assert_eq!(estimate_duration(&entry, &topo), 17.0);
}
#[test]
fn select_empty() {
assert!(select(&[], 100.0).is_empty());
}
#[test]
fn select_zero_budget() {
let candidates = vec![TestCandidate {
name: "t1".into(),
features: 0xFF,
estimated_secs: 10.0,
}];
assert!(select(&candidates, 0.0).is_empty());
}
#[test]
fn select_single_fits() {
let candidates = vec![TestCandidate {
name: "t1".into(),
features: 0xFF,
estimated_secs: 10.0,
}];
let sel = select(&candidates, 20.0);
assert_eq!(sel, vec![0]);
}
#[test]
fn select_single_too_expensive() {
let candidates = vec![TestCandidate {
name: "t1".into(),
features: 0xFF,
estimated_secs: 30.0,
}];
let sel = select(&candidates, 20.0);
assert!(sel.is_empty());
}
#[test]
fn select_prefers_coverage_per_second() {
let candidates = vec![
TestCandidate {
name: "t1".into(),
features: 0b1111,
estimated_secs: 20.0,
},
TestCandidate {
name: "t2".into(),
features: 0b110000,
estimated_secs: 5.0,
},
];
let sel = select(&candidates, 25.0);
assert_eq!(sel, vec![0, 1]); }
#[test]
fn select_budget_constraint() {
let candidates = vec![
TestCandidate {
name: "t1".into(),
features: 0b1111,
estimated_secs: 15.0,
},
TestCandidate {
name: "t2".into(),
features: 0b110000,
estimated_secs: 15.0,
},
];
let sel = select(&candidates, 20.0);
assert_eq!(sel, vec![0]);
}
#[test]
fn select_marginal_coverage_decreases() {
let candidates = vec![
TestCandidate {
name: "t1".into(),
features: 0b1111,
estimated_secs: 5.0,
},
TestCandidate {
name: "t2".into(),
features: 0b1111,
estimated_secs: 5.0,
},
TestCandidate {
name: "t3".into(),
features: 0b110000,
estimated_secs: 5.0,
},
];
let sel = select(&candidates, 100.0);
assert_eq!(sel.len(), 2);
assert!(sel.contains(&2)); }
#[test]
fn select_sorted_output() {
let candidates = vec![
TestCandidate {
name: "t1".into(),
features: 0b01,
estimated_secs: 10.0,
},
TestCandidate {
name: "t2".into(),
features: 0b10,
estimated_secs: 5.0,
},
];
let sel = select(&candidates, 100.0);
assert_eq!(sel, vec![0, 1]);
}
#[test]
fn djb2_hash_one_hot_4() {
for name in &["eevdf", "scx_mitosis", "scx_rusty", "", "x"] {
let one_hot = 1u64 << (djb2_hash(name) % 4);
assert_eq!(one_hot.count_ones(), 1);
assert!(one_hot < 16);
}
}
#[test]
fn djb2_hash_different_names_differ() {
let h1 = djb2_hash("eevdf");
let h2 = djb2_hash("scx_mitosis");
assert_ne!(h1, h2);
}
#[test]
fn select_different_features_both_selected() {
let candidates = vec![
TestCandidate {
name: "sched_a_test1".into(),
features: 0b0001,
estimated_secs: 5.0,
},
TestCandidate {
name: "sched_b_test1".into(),
features: 0b0010 | (1 << 3),
estimated_secs: 5.0,
},
];
let sel = select(&candidates, 100.0);
assert_eq!(sel, vec![0, 1]);
}
#[test]
fn select_tie_broken_by_name() {
let candidates = vec![
TestCandidate {
name: "zzz".into(),
features: 0b1111,
estimated_secs: 10.0,
},
TestCandidate {
name: "aaa".into(),
features: 0b1111,
estimated_secs: 10.0,
},
];
let sel = select(&candidates, 15.0);
assert_eq!(sel, vec![1]); }
#[test]
fn selection_stats_basic() {
let candidates = vec![
TestCandidate {
name: "t1".into(),
features: 0b0011,
estimated_secs: 5.0,
},
TestCandidate {
name: "t2".into(),
features: 0b1100,
estimated_secs: 5.0,
},
TestCandidate {
name: "t3".into(),
features: 0b1111,
estimated_secs: 5.0,
},
];
let sel = vec![0, 1];
let stats = selection_stats(&candidates, &sel, 100.0);
assert_eq!(stats.selected, 2);
assert_eq!(stats.total, 3);
assert_eq!(stats.budget_used, 10.0);
assert_eq!(stats.budget_total, 100.0);
assert_eq!(stats.bits_covered, 4);
assert_eq!(stats.bits_possible, 4);
}
#[test]
fn estimate_duration_host_only() {
let entry = KtstrTestEntry {
duration: std::time::Duration::from_secs(5),
host_only: true,
..KtstrTestEntry::DEFAULT
};
let topo = Topology {
sockets: 14,
cores_per_socket: 9,
threads_per_core: 2,
};
assert_eq!(estimate_duration(&entry, &topo), 7.0);
}
#[test]
fn select_zero_cost_selected_first() {
let candidates = vec![
TestCandidate {
name: "free".into(),
features: 0b0001,
estimated_secs: 0.0,
},
TestCandidate {
name: "expensive".into(),
features: 0b0010,
estimated_secs: 100.0,
},
];
let sel = select(&candidates, 50.0);
assert_eq!(sel, vec![0]);
}
}