Coverage Report

Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/home/noah/src/trueno/src/tuner/mod.rs
Line
Count
Source
1
//! ML-Based ComputeBrick Tuner
2
//!
3
//! Implements learned cost models for kernel selection and throughput prediction.
4
//! See: `docs/specifications/ml-tuner-bricks.md`
5
//!
6
//! # Architecture
7
//!
8
//! ```text
9
//! BrickProfiler → FeatureExtractor → TunerModel → Recommendations
10
//! ```
11
//!
12
//! # Scientific Foundations
13
//!
14
//! - Chen et al. (2018). "TVM: An Automated End-to-End Optimizing Compiler." OSDI '18.
15
//! - Williams et al. (2009). "Roofline: An Insightful Visual Performance Model." CACM.
16
//! - Friedman (2001). "Greedy Function Approximation: A Gradient Boosting Machine."
17
//!
18
//! # Example
19
//!
20
//! ```rust,ignore
21
//! use trueno::tuner::{BrickTuner, TunerFeatures};
22
//!
23
//! let features = TunerFeatures::builder()
24
//!     .model_params_b(1.5)
25
//!     .hidden_dim(1536)
26
//!     .batch_size(4)
27
//!     .quant_type(QuantType::Q4K)
28
//!     .build();
29
//!
30
//! let tuner = BrickTuner::load_or_default();
31
//! let recommendation = tuner.recommend(&features);
32
//! println!("Predicted: {} tok/s", recommendation.throughput.predicted_tps);
33
//! ```
34
35
// Submodules
36
mod brick_tuner;
37
mod data_collector;
38
pub mod error;
39
mod evolution;
40
mod features;
41
pub(crate) mod helpers;
42
mod models;
43
pub mod pretrained;
44
mod types;
45
46
// Re-export all public types
47
pub use brick_tuner::{BrickTuner, ExperimentSuggestion, TunerRecommendation};
48
pub use data_collector::{
49
    ConceptDriftStatus, TrainingSample, TrainingStats, TunerDataCollector, UserFeedback,
50
};
51
pub use error::TunerError;
52
pub use evolution::{CalibrationResult, KernelArm, KernelBandit, OnlineLearner};
53
pub use features::{FeatureExtractor, RunConfig, TunerFeatures, TunerFeaturesBuilder};
54
pub use models::{
55
    BottleneckClassifier, BottleneckPrediction, KernelClassifier, KernelRecommendation,
56
    ThroughputPrediction, ThroughputRegressor,
57
};
58
pub use types::{BottleneckClass, KernelType, QuantType};
59
60
// Re-export helpers for tests (crate-internal)
61
#[cfg(test)]
62
pub(crate) use helpers::{chrono_lite_now, pad_right};
63
64
// ============================================================================
65
// BrickProfiler Integration
66
// ============================================================================
67
68
use crate::brick::BrickProfiler;
69
70
impl BrickProfiler {
71
    /// Get ML-based tuning recommendations.
72
    ///
73
    /// Extracts features from current profile and returns recommendations.
74
0
    pub fn get_tuner_recommendations(&self, config: &RunConfig) -> Option<TunerRecommendation> {
75
0
        if !self.is_enabled() {
76
0
            return None;
77
0
        }
78
79
        // Create feature extractor
80
0
        let extractor = FeatureExtractor::new();
81
82
        // Extract features
83
0
        let features = extractor.extract(self, config);
84
85
        // Get recommendation from global tuner
86
0
        let tuner = BrickTuner::new();
87
0
        Some(tuner.recommend(&features))
88
0
    }
89
90
    /// Print tuner recommendations to console.
91
0
    pub fn print_tuner_recommendations(&self, config: &RunConfig) {
92
0
        if let Some(rec) = self.get_tuner_recommendations(config) {
93
0
            let tuner = BrickTuner::new();
94
0
            tuner.print_recommendation(&rec);
95
0
        } else {
96
0
            println!("Tuner recommendations not available (profiler disabled)");
97
0
        }
98
0
    }
99
100
    /// Get tokens per second from profiler.
101
0
    pub fn tokens_per_sec(&self) -> Option<f32> {
102
0
        let total_ns = self.total_ns();
103
0
        let total_tokens = self.total_tokens();
104
0
        if total_ns == 0 || total_tokens == 0 {
105
0
            return None;
106
0
        }
107
0
        Some(total_tokens as f32 * 1e9 / total_ns as f32)
108
0
    }
109
}
110
111
// ============================================================================
112
// Tests
113
// ============================================================================
114
115
#[cfg(test)]
116
mod tests;