laddu_runtime/cpu/
prepared.rs

1use std::{mem::size_of, sync::Arc};
2
3use laddu_data::{
4    BatchLayout,
5    data::{CacheStorage, Dataset},
6    schema::Precision as DataPrecision,
7};
8use laddu_expr::{ExprId, ValueKind};
9use laddu_memory::{FootprintOverflow, MemoryFootprint};
10use num::complex::Complex64;
11
12use super::cache::{CachedFactorSlot, CachedSlot, CachedSolveRowSlot};
13use super::{
14    CpuCachedBatch, CpuCachedDataset, CpuPlan, CpuPreparedDataset, DynamicLu, RuntimeError,
15    RuntimeResult,
16};
17use crate::execution::Execution;
18use crate::preparation::{DatasetPreparation, LocalDatasetStats};
19
20impl CpuPlan {
21    /// Materializes all event-dependent caches for a dataset.
22    ///
23    /// # Errors
24    ///
25    /// Returns [`RuntimeError`] when the dataset cannot be read, a batch schema
26    /// is incompatible, cache construction fails, or a matrix is singular.
27    pub fn cache_dataset(&self, dataset: &Dataset) -> RuntimeResult<CpuCachedDataset> {
28        self.cache_dataset_with_plan(dataset, dataset.read_plan())
29    }
30
31    /// Estimates retained compiled-cache bytes for `events`.
32    pub fn cache_memory_estimate(&self, events: usize) -> usize {
33        self.cache_memory_footprint()
34            .map(|footprint| usize::try_from(footprint.peak_bytes(events)).unwrap_or(usize::MAX))
35            .unwrap_or(usize::MAX)
36    }
37
38    fn cache_memory_footprint(&self) -> Result<MemoryFootprint, FootprintOverflow> {
39        let mut fixed = MemoryFootprint::fixed(0);
40        for (count, bytes) in [
41            (self.cache_plan.entries().len(), size_of::<CachedSlot>()),
42            (self.factor_matrices.len(), size_of::<CachedFactorSlot>()),
43            (self.solve_row_keys.len(), size_of::<CachedSolveRowSlot>()),
44            (self.cache_plan.entries().len(), size_of::<ExprId>()),
45            (self.factor_matrices.len(), size_of::<ExprId>()),
46            (
47                self.solve_row_keys.len(),
48                size_of::<(ExprId, usize, usize)>(),
49            ),
50        ] {
51            fixed = fixed.checked_add(
52                MemoryFootprint::from_usize_checked(bytes, 0)?.checked_scale_usize(count)?,
53            )?;
54        }
55
56        let mut per_event = MemoryFootprint::per_event(size_of::<f64>() as u64);
57        for entry in self.cache_plan.entries() {
58            let bytes = match entry.value_kind() {
59                ValueKind::Real => MemoryFootprint::per_event(size_of::<f64>() as u64),
60                ValueKind::Complex => MemoryFootprint::per_event(size_of::<Complex64>() as u64),
61                ValueKind::Vector { len } => {
62                    MemoryFootprint::per_event(size_of::<Complex64>() as u64)
63                        .checked_scale_usize(len)?
64                }
65                ValueKind::Matrix { rows, cols } => {
66                    MemoryFootprint::per_event(size_of::<Complex64>() as u64)
67                        .checked_scale_usize(rows)?
68                        .checked_scale_usize(cols)?
69                }
70            };
71            per_event = per_event.checked_add(bytes)?;
72        }
73        for (_, dimension) in &self.factor_matrices {
74            let bytes = MemoryFootprint::per_event(size_of::<DynamicLu>() as u64)
75                .checked_add(
76                    MemoryFootprint::per_event(size_of::<Complex64>() as u64)
77                        .checked_scale_usize(*dimension)?
78                        .checked_scale_usize(*dimension)?,
79                )?
80                .checked_add(
81                    MemoryFootprint::per_event(size_of::<usize>() as u64)
82                        .checked_scale_usize(*dimension)?,
83                )?;
84            per_event = per_event.checked_add(bytes)?;
85        }
86        for (_, _, dimension) in &self.solve_row_keys {
87            per_event = per_event.checked_add(
88                MemoryFootprint::per_event(size_of::<Complex64>() as u64)
89                    .checked_scale_usize(*dimension)?,
90            )?;
91        }
92        fixed.checked_add(per_event)
93    }
94
95    fn cache_dataset_with_plan(
96        &self,
97        dataset: &Dataset,
98        read_plan: laddu_data::io::ReadPlan,
99    ) -> RuntimeResult<CpuCachedDataset> {
100        let mut batches = Vec::new();
101        let mut sum_weights = 0.0;
102        for batch in dataset
103            .stream_with_plan(read_plan)
104            .map_err(|err| RuntimeError::Data(err.to_string()))?
105        {
106            let batch = batch.map_err(|err| RuntimeError::Data(err.to_string()))?;
107            let cached = CpuCachedBatch::from_cache(self.cache_event_batch(&batch)?);
108            sum_weights += cached.sum_weights();
109            batches.push(cached);
110        }
111        Ok(CpuCachedDataset::from_parts(batches, sum_weights))
112    }
113
114    /// Prepares a dataset according to its cache-storage policy.
115    ///
116    /// # Errors
117    ///
118    /// Returns [`RuntimeError`] when dataset reading or cache construction
119    /// fails, or another distributed worker reports failure.
120    ///
121    /// # Panics
122    ///
123    /// Panics if successful preparation planning fails to record its staging
124    /// memory decision, which would violate the planner invariant.
125    pub fn prepare_dataset(
126        &self,
127        execution: &Execution,
128        dataset: &Dataset,
129    ) -> RuntimeResult<CpuPreparedDataset> {
130        let preparation = DatasetPreparation::new(execution, dataset);
131        let mut planning = preparation.runtime_plan()?;
132        let initial_read_plan = planning.read_plan();
133        let schema = dataset
134            .schema()
135            .map_err(|error| RuntimeError::Data(error.to_string()))?;
136        let source_footprint = BatchLayout::from_schema(&schema)
137            .schema_footprint(DataPrecision::F64)
138            .map_err(|error| RuntimeError::Data(format!("source working-set overflow: {error}")))?;
139        let cache_footprint = self
140            .cache_memory_footprint()
141            .map_err(|error| RuntimeError::Data(format!("cache working-set overflow: {error}")))?;
142        let source_bytes_per_event =
143            usize::try_from(source_footprint.bytes_per_event).unwrap_or(usize::MAX);
144        let cache_zero = usize::try_from(cache_footprint.fixed_bytes).unwrap_or(usize::MAX);
145        let cache_bytes_per_event =
146            usize::try_from(cache_footprint.bytes_per_event).unwrap_or(usize::MAX);
147        let local_event_limit = planning.event_limit();
148        let host_remaining = execution.host_memory().remaining();
149        let resident_plan = resident_cache_plan(
150            cache_zero,
151            cache_bytes_per_event,
152            source_bytes_per_event.saturating_mul(2),
153            local_event_limit,
154            usize::try_from(host_remaining).unwrap_or(usize::MAX),
155        );
156        let source_staging = source_footprint
157            .checked_scale(2)
158            .and_then(|footprint| {
159                MemoryFootprint::fixed(footprint.bytes_per_event)
160                    .checked_add(MemoryFootprint::per_event(cache_footprint.bytes_per_event))
161            })
162            .map_err(|error| RuntimeError::Data(format!("source working-set overflow: {error}")))?;
163        let minimum = MemoryFootprint::fixed(cache_footprint.fixed_bytes)
164            .checked_add(source_staging)
165            .map_err(|error| RuntimeError::Data(format!("cache working-set overflow: {error}")))?
166            .peak_bytes(local_event_limit);
167        let resident_bytes = resident_plan
168            .map(|(bytes, _)| u64::try_from(bytes).unwrap_or(u64::MAX))
169            .unwrap_or(minimum);
170        planning.select_storage(
171            dataset.memory_policy(),
172            "host",
173            resident_plan.is_some(),
174            resident_bytes,
175            host_remaining,
176        )?;
177        let requested_storage = planning.storage();
178        planning.reserve_storage(Some(execution.host_memory()), || {
179            RuntimeError::Data("CPU execution has no host memory pool".into())
180        })?;
181        let persistent_lease = planning.take_memory_lease();
182        let available_for_batch = execution.host_memory().remaining();
183        // Sources may hold the current decoded batch plus one bounded
184        // prefetched batch. A resident cache is already covered by its
185        // persistent lease; streaming additionally needs one transient cache.
186        let transient_footprint = if requested_storage == CacheStorage::Streaming {
187            cache_footprint
188                .checked_add(source_footprint.checked_scale(2).map_err(|error| {
189                    RuntimeError::Data(format!("source working-set overflow: {error}"))
190                })?)
191                .map_err(|error| {
192                    RuntimeError::Data(format!("cache working-set overflow: {error}"))
193                })?
194        } else {
195            source_footprint.checked_scale(2).map_err(|error| {
196                RuntimeError::Data(format!("source working-set overflow: {error}"))
197            })?
198        };
199        let decision = planning.fit_staging(
200            "CPU prepared dataset",
201            transient_footprint,
202            available_for_batch,
203            if requested_storage == CacheStorage::Resident {
204                "resident"
205            } else {
206                "streaming"
207            },
208        )?;
209        planning.clamp_read_plan(initial_read_plan.chunk_size, decision.chunk_events);
210        let decision = planning
211            .take_decisions()
212            .into_iter()
213            .next()
214            .expect("CPU preparation records one staging decision");
215        execution.record_memory_decision(decision.clone());
216        let read_plan = planning.read_plan();
217        match requested_storage {
218            CacheStorage::Resident => {
219                let dataset =
220                    preparation.coordinate(self.cache_dataset_with_plan(dataset, read_plan))?;
221                let stats = preparation.finish_stats(
222                    LocalDatasetStats::new(
223                        dataset.len(),
224                        dataset.batches().len(),
225                        dataset.sum_weights(),
226                    ),
227                    dataset.resident_bytes(),
228                    requested_storage,
229                );
230                let memory_lease = persistent_lease.ok_or_else(|| {
231                    RuntimeError::Data(
232                        "resident dataset preparation did not reserve host memory".into(),
233                    )
234                })?;
235                Ok(CpuPreparedDataset::Resident {
236                    dataset: Arc::new(dataset),
237                    stats,
238                    memory_lease,
239                })
240            }
241            CacheStorage::Streaming => {
242                let local = preparation.scan()?;
243                Ok(CpuPreparedDataset::Streaming {
244                    dataset: dataset.clone(),
245                    stats: preparation.finish_stats(local, 0, requested_storage),
246                    read_plan,
247                    transient_bytes: decision.estimated_peak_bytes,
248                })
249            }
250        }
251    }
252}
253
254pub(super) fn resident_cache_plan(
255    fixed_per_batch: usize,
256    cache_bytes_per_event: usize,
257    source_bytes_per_event: usize,
258    events: usize,
259    available: usize,
260) -> Option<(usize, usize)> {
261    if events == 0 {
262        return Some((fixed_per_batch, 1));
263    }
264    let event_cache = cache_bytes_per_event.checked_mul(events)?;
265    let minimum = event_cache
266        .checked_add(fixed_per_batch)?
267        .checked_add(source_bytes_per_event)?;
268    if minimum > available {
269        return None;
270    }
271    let mut chunk = events;
272    for _ in 0..16 {
273        let batches = events.saturating_add(chunk - 1) / chunk;
274        let resident = event_cache.checked_add(fixed_per_batch.checked_mul(batches)?)?;
275        let next = available
276            .saturating_sub(resident)
277            .checked_div(source_bytes_per_event.max(1))?
278            .min(events);
279        if next == 0 {
280            return None;
281        }
282        if next == chunk {
283            return Some((resident, chunk));
284        }
285        chunk = next;
286    }
287    let batches = events.saturating_add(chunk - 1) / chunk;
288    let resident = event_cache.checked_add(fixed_per_batch.checked_mul(batches)?)?;
289    (resident.checked_add(source_bytes_per_event.checked_mul(chunk)?)? <= available)
290        .then_some((resident, chunk))
291}
laddu_runtime/cpu/prepared.rs

laddu_runtime/cpu/
prepared.rs
