From cb457f898a05fcb8fe6f7e3b1880f43ddf20d027 Mon Sep 17 00:00:00 2001
From: Louis <louisoes05@gmail.com>
Date: Wed, 12 Aug 2026 13:30:58 +0800
Subject: [PATCH] fix(runtime): discover cached ROCm installations

---
 crates/koharu-runtime/src/hardware/hip.rs | 69 +++++++++++++++++++----
 1 file changed, 58 insertions(+), 11 deletions(-)

diff --git a/crates/koharu-runtime/src/hardware/hip.rs b/crates/koharu-runtime/src/hardware/hip.rs
index cf1288fd5..84ef6bb2a 100644
--- a/crates/koharu-runtime/src/hardware/hip.rs
+++ b/crates/koharu-runtime/src/hardware/hip.rs
@@ -1,7 +1,9 @@
-use std::ffi::c_void;
+use std::{ffi::c_void, path::PathBuf};
 
 use libloading::Library;
 
+use crate::Store;
+
 const BUFFER_SIZE: usize = 64 * 1024;
 type GetProperties = unsafe extern "C" fn(*mut c_void, i32) -> i32;
 
@@ -9,16 +11,7 @@ type GetProperties = unsafe extern "C" fn(*mut c_void, i32) -> i32;
 struct Properties([u8; BUFFER_SIZE]);
 
 pub(super) fn probe() -> Option<String> {
-    let names: &[&str] = if cfg!(target_os = "windows") {
-        &["amdhip64.dll", "amdhip64_7.dll"]
-    } else if cfg!(target_os = "linux") {
-        &["libamdhip64.so", "libamdhip64.so.7"]
-    } else {
-        &[]
-    };
-    let library = names
-        .iter()
-        .find_map(|name| unsafe { Library::new(name).ok() })?;
+    let library = library_candidates().find_map(|path| unsafe { Library::new(path).ok() })?;
     let get = unsafe {
         library
             .get::<GetProperties>(b"hipGetDeviceProperties\0")
@@ -31,6 +24,48 @@ pub(super) fn probe() -> Option<String> {
     target(&properties.0).map(str::to_owned)
 }
 
+fn library_candidates() -> impl Iterator<Item = PathBuf> {
+    let names: &[&str] = if cfg!(target_os = "windows") {
+        &["amdhip64.dll", "amdhip64_7.dll"]
+    } else if cfg!(target_os = "linux") {
+        &["libamdhip64.so", "libamdhip64.so.7"]
+    } else {
+        &[]
+    };
+    let system = names.iter().map(PathBuf::from);
+    let installed = installed_library_candidates(names);
+    system.chain(installed)
+}
+
+#[cfg(target_os = "windows")]
+fn installed_library_candidates(names: &[&str]) -> impl Iterator<Item = PathBuf> {
+    let root = Store::root().join("rocm");
+    let versions = std::fs::read_dir(root)
+        .into_iter()
+        .flatten()
+        .filter_map(Result::ok)
+        .filter(|entry| entry.file_type().is_ok_and(|file_type| file_type.is_dir()));
+    let targets = versions.flat_map(|version| {
+        std::fs::read_dir(version.path())
+            .into_iter()
+            .flatten()
+            .filter_map(Result::ok)
+            .filter(|entry| entry.file_type().is_ok_and(|file_type| file_type.is_dir()))
+            .collect::<Vec<_>>()
+    });
+    targets.flat_map(move |target| {
+        names
+            .iter()
+            .map(move |name| target.path().join("core").join("bin").join(name))
+            .collect::<Vec<_>>()
+    })
+}
+
+#[cfg(not(target_os = "windows"))]
+fn installed_library_candidates(_names: &[&str]) -> impl Iterator<Item = PathBuf> {
+    std::iter::empty()
+}
+
 fn target(properties: &[u8]) -> Option<&str> {
     properties
         .windows(3)
@@ -59,4 +94,16 @@ mod tests {
     fn extracts_the_gfx_architecture() {
         assert_eq!(target(b"Radeon\0gfx1201:sramecc-\0"), Some("gfx1201"));
     }
+
+    #[test]
+    fn system_candidates_are_kept_first() {
+        let candidates = library_candidates().take(2).collect::<Vec<_>>();
+        assert_eq!(
+            candidates,
+            [
+                PathBuf::from("amdhip64.dll"),
+                PathBuf::from("amdhip64_7.dll")
+            ]
+        );
+    }
 }