feat(onnx): add DirectML/MIGraphX/CoreML execution provider support

colgrep/src/index/mod.rs

@@ -999,11 +999,10 @@ impl IndexBuilder {
                     }
                 }
             };
-            #[cfg(not(feature = "cuda"))]
+            #[cfg(not(any(feature = "cuda", feature = "directml", feature = "migraphx", feature = "coreml")))]
             let (num_sessions, execution_provider) = {
-                let _ = num_units; // Silence unused warning when cuda feature is disabled
+                let _ = num_units;
 
-                // Initialize ONNX Runtime (CPU-only build)
                 crate::onnx_runtime::ensure_onnx_runtime()
                     .context("Failed to initialize ONNX Runtime")?;
 
@@ -1014,8 +1013,31 @@ impl IndexBuilder {
                 )
             };
 
+            #[cfg(any(feature = "directml", feature = "migraphx", feature = "coreml"))]
+            #[cfg(not(feature = "cuda"))]
+            let (num_sessions, execution_provider) = {
+                let _ = num_units;
+
+                crate::onnx_runtime::ensure_onnx_runtime()
+                    .context("Failed to initialize ONNX Runtime")?;
+
+                let provider = if cfg!(feature = "directml") {
+                    ExecutionProvider::DirectML
+                } else if cfg!(feature = "migraphx") {
+                    ExecutionProvider::MIGraphX
+                } else {
+                    ExecutionProvider::CoreML
+                };
+
+                (
+                    self.parallel_sessions
+                        .unwrap_or_else(crate::config::get_default_cpu_parallel_sessions),
+                    provider,
+                )
+            };
+
             // Print model info after ONNX runtime is initialized (and any potential re-exec)
-            eprintln!("🤖 Model: {}", self.model_id);
+            eprintln!("🤖 Model: {} ({})", self.model_id, execution_provider);
             eprintln!("📂 Building index...");
 
             // Use runtime default for batch size (respects cuDNN availability)
@@ -3286,6 +3308,10 @@ impl Searcher {
             AccelerationMode::Auto => {
                 if cfg!(feature = "coreml") {
                     ExecutionProvider::CoreML
+                } else if cfg!(feature = "directml") {
+                    ExecutionProvider::DirectML
+                } else if cfg!(feature = "migraphx") {
+                    ExecutionProvider::MIGraphX
                 } else {
                     ExecutionProvider::Cpu
                 }
@@ -3362,6 +3388,10 @@ impl Searcher {
             AccelerationMode::Auto => {
                 if cfg!(feature = "coreml") {
                     ExecutionProvider::CoreML
+                } else if cfg!(feature = "directml") {
+                    ExecutionProvider::DirectML
+                } else if cfg!(feature = "migraphx") {
+                    ExecutionProvider::MIGraphX
                 } else {
                     ExecutionProvider::Cpu
                 }

colgrep/src/onnx_runtime.rs

@@ -41,16 +41,11 @@ const ORT_LIB_NAME: &str = "libonnxruntime.so";
 #[cfg(target_os = "windows")]
 const ORT_LIB_NAME: &str = "onnxruntime.dll";
 
-/// Whether to use GPU (CUDA) version of ONNX Runtime
-#[cfg(feature = "cuda")]
-const USE_GPU: bool = true;
-#[cfg(not(feature = "cuda"))]
-const USE_GPU: bool = false;
 
 /// Subdirectory name for caching (gpu vs cpu)
-#[cfg(feature = "cuda")]
+#[cfg(any(feature = "cuda", feature = "directml"))]
 const ORT_CACHE_SUBDIR: &str = "gpu";
-#[cfg(not(feature = "cuda"))]
+#[cfg(not(any(feature = "cuda", feature = "directml")))]
 const ORT_CACHE_SUBDIR: &str = "cpu";
 
 /// Ensure ONNX Runtime is available.
@@ -557,7 +552,10 @@ fn download_onnx_runtime() -> Result<PathBuf> {
         && cache_dir.join("onnxruntime_providers_shared.dll").exists()
         && cache_dir.join("onnxruntime_providers_cuda.dll").exists();
 
-    #[cfg(not(feature = "cuda"))]
+    #[cfg(all(feature = "directml", not(feature = "cuda")))]
+    let already_cached = lib_path.exists();
+
+    #[cfg(not(any(feature = "cuda", feature = "directml")))]
     let already_cached = lib_path.exists();
 
     if already_cached {
@@ -568,11 +566,12 @@ fn download_onnx_runtime() -> Result<PathBuf> {
 
     let (url, files_to_extract) = get_download_info()?;
 
-    if USE_GPU {
-        eprintln!("⚙️  Runtime: ONNX {} (GPU/CUDA)", ORT_VERSION);
-    } else {
-        eprintln!("⚙️  Runtime: ONNX {} (CPU)", ORT_VERSION);
-    }
+    #[cfg(feature = "cuda")]
+    eprintln!("⚙️  Runtime: ONNX {} (GPU/CUDA)", ORT_VERSION);
+    #[cfg(all(feature = "directml", not(feature = "cuda")))]
+    eprintln!("⚙️  Runtime: ONNX {} (GPU/DirectML)", ORT_VERSION);
+    #[cfg(not(any(feature = "cuda", feature = "directml")))]
+    eprintln!("⚙️  Runtime: ONNX {} (CPU)", ORT_VERSION);
 
     // Download archive
     let response = ureq::get(&url)
@@ -593,6 +592,24 @@ type FileToExtract = (String, String);
 
 /// Get download URL and files to extract for current platform
 fn get_download_info() -> Result<(String, Vec<FileToExtract>)> {
+    // DirectML: download from NuGet (Microsoft GPU package does not include DirectML)
+    #[cfg(all(target_os = "windows", target_arch = "x86_64", feature = "directml", not(feature = "cuda")))]
+    return Ok((
+        format!(
+            "https://www.nuget.org/api/v2/package/Microsoft.ML.OnnxRuntime.DirectML/{}",
+            ORT_VERSION
+        ),
+        vec![
+            (
+                "runtimes/win-x64/native/onnxruntime.dll".to_string(),
+                "onnxruntime.dll".to_string(),
+            ),
+        ],
+    ));
+
+    // All other configurations: download from GitHub releases
+    #[cfg(not(all(target_os = "windows", target_arch = "x86_64", feature = "directml", not(feature = "cuda"))))]
+    {
     let base = format!(
         "https://github.com/microsoft/onnxruntime/releases/download/v{}",
         ORT_VERSION
@@ -694,7 +711,7 @@ fn get_download_info() -> Result<(String, Vec<FileToExtract>)> {
         )
     };
 
-    #[cfg(all(target_os = "windows", target_arch = "x86_64", not(feature = "cuda")))]
+    #[cfg(all(target_os = "windows", target_arch = "x86_64", not(any(feature = "cuda", feature = "directml"))))]
     let (archive, files) = (
         format!("onnxruntime-win-x64-{}.zip", ORT_VERSION),
         vec![(
@@ -715,6 +732,7 @@ fn get_download_info() -> Result<(String, Vec<FileToExtract>)> {
     ));
 
     Ok((format!("{}/{}", base, archive), files))
+    }
 }
 
 /// Extract libraries from tgz archive

next-plaid-onnx/src/lib.rs

@@ -198,6 +198,23 @@ pub enum ExecutionProvider {
     MIGraphX,
 }
 
+impl std::fmt::Display for ExecutionProvider {
+    /// Short user-facing label matching the tokens used in onnx_runtime.rs
+    /// download messages (e.g. "CPU", "CUDA", "DirectML"). Stable across
+    /// releases; do not change without updating callers that log the label.
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Auto => f.write_str("auto"),
+            Self::Cpu => f.write_str("CPU"),
+            Self::Cuda => f.write_str("CUDA"),
+            Self::TensorRT => f.write_str("TensorRT"),
+            Self::CoreML => f.write_str("CoreML"),
+            Self::DirectML => f.write_str("DirectML"),
+            Self::MIGraphX => f.write_str("MIGraphX"),
+        }
+    }
+}
+
 fn configure_execution_provider(
     builder: SessionBuilder,
     provider: ExecutionProvider,
@@ -2403,6 +2420,19 @@ mod tests {
         assert_eq!(debug_str, "Cuda");
     }
 
+    #[test]
+    fn test_execution_provider_display() {
+        // Labels are part of the user-facing surface (e.g. colgrep's
+        // `Model: <id> (<backend>)` line); lock them in.
+        assert_eq!(format!("{}", ExecutionProvider::Auto), "auto");
+        assert_eq!(format!("{}", ExecutionProvider::Cpu), "CPU");
+        assert_eq!(format!("{}", ExecutionProvider::Cuda), "CUDA");
+        assert_eq!(format!("{}", ExecutionProvider::TensorRT), "TensorRT");
+        assert_eq!(format!("{}", ExecutionProvider::CoreML), "CoreML");
+        assert_eq!(format!("{}", ExecutionProvider::DirectML), "DirectML");
+        assert_eq!(format!("{}", ExecutionProvider::MIGraphX), "MIGraphX");
+    }
+
     // =========================================================================
     // Pool embeddings tests
     // =========================================================================