From a5a7d65dfe68a9eebe4a01f6080969797a5e0848 Mon Sep 17 00:00:00 2001
From: Seunghoon Lee <lshqqytiger@naver.com>
Date: Fri, 28 Jun 2024 16:39:54 +0900
Subject: [PATCH] wsl rocm preload hsa runtime library

---
 installer.py               | 14 ++++++++++++++
 modules/zluda_installer.py |  2 +-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/installer.py b/installer.py
index 4aeecd7e6..fefee8d55 100644
--- a/installer.py
+++ b/installer.py
@@ -680,6 +680,20 @@ def check_torch():
         install('onnxruntime-gpu', 'onnxruntime-gpu', ignore=True, quiet=True)
     elif is_rocm_available(allow_rocm):
         torch_command = install_rocm_zluda(torch_command)
+
+        # WSL ROCm
+        if os.environ.get('WSL_DISTRO_NAME', None) is not None:
+            import ctypes
+            try:
+                # Preload stdc++ library. This will ignore Anaconda stdc++ library.
+                ctypes.CDLL("/lib/x86_64-linux-gnu/libstdc++.so.6", mode=ctypes.RTLD_GLOBAL)
+            except OSError:
+                pass
+            try:
+                # Preload HSA Runtime library.
+                ctypes.CDLL("/opt/rocm/lib/libhsa-runtime64.so", mode=ctypes.RTLD_GLOBAL)
+            except OSError:
+                log.error("Failed to preload HSA Runtime library.")
     elif is_ipex_available(allow_ipex):
         torch_command = install_ipex(torch_command)
     elif allow_openvino and args.use_openvino:
diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py
index fb7d94780..17956d439 100644
--- a/modules/zluda_installer.py
+++ b/modules/zluda_installer.py
@@ -33,7 +33,7 @@ def install(zluda_path: os.PathLike) -> None:
     if os.path.exists(zluda_path):
         return
 
-    if platform.system() != 'Windows': # TODO
+    if platform.system() != 'Windows': # Windows-only. (PyTorch should be rebuilt on Linux)
         return
 
     with open('_zluda', 'wb') as file: