use ort::session::builder::{BuilderResult, SessionBuilder};
#[allow(dead_code)]
pub fn gpu_execution_provider() -> Option<&'static str> {
#[cfg(feature = "gpu-cuda")]
{
return Some("CUDAExecutionProvider");
}
#[cfg(feature = "gpu-coreml")]
{
return Some("CoreMLExecutionProvider");
}
#[cfg(feature = "gpu-directml")]
{
return Some("DmlExecutionProvider");
}
#[cfg(feature = "gpu-rocm")]
{
return Some("ROCMExecutionProvider");
}
#[allow(unreachable_code)]
None
}
#[allow(clippy::result_large_err)]
pub fn apply_gpu_ep(builder: SessionBuilder) -> BuilderResult {
#[cfg(feature = "gpu-cuda")]
{
tracing::info!("registering CUDAExecutionProvider");
return builder.with_execution_providers([ort::ep::CUDA::default().build()]);
}
#[cfg(feature = "gpu-coreml")]
{
tracing::info!("registering CoreMLExecutionProvider");
return builder.with_execution_providers([ort::ep::CoreML::default().build()]);
}
#[cfg(feature = "gpu-directml")]
{
tracing::info!("registering DmlExecutionProvider");
return builder.with_execution_providers([ort::ep::DirectML::default().build()]);
}
#[cfg(feature = "gpu-rocm")]
{
tracing::info!("registering ROCMExecutionProvider");
return builder.with_execution_providers([ort::ep::ROCm::default().build()]);
}
#[allow(unreachable_code)]
Ok(builder)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_build_is_cpu_only() {
assert!(
gpu_execution_provider().is_none(),
"no GPU execution provider should be selected by default"
);
}
#[test]
fn apply_gpu_ep_is_no_op_in_cpu_build() {
let builder = ort::session::Session::builder().expect("failed to create session builder");
apply_gpu_ep(builder).expect("CPU build apply_gpu_ep must return Ok");
}
}