sokuto_runtime/engine/
provider.rs

1//! # Execution Provider 解決モジュール
2//!
3//! 利用可能なアクセラレータ(CoreML / CUDA / TensorRT)の登録を試行し、
4//! 初期化失敗時やドライバ未検出時には CPU バックエンドへ安全にフォールバックする。
5
6use ort::ep::ExecutionProvider as OrtExecutionProvider;
7use ort::session::builder::SessionBuilder;
8
9use crate::engine::config::ExecutionProvider;
10use crate::error::{Result, RuntimeError};
11
12/// 指定された優先順位リストに従って Execution Provider をセッションビルダーに登録する。
13///
14/// 実際に登録に成功した最優先の Execution Provider を返却する。
15/// すべてのアクセラレータが失敗または未コンパイルの場合は、必ず CPU へフォールバックする。
16pub fn register_execution_providers(
17    builder: &mut SessionBuilder,
18    preferred: &[ExecutionProvider],
19    enable_mem_arena: bool,
20) -> Result<ExecutionProvider> {
21    for provider in preferred {
22        match provider {
23            ExecutionProvider::Auto => {
24                // Auto の場合はプラットフォーム標準の優先順位を再帰的に試行
25                #[cfg(target_os = "macos")]
26                let sub_preferred = [ExecutionProvider::CoreML, ExecutionProvider::CPU];
27                #[cfg(not(target_os = "macos"))]
28                let sub_preferred = [
29                    ExecutionProvider::TensorRT,
30                    ExecutionProvider::CUDA,
31                    ExecutionProvider::CPU,
32                ];
33                if let Ok(active) =
34                    register_execution_providers(builder, &sub_preferred, enable_mem_arena)
35                {
36                    return Ok(active);
37                }
38            }
39            ExecutionProvider::CoreML => {
40                #[cfg(feature = "coreml")]
41                {
42                    use ort::ep::coreml::CoreML;
43                    let ep = CoreML::default();
44                    match ep.is_available() {
45                        Ok(true) => match ep.register(builder) {
46                            Ok(()) => {
47                                tracing::info!(
48                                    "CoreML Execution Provider が正常に登録されました。"
49                                );
50                                return Ok(ExecutionProvider::CoreML);
51                            }
52                            Err(e) => {
53                                tracing::warn!(
54                                    "CoreML の登録に失敗しました ({e})。CPU にフォールバックします。"
55                                );
56                            }
57                        },
58                        Ok(false) => {
59                            tracing::info!(
60                                "CoreML は利用可能ではありません。次の EP を試行します。"
61                            );
62                        }
63                        Err(e) => {
64                            tracing::warn!("CoreML 利用可否確認中にエラーが発生しました: {e}。");
65                        }
66                    }
67                }
68                #[cfg(not(feature = "coreml"))]
69                {
70                    tracing::debug!(
71                        "CoreML feature は未コンパイルです。CPU にフォールバックします。"
72                    );
73                }
74            }
75            ExecutionProvider::CUDA => {
76                #[cfg(feature = "cuda")]
77                {
78                    use ort::ep::cuda::CUDA;
79                    let ep = CUDA::default();
80                    match ep.is_available() {
81                        Ok(true) => match ep.register(builder) {
82                            Ok(()) => {
83                                tracing::info!("CUDA Execution Provider が正常に登録されました。");
84                                return Ok(ExecutionProvider::CUDA);
85                            }
86                            Err(e) => {
87                                tracing::warn!(
88                                    "CUDA の登録に失敗しました ({e})。CPU にフォールバックします。"
89                                );
90                            }
91                        },
92                        Ok(false) => {
93                            tracing::info!("CUDA は利用可能ではありません。次の EP を試行します。");
94                        }
95                        Err(e) => {
96                            tracing::warn!("CUDA 利用可否確認中にエラーが発生しました: {e}。");
97                        }
98                    }
99                }
100                #[cfg(not(feature = "cuda"))]
101                {
102                    tracing::debug!(
103                        "CUDA feature は未コンパイルです。CPU にフォールバックします。"
104                    );
105                }
106            }
107            ExecutionProvider::TensorRT => {
108                #[cfg(feature = "tensorrt")]
109                {
110                    use ort::ep::tensorrt::TensorRT;
111                    let ep = TensorRT::default();
112                    match ep.is_available() {
113                        Ok(true) => match ep.register(builder) {
114                            Ok(()) => {
115                                tracing::info!(
116                                    "TensorRT Execution Provider が正常に登録されました。"
117                                );
118                                return Ok(ExecutionProvider::TensorRT);
119                            }
120                            Err(e) => {
121                                tracing::warn!(
122                                    "TensorRT の登録に失敗しました ({e})。CPU にフォールバックします。"
123                                );
124                            }
125                        },
126                        Ok(false) => {
127                            tracing::info!(
128                                "TensorRT は利用可能ではありません。次の EP を試行します。"
129                            );
130                        }
131                        Err(e) => {
132                            tracing::warn!("TensorRT 利用可否確認中にエラーが発生しました: {e}。");
133                        }
134                    }
135                }
136                #[cfg(not(feature = "tensorrt"))]
137                {
138                    tracing::debug!(
139                        "TensorRT feature は未コンパイルです。CPU にフォールバックします。"
140                    );
141                }
142            }
143            ExecutionProvider::CPU => {
144                let ep = ort::ep::CPU::default().with_arena_allocator(enable_mem_arena);
145                ep.register(builder).map_err(|e| {
146                    RuntimeError::ExecutionProviderError(format!(
147                        "CPUExecutionProvider の登録に失敗しました: {e}。"
148                    ))
149                })?;
150                tracing::info!(
151                    "CPU Execution Provider が正常に登録されました (メモリアリーナ: {enable_mem_arena})。"
152                );
153                return Ok(ExecutionProvider::CPU);
154            }
155        }
156    }
157
158    // いずれも成功しなかった場合の最終フォールバック
159    let ep = ort::ep::CPU::default().with_arena_allocator(enable_mem_arena);
160    ep.register(builder).map_err(|e| {
161        RuntimeError::ExecutionProviderError(format!(
162            "CPUExecutionProvider への最終フォールバックに失敗しました: {e}。"
163        ))
164    })?;
165    tracing::info!(
166        "CPU Execution Provider への最終フォールバックが完了しました (メモリアリーナ: {enable_mem_arena})。"
167    );
168    Ok(ExecutionProvider::CPU)
169}
sokuto_runtime/engine/provider.rs

sokuto_runtime/engine/
provider.rs
