trueno/tuner/
mod.rs

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
36mod brick_tuner;
37mod data_collector;
38pub mod error;
39mod evolution;
40mod features;
41pub(crate) mod helpers;
42mod models;
43pub mod pretrained;
44mod types;
45
46// Re-export all public types
47pub use brick_tuner::{BrickTuner, ExperimentSuggestion, TunerRecommendation};
48pub use data_collector::{
49    ConceptDriftStatus, TrainingSample, TrainingStats, TunerDataCollector, UserFeedback,
50};
51pub use error::TunerError;
52pub use evolution::{CalibrationResult, KernelArm, KernelBandit, OnlineLearner};
53pub use features::{FeatureExtractor, RunConfig, TunerFeatures, TunerFeaturesBuilder};
54pub use models::{
55    BottleneckClassifier, BottleneckPrediction, KernelClassifier, KernelRecommendation,
56    ThroughputPrediction, ThroughputRegressor,
57};
58pub use types::{BottleneckClass, KernelType, QuantType};
59
60// Re-export helpers for tests (crate-internal)
61#[cfg(test)]
62pub(crate) use helpers::{chrono_lite_now, pad_right};
63
64// ============================================================================
65// BrickProfiler Integration
66// ============================================================================
67
68use crate::brick::BrickProfiler;
69
70impl BrickProfiler {
71    /// Get ML-based tuning recommendations.
72    ///
73    /// Extracts features from current profile and returns recommendations.
74    pub fn get_tuner_recommendations(&self, config: &RunConfig) -> Option<TunerRecommendation> {
75        if !self.is_enabled() {
76            return None;
77        }
78
79        // Create feature extractor
80        let extractor = FeatureExtractor::new();
81
82        // Extract features
83        let features = extractor.extract(self, config);
84
85        // Get recommendation from global tuner
86        let tuner = BrickTuner::new();
87        Some(tuner.recommend(&features))
88    }
89
90    /// Print tuner recommendations to console.
91    pub fn print_tuner_recommendations(&self, config: &RunConfig) {
92        if let Some(rec) = self.get_tuner_recommendations(config) {
93            let tuner = BrickTuner::new();
94            tuner.print_recommendation(&rec);
95        } else {
96            println!("Tuner recommendations not available (profiler disabled)");
97        }
98    }
99
100    /// Get tokens per second from profiler.
101    pub fn tokens_per_sec(&self) -> Option<f32> {
102        let total_ns = self.total_ns();
103        let total_tokens = self.total_tokens();
104        if total_ns == 0 || total_tokens == 0 {
105            return None;
106        }
107        Some(total_tokens as f32 * 1e9 / total_ns as f32)
108    }
109}
110
111// ============================================================================
112// Tests
113// ============================================================================
114
115#[cfg(test)]
116mod tests;
trueno/tuner/mod.rs

trueno/tuner/
mod.rs
