automatically download zluda, add zluda in PATH

This commit is contained in:
Seunghoon Lee
2024-02-28 18:15:58 +09:00
parent e3c1432502
commit e0754f7aca
+26 -1
View File
@@ -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