From e0754f7acafefe2045369e5ee87cd58a5c51b7d8 Mon Sep 17 00:00:00 2001
From: Seunghoon Lee <lshqqytiger@naver.com>
Date: Wed, 28 Feb 2024 18:15:58 +0900
Subject: [PATCH] automatically download zluda, add zluda in PATH

---
 installer.py | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/installer.py b/installer.py
index bccd86d30..a132b7e36 100644
--- a/installer.py
+++ b/installer.py
@@ -481,7 +481,27 @@ def check_torch():
         if args.use_zluda:
             torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.2.0 torchvision --index-url https://download.pytorch.org/whl/cu118')
             log.warning("ZLUDA support: experimental")
+            paths = os.environ.get('PATH', '.')
             zluda_need_dll_patch = is_windows and not installed('torch')
+            zluda_path = find_zluda()
+            if zluda_path is None:
+                import urllib.request
+                if is_windows:
+                    import zipfile
+                    archive_type = zipfile.ZipFile
+                    zluda_url = 'https://github.com/lshqqytiger/ZLUDA/releases/download/v3.5-win/ZLUDA-windows-amd64.zip'
+                else:
+                    import tarfile
+                    archive_type = tarfile.TarFile
+                    zluda_url = 'https://github.com/vosen/ZLUDA/releases/download/v3/zluda-3-linux.tar.gz'
+                urllib.request.urlretrieve(zluda_url, '_zluda')
+                with archive_type('_zluda', 'r') as f:
+                    f.extractall('.zluda')
+                zluda_path = os.path.abspath('./.zluda')
+                os.remove('_zluda')
+            log.debug(f'Found ZLUDA in {zluda_path}')
+            if zluda_path not in paths:
+                os.environ['PATH'] += ';' + zluda_path
         elif is_windows: # TODO TBD after ROCm for Windows is released
             log.warning("HIP SDK is detected, but no Torch release for Windows available")
             log.info("For ZLUDA support specify '--use-zluda'")
@@ -909,7 +929,7 @@ def get_onnxruntime_source_for_rocm(rocm_ver):
         return 'onnxruntime-gpu'
 
 
-def patch_zluda():
+def find_zluda():
     zluda_path = os.environ.get('ZLUDA', None)
     if zluda_path is None:
         paths = os.environ.get('PATH', '').split(';')
@@ -917,6 +937,11 @@ def patch_zluda():
             if os.path.exists(os.path.join(path, 'zluda_redirect.dll')):
                 zluda_path = path
                 break
+    return zluda_path
+
+
+def patch_zluda():
+    zluda_path = find_zluda()
     if zluda_path is None:
         log.warning('Failed to automatically patch torch with ZLUDA. Could not find ZLUDA from PATH.')
         return