mesh_sieve/accelerator/
reduction.rs

1//! Persistent scalar reductions over device-resident vectors.
2
3use super::{AcceleratorBackend, AcceleratorError, CpuBackend, DeviceBuffer, FvmScalar};
4
5/// Reusable one-scalar output allocation for vector reductions.
6pub struct DeviceReduction<T: FvmScalar, B: AcceleratorBackend> {
7    /// Maximum accepted input length.
8    pub(crate) input_len: usize,
9    /// Resident one-element reduction result.
10    pub(crate) result: B::Buffer<T>,
11    pub(crate) workspace: B::Buffer<T>,
12    pub(crate) backend_id: u64,
13}
14
15impl<T: FvmScalar, B: AcceleratorBackend> DeviceReduction<T, B> {
16    /// Allocate a reduction workspace for vectors of exactly `input_len` values.
17    pub fn new(backend: &B, input_len: usize) -> Result<Self, AcceleratorError> {
18        let result = backend
19            .allocate(1)
20            .map_err(|e| AcceleratorError::AllocationFailed {
21                bytes: std::mem::size_of::<T>(),
22                reason: e.to_string(),
23            })?;
24        let workspace_len = input_len.div_ceil(256).clamp(1, 4096);
25        let workspace =
26            backend
27                .allocate(workspace_len)
28                .map_err(|e| AcceleratorError::AllocationFailed {
29                    bytes: workspace_len.saturating_mul(std::mem::size_of::<T>()),
30                    reason: e.to_string(),
31                })?;
32        Ok(Self {
33            input_len,
34            result,
35            workspace,
36            backend_id: backend.identity(),
37        })
38    }
39
40    /// Download the most recently computed scalar.
41    pub fn download(&self, backend: &B) -> Result<T, AcceleratorError> {
42        let mut value = [T::zeroed()];
43        backend
44            .download(&self.result, &mut value)
45            .map_err(|e| AcceleratorError::DeviceTransferFailed(e.to_string()))?;
46        Ok(value[0])
47    }
48
49    pub(crate) fn validate(&self, found: usize) -> Result<(), AcceleratorError> {
50        if self.input_len == found {
51            Ok(())
52        } else {
53            Err(AcceleratorError::LengthMismatch {
54                expected: self.input_len,
55                found,
56            })
57        }
58    }
59}
60
61impl<T: FvmScalar> DeviceReduction<T, CpuBackend> {
62    /// Deterministic sum in input order.
63    pub fn sum(
64        &mut self,
65        input: &<CpuBackend as AcceleratorBackend>::Buffer<T>,
66    ) -> Result<T, AcceleratorError> {
67        self.validate(input.len())?;
68        let value = input
69            .as_slice()
70            .iter()
71            .fold(0.0, |sum, value| sum + value.to_f64());
72        self.result.as_mut_slice()[0] = T::from_f64(value);
73        Ok(T::from_f64(value))
74    }
75
76    /// Deterministic dot product in input order.
77    pub fn dot(
78        &mut self,
79        lhs: &<CpuBackend as AcceleratorBackend>::Buffer<T>,
80        rhs: &<CpuBackend as AcceleratorBackend>::Buffer<T>,
81    ) -> Result<T, AcceleratorError> {
82        self.validate(lhs.len())?;
83        self.validate(rhs.len())?;
84        let value = lhs
85            .as_slice()
86            .iter()
87            .zip(rhs.as_slice())
88            .fold(0.0, |sum, (&a, &b)| sum + a.to_f64() * b.to_f64());
89        self.result.as_mut_slice()[0] = T::from_f64(value);
90        Ok(T::from_f64(value))
91    }
92
93    /// Euclidean norm.
94    pub fn l2_norm(
95        &mut self,
96        input: &<CpuBackend as AcceleratorBackend>::Buffer<T>,
97    ) -> Result<T, AcceleratorError> {
98        self.validate(input.len())?;
99        let value = input
100            .as_slice()
101            .iter()
102            .fold(0.0, |sum, value| sum + value.to_f64() * value.to_f64())
103            .sqrt();
104        self.result.as_mut_slice()[0] = T::from_f64(value);
105        Ok(T::from_f64(value))
106    }
107
108    /// Maximum absolute value, or zero for an empty vector.
109    pub fn max_abs(
110        &mut self,
111        input: &<CpuBackend as AcceleratorBackend>::Buffer<T>,
112    ) -> Result<T, AcceleratorError> {
113        self.validate(input.len())?;
114        let value = input
115            .as_slice()
116            .iter()
117            .fold(0.0_f64, |max, value| max.max(value.to_f64().abs()));
118        self.result.as_mut_slice()[0] = T::from_f64(value);
119        Ok(T::from_f64(value))
120    }
121}
mesh_sieve/accelerator/reduction.rs

mesh_sieve/accelerator/
reduction.rs
