rnn/runtime/
budget.rs

1use super::{estimate_runtime_memory, RuntimeError, RuntimeEstimate, RuntimeProfile};
2use crate::model_config::TransformerConfig;
3
4#[derive(Clone, Copy, Debug, PartialEq, Eq)]
5pub struct BudgetFit {
6    pub fits: bool,
7    pub required_bytes: usize,
8    pub budget_bytes: usize,
9    pub headroom_bytes: usize,
10}
11
12pub fn check_runtime_budget(
13    config: &TransformerConfig,
14    profile: RuntimeProfile,
15    budget_bytes: usize,
16) -> Result<BudgetFit, RuntimeError> {
17    let estimate = estimate_runtime_memory(config, profile)?;
18    Ok(fit_from_estimate(estimate, budget_bytes))
19}
20
21pub fn fit_from_estimate(estimate: RuntimeEstimate, budget_bytes: usize) -> BudgetFit {
22    if estimate.total_bytes <= budget_bytes {
23        BudgetFit {
24            fits: true,
25            required_bytes: estimate.total_bytes,
26            budget_bytes,
27            headroom_bytes: budget_bytes - estimate.total_bytes,
28        }
29    } else {
30        BudgetFit {
31            fits: false,
32            required_bytes: estimate.total_bytes,
33            budget_bytes,
34            headroom_bytes: 0,
35        }
36    }
37}
rnn/runtime/budget.rs

rnn/runtime/
budget.rs
