use crate::compute::ComputeBackend;
#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub enum SplitRatio {
#[default]
Measured,
Fixed(f64),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum CoExecution {
#[default]
NativePrecision,
MatchGpuPrecision,
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct SplitPlan {
pub lanes: usize,
pub cpu_lanes: usize,
pub gpu_lanes: usize,
pub cpu_backend: ComputeBackend,
pub co_execution: CoExecution,
}
impl SplitPlan {
#[must_use]
pub fn for_lanes(lanes: usize, ratio: SplitRatio, co_execution: CoExecution) -> Self {
let cpu_backend = ComputeBackend::CpuMulti.resolve();
let gpu_fraction = if ComputeBackend::Gpu.is_available() {
match ratio {
SplitRatio::Measured => 0.0,
SplitRatio::Fixed(f) => f.clamp(0.0, 1.0),
}
} else {
0.0
};
let gpu_lanes = ((lanes as f64) * gpu_fraction).floor() as usize;
let gpu_lanes = gpu_lanes.min(lanes);
SplitPlan {
lanes,
cpu_lanes: lanes - gpu_lanes,
gpu_lanes,
cpu_backend,
co_execution,
}
}
#[must_use]
pub fn is_co_executing(&self) -> bool {
self.cpu_lanes > 0 && self.gpu_lanes > 0
}
#[must_use]
pub fn precision_note(&self) -> &'static str {
if !self.is_co_executing() {
if self.gpu_lanes > 0 {
"all lanes computed on the GPU in f32"
} else {
"all lanes computed on the CPU in f64"
}
} else {
match self.co_execution {
CoExecution::NativePrecision => {
"MIXED PRECISION: leading lanes are f64 (CPU), trailing lanes are f32 (GPU); \
accuracy is not uniform across the output array"
}
CoExecution::MatchGpuPrecision => {
"uniform f32-grade accuracy: CPU lanes deliberately rounded to f32 to match \
the GPU half"
}
}
}
}
pub fn harmonise_cpu_half(&self, cpu_values: &mut [f64]) {
if self.co_execution == CoExecution::MatchGpuPrecision {
for v in cpu_values.iter_mut() {
*v = (*v as f32) as f64;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn split_is_exhaustive_and_deterministic() {
for lanes in [0, 1, 63, 1024, 1_000_003] {
let a = SplitPlan::for_lanes(lanes, SplitRatio::Fixed(0.3), CoExecution::default());
let b = SplitPlan::for_lanes(lanes, SplitRatio::Fixed(0.3), CoExecution::default());
assert_eq!(a, b, "the same inputs must plan identically every time");
assert_eq!(
a.cpu_lanes + a.gpu_lanes,
lanes,
"every lane must be assigned exactly once"
);
}
}
#[test]
fn measured_ratio_is_cpu_only_on_this_hardware() {
let p = SplitPlan::for_lanes(1 << 20, SplitRatio::Measured, CoExecution::default());
assert_eq!(p.gpu_lanes, 0);
assert_eq!(p.cpu_lanes, 1 << 20);
assert!(!p.is_co_executing());
}
#[test]
fn no_gpu_means_every_lane_goes_to_the_cpu() {
let p = SplitPlan::for_lanes(1000, SplitRatio::Fixed(0.9), CoExecution::default());
if !ComputeBackend::Gpu.is_available() {
assert_eq!(p.gpu_lanes, 0, "no adapter must degrade to a CPU-only plan");
assert_eq!(p.cpu_lanes, 1000);
} else {
assert_eq!(p.gpu_lanes, 900);
assert_eq!(p.cpu_lanes, 100);
}
}
#[test]
fn ratio_is_clamped_not_wrapped() {
let hi = SplitPlan::for_lanes(100, SplitRatio::Fixed(5.0), CoExecution::default());
let lo = SplitPlan::for_lanes(100, SplitRatio::Fixed(-2.0), CoExecution::default());
assert!(hi.gpu_lanes <= 100);
assert_eq!(lo.gpu_lanes, 0);
assert_eq!(lo.cpu_lanes, 100);
}
#[test]
fn remainder_lane_goes_to_the_more_accurate_side() {
let p = SplitPlan::for_lanes(10, SplitRatio::Fixed(0.35), CoExecution::default());
if ComputeBackend::Gpu.is_available() {
assert_eq!(p.gpu_lanes, 3);
assert_eq!(p.cpu_lanes, 7);
}
}
#[test]
fn precision_note_names_the_mixed_case() {
let cpu_only = SplitPlan::for_lanes(100, SplitRatio::Fixed(0.0), CoExecution::default());
assert!(cpu_only.precision_note().contains("f64"));
let mixed = SplitPlan {
lanes: 100,
cpu_lanes: 50,
gpu_lanes: 50,
cpu_backend: ComputeBackend::Serial,
co_execution: CoExecution::NativePrecision,
};
assert!(
mixed.precision_note().contains("MIXED PRECISION"),
"the mixed-precision caveat must be stated, not implied"
);
let matched = SplitPlan {
co_execution: CoExecution::MatchGpuPrecision,
..mixed
};
assert!(matched.precision_note().contains("uniform f32"));
}
#[test]
fn harmonise_rounds_only_when_asked() {
let original = std::f64::consts::PI;
let mixed = SplitPlan {
lanes: 2,
cpu_lanes: 1,
gpu_lanes: 1,
cpu_backend: ComputeBackend::Serial,
co_execution: CoExecution::NativePrecision,
};
let mut v = [original];
mixed.harmonise_cpu_half(&mut v);
assert_eq!(v[0], original, "NativePrecision must not touch the values");
let matched = SplitPlan {
co_execution: CoExecution::MatchGpuPrecision,
..mixed
};
let mut v = [original];
matched.harmonise_cpu_half(&mut v);
assert_eq!(v[0], (original as f32) as f64);
assert_ne!(v[0], original, "MatchGpuPrecision must actually round");
}
}