sokuto_runtime/engine/
config.rs

1//! # 推論セッション設定モジュール
2//!
3//! ONNX Runtime セッションの Execution Provider 優先度、
4//! スレッド並列数、グラフ最適化レベル、メモリアリーナなどの構成パラメータを定義する。
5
6use std::path::PathBuf;
7
8use ort::session::builder::GraphOptimizationLevel;
9
10/// システムの物理コア数とセッションプール数から最適な `intra_threads` を算出する。
11///
12/// ハイパースレッディング (SMT) による論理コアではなく、物理コア数 (`num_cpus::get_physical`) を基準とし、
13/// スレッド間の演算器競合およびキャッシュスラッシングを防止する。
14///
15/// # 計算式
16/// `max(1, min(max_cap, physical_cores / pool_size))`
17///
18/// # 引数
19/// - `pool_size`: セッションプール数。
20/// - `max_cap`: 単一セッションに割り当てるスレッド数の上限値。
21pub fn auto_intra_threads(pool_size: usize, max_cap: usize) -> usize {
22    let physical = num_cpus::get_physical().max(1);
23    let pool = pool_size.max(1);
24    let per_session = physical / pool;
25    per_session.clamp(1, max_cap)
26}
27
28/// 利用可能な Execution Provider (実行バックエンド) の指定。
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
30pub enum ExecutionProvider {
31    /// 自動選択 (利用可能なアクセラレータを優先し、利用不可時は CPU へフォールバック)。
32    #[default]
33    Auto,
34    /// Apple Silicon 向け CoreML バックエンド。
35    CoreML,
36    /// NVIDIA GPU 向け CUDA バックエンド。
37    CUDA,
38    /// NVIDIA GPU 向け TensorRT バックエンド。
39    TensorRT,
40    /// CPU バックエンド (高移植性・標準フォールバック)。
41    CPU,
42}
43
44impl ExecutionProvider {
45    /// バックエンドの表示名を取得する。
46    pub fn name(&self) -> &'static str {
47        match self {
48            Self::Auto => "Auto",
49            Self::CoreML => "CoreML",
50            Self::CUDA => "CUDA",
51            Self::TensorRT => "TensorRT",
52            Self::CPU => "CPU",
53        }
54    }
55}
56
57/// グラフ最適化レベル。
58#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
59pub enum OptimizationLevel {
60    /// 最適化を無効化。
61    Disable,
62    /// 基本的な最適化 (ノード結合、定数畳み込み等)。
63    Basic,
64    /// 拡張最適化 (レイヤー正規化融合、GELU融合等)。
65    Extended,
66    /// すべての最適化を適用 (デフォルト)。
67    #[default]
68    All,
69}
70
71impl OptimizationLevel {
72    /// `ort` クレートの `GraphOptimizationLevel` に変換する。
73    pub fn to_ort_level(self) -> GraphOptimizationLevel {
74        match self {
75            Self::Disable => GraphOptimizationLevel::Disable,
76            Self::Basic => GraphOptimizationLevel::Level1,
77            Self::Extended => GraphOptimizationLevel::Level2,
78            Self::All => GraphOptimizationLevel::Level3,
79        }
80    }
81}
82
83/// ONNX Runtime セッション構築設定。
84#[derive(Debug, Clone)]
85pub struct SessionConfig {
86    /// 優先的に試行する Execution Provider の順序付きリスト。
87    pub preferred_providers: Vec<ExecutionProvider>,
88    /// 単一オペレータ内の並列スレッド数 (intra_op)。None の場合は物理コア数から自動算出。
89    pub intra_threads: Option<usize>,
90    /// 複数オペレータ間の並列スレッド数 (inter_op)。None の場合は自動設定 (通常は 1)。
91    pub inter_threads: Option<usize>,
92    /// グラフ内の並列実行モード (false で Sequential、true で Parallel)。
93    pub parallel_execution: bool,
94    /// グラフ最適化レベル。
95    pub optimization_level: OptimizationLevel,
96    /// CPU メモリアリーナの有効化。
97    pub enable_mem_arena: bool,
98    /// 並行推論用セッションプールサイズ (デフォルト: 1)。
99    pub pool_size: usize,
100    /// 最適化済みモデルグラフのディスク保存・キャッシュパス (任意)。
101    pub optimized_model_path: Option<PathBuf>,
102}
103
104impl Default for SessionConfig {
105    fn default() -> Self {
106        #[cfg(target_os = "macos")]
107        let default_providers = vec![ExecutionProvider::CoreML, ExecutionProvider::CPU];
108
109        #[cfg(not(target_os = "macos"))]
110        let default_providers = vec![
111            ExecutionProvider::TensorRT,
112            ExecutionProvider::CUDA,
113            ExecutionProvider::CPU,
114        ];
115
116        Self {
117            preferred_providers: default_providers,
118            intra_threads: None,
119            inter_threads: Some(1),
120            parallel_execution: false,
121            optimization_level: OptimizationLevel::All,
122            enable_mem_arena: true,
123            pool_size: 1,
124            optimized_model_path: None,
125        }
126    }
127}
128
129impl SessionConfig {
130    /// ホスト環境の物理コア数に基づき、スレッド数を自動計算した設定を生成する。
131    pub fn auto() -> Self {
132        let mut config = Self::default();
133        config.intra_threads = Some(auto_intra_threads(config.pool_size, 8));
134        config
135    }
136
137    /// CPU 推論に特化した設定を生成する。
138    pub fn cpu_only() -> Self {
139        Self {
140            preferred_providers: vec![ExecutionProvider::CPU],
141            intra_threads: None,
142            inter_threads: Some(1),
143            parallel_execution: false,
144            optimization_level: OptimizationLevel::All,
145            enable_mem_arena: true,
146            pool_size: 1,
147            optimized_model_path: None,
148        }
149    }
150
151    /// Tier 1 (130M-INT8) モデル向けの推奨セッション設定を生成する。
152    ///
153    /// 軽量モデルの短時間フォワードパス (~12ms) においてスレッド同期オーバーヘッドを抑えるため、
154    /// `intra_threads` の上限を 4 スレッドに制限する。
155    pub fn for_tier1() -> Self {
156        let mut config = Self::cpu_only();
157        config.intra_threads = Some(auto_intra_threads(config.pool_size, 4));
158        config
159    }
160
161    /// Tier 2 (310M-INT8) モデル向けの推奨セッション設定を生成する。
162    ///
163    /// 計算負荷の高い GEMM 演算を高速化するため、物理コア数に応じて最大 8 スレッドまで割り当てる。
164    pub fn for_tier2() -> Self {
165        let mut config = Self::cpu_only();
166        config.intra_threads = Some(auto_intra_threads(config.pool_size, 8));
167        config
168    }
169
170    /// 現在のセッションプール数に基づき、スレッド数を自動計算して適用する。
171    pub fn with_auto_threads(mut self) -> Self {
172        self.intra_threads = Some(auto_intra_threads(self.pool_size, 8));
173        self
174    }
175
176    /// スレッド数を明示指定したビルダーメソッド。
177    pub fn with_threads(mut self, intra: Option<usize>, inter: Option<usize>) -> Self {
178        self.intra_threads = intra;
179        self.inter_threads = inter;
180        self
181    }
182
183    /// 並列実行フラグを指定したビルダーメソッド。
184    pub fn with_parallel_execution(mut self, parallel: bool) -> Self {
185        self.parallel_execution = parallel;
186        self
187    }
188
189    /// 最適化レベルを指定したビルダーメソッド。
190    pub fn with_optimization_level(mut self, level: OptimizationLevel) -> Self {
191        self.optimization_level = level;
192        self
193    }
194
195    /// メモリアリーナ有効化フラグを指定したビルダーメソッド。
196    pub fn with_memory_arena(mut self, enable: bool) -> Self {
197        self.enable_mem_arena = enable;
198        self
199    }
200
201    /// セッションプールサイズを指定したビルダーメソッド。
202    pub fn with_pool_size(mut self, size: usize) -> Self {
203        self.pool_size = size.max(1);
204        self
205    }
206
207    /// 最適化済みモデルの保存・キャッシュパスを指定したビルダーメソッド。
208    pub fn with_optimized_model_path(mut self, path: impl Into<PathBuf>) -> Self {
209        self.optimized_model_path = Some(path.into());
210        self
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    #[test]
219    fn test_auto_intra_threads() {
220        let physical = num_cpus::get_physical();
221        assert!(physical >= 1);
222
223        let t1 = auto_intra_threads(1, 4);
224        assert!((1..=4).contains(&t1));
225
226        let t2 = auto_intra_threads(1, 8);
227        assert!((1..=8).contains(&t2));
228
229        let t_pool = auto_intra_threads(4, 8);
230        assert!(t_pool >= 1);
231    }
232
233    #[test]
234    fn test_tier_presets() {
235        let cfg1 = SessionConfig::for_tier1();
236        assert_eq!(cfg1.preferred_providers, vec![ExecutionProvider::CPU]);
237        assert_eq!(cfg1.inter_threads, Some(1));
238        assert!(!cfg1.parallel_execution);
239        assert!(cfg1.intra_threads.is_some());
240        assert!(cfg1.intra_threads.unwrap() <= 4);
241
242        let cfg2 = SessionConfig::for_tier2();
243        assert_eq!(cfg2.preferred_providers, vec![ExecutionProvider::CPU]);
244        assert_eq!(cfg2.inter_threads, Some(1));
245        assert!(!cfg2.parallel_execution);
246        assert!(cfg2.intra_threads.is_some());
247        assert!(cfg2.intra_threads.unwrap() <= 8);
248    }
249}
sokuto_runtime/engine/config.rs

sokuto_runtime/engine/
config.rs
