#[allow(unused_imports)]
use super::functions::*;
#[allow(unused_imports)]
use super::functions_2::*;
#[allow(dead_code)]
pub struct LoadBalancePlan {
pub ranges: Vec<std::ops::Range<usize>>,
pub weights: Vec<f64>,
}
impl LoadBalancePlan {
#[allow(dead_code)]
pub fn num_workers(&self) -> usize {
self.ranges.len()
}
#[allow(dead_code)]
pub fn max_weight(&self) -> f64 {
self.weights
.iter()
.copied()
.fold(f64::NEG_INFINITY, f64::max)
}
#[allow(dead_code)]
pub fn imbalance_ratio(&self) -> f64 {
if self.weights.is_empty() {
return 1.0;
}
let total: f64 = self.weights.iter().sum();
let avg = total / self.weights.len() as f64;
if avg < 1e-15 {
return 1.0;
}
self.max_weight() / avg
}
}
#[allow(dead_code)]
pub struct WorkStealQueue<T> {
pub(super) items: std::collections::VecDeque<T>,
}
#[allow(dead_code)]
impl<T: Send> WorkStealQueue<T> {
pub fn new() -> Self {
Self {
items: std::collections::VecDeque::new(),
}
}
pub fn push(&mut self, task: T) {
self.items.push_back(task);
}
pub fn pop(&mut self) -> Option<T> {
self.items.pop_back()
}
pub fn steal(&mut self) -> Option<T> {
self.items.pop_front()
}
pub fn len(&self) -> usize {
self.items.len()
}
pub fn is_empty(&self) -> bool {
self.items.is_empty()
}
}
#[allow(dead_code)]
pub struct WorkGroupConfig {
pub preferred_size: usize,
pub max_size: usize,
pub min_size: usize,
}
impl WorkGroupConfig {
#[allow(dead_code)]
pub fn new(preferred_size: usize) -> Self {
Self {
preferred_size: preferred_size.max(1),
max_size: 1024,
min_size: 32,
}
}
#[allow(dead_code)]
pub fn cpu_default() -> Self {
let threads = rayon::current_num_threads().max(1);
Self {
preferred_size: 64,
max_size: 1024,
min_size: threads,
}
}
#[allow(dead_code)]
pub fn optimal_size(&self, total: usize) -> usize {
if total == 0 {
return self.min_size;
}
if total <= self.preferred_size {
return total.max(self.min_size).min(self.max_size);
}
let preferred_groups = total.div_ceil(self.preferred_size);
let preferred_waste = preferred_groups * self.preferred_size - total;
let preferred_waste_ratio = preferred_waste as f64 / total as f64;
if preferred_waste_ratio < 0.25 {
return self.preferred_size;
}
let mut best_size = self.preferred_size;
let mut best_waste = preferred_waste;
for candidate in (self.min_size..=self.max_size).step_by(self.min_size) {
let groups = total.div_ceil(candidate);
let waste = groups * candidate - total;
if waste < best_waste {
best_waste = waste;
best_size = candidate;
}
}
best_size
}
#[allow(dead_code)]
pub fn num_groups(&self, total: usize) -> usize {
let size = self.optimal_size(total);
total.div_ceil(size)
}
#[allow(dead_code)]
pub fn group_ranges(&self, total: usize) -> Vec<std::ops::Range<usize>> {
let size = self.optimal_size(total);
(0..total)
.step_by(size.max(1))
.map(|start| start..(start + size).min(total))
.collect()
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
#[allow(dead_code)]
pub enum LoadBalanceStrategy {
Static,
Weighted,
Guided,
}
pub struct WorkChunker {
pub n: usize,
pub chunk_size: usize,
}
impl WorkChunker {
pub fn new(n: usize) -> Self {
let threads = rayon::current_num_threads().max(1);
let chunk_size = n / threads + 1;
Self { n, chunk_size }
}
pub fn chunks(&self) -> Vec<std::ops::Range<usize>> {
let cs = self.chunk_size.max(1);
(0..self.n)
.step_by(cs)
.map(|start| start..(start + cs).min(self.n))
.collect()
}
}