tropical_gemm/core/
workspace.rs

1use crate::types::TropicalScalar;
2
3/// Reusable CPU packing storage, shared by sequential calls through an exclusive borrow.
4///
5/// Each parallel output task gets its own pair of buffers. Storage grows to the
6/// largest panels and task count used so far; call [`Self::clear`] to release it.
7/// This reuses packing allocations, not result matrices or Rayon task metadata.
8/// The scalar parameter permits reuse across semirings with the same scalar type.
9///
10/// ```
11/// use tropical_gemm::{GemmWorkspace, TropicalGemm, TropicalMaxPlus};
12/// let mut workspace = GemmWorkspace::<f32>::new();
13/// let mut c = vec![TropicalMaxPlus(0.0); 4];
14/// for _ in 0..3 {
15///     TropicalGemm::new(2, 2, 2).execute_with_workspace(
16///         &[1.0; 4], 2, &[2.0; 4], 2, &mut c, 2, &mut workspace,
17///     );
18/// }
19/// assert!(c.iter().all(|v| v.0 == 3.0));
20/// ```
21#[derive(Debug)]
22pub struct GemmWorkspace<T: TropicalScalar> {
23    buffers: Vec<PackingBuffers<T>>,
24}
25
26impl<T: TropicalScalar> Default for GemmWorkspace<T> {
27    fn default() -> Self {
28        Self {
29            buffers: Vec::new(),
30        }
31    }
32}
33
34impl<T: TropicalScalar> GemmWorkspace<T> {
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    /// Bytes reserved by scalar packing buffers (excluding task metadata).
40    pub fn capacity_bytes(&self) -> usize {
41        self.buffers
42            .iter()
43            .map(|b| {
44                b.a.capacity() * std::mem::size_of::<T>()
45                    + b.b.capacity() * std::mem::size_of::<T>()
46            })
47            .sum()
48    }
49
50    /// Release all packing buffers. The next call will allocate them again.
51    pub fn clear(&mut self) {
52        self.buffers.clear();
53    }
54
55    pub(super) fn tasks(&mut self, count: usize) -> &mut [PackingBuffers<T>] {
56        self.buffers
57            .resize_with(count.max(self.buffers.len()), PackingBuffers::default);
58        &mut self.buffers[..count]
59    }
60}
61
62#[derive(Debug)]
63pub(super) struct PackingBuffers<T: TropicalScalar> {
64    pub a: Vec<T>,
65    pub b: Vec<T>,
66}
67
68impl<T: TropicalScalar> Default for PackingBuffers<T> {
69    fn default() -> Self {
70        Self {
71            a: Vec::new(),
72            b: Vec::new(),
73        }
74    }
75}
76
77impl<T: TropicalScalar> PackingBuffers<T> {
78    pub fn prepare(&mut self, a_len: usize, b_len: usize) {
79        // Packing overwrites every used element, including edge padding. Old
80        // contents need no clearing, and unused capacity never enters a kernel.
81        if self.a.len() < a_len {
82            self.a.resize(a_len, T::scalar_zero());
83        }
84        if self.b.len() < b_len {
85            self.b.resize(b_len, T::scalar_zero());
86        }
87    }
88}
tropical_gemm/core/workspace.rs

tropical_gemm/core/
workspace.rs
