From f46ee37f3ac2409a471ee339af2352f989ec5147 Mon Sep 17 00:00:00 2001
From: Seunghoon Lee <lshqqytiger@naver.com>
Date: Mon, 31 Mar 2025 22:49:38 +0900
Subject: [PATCH] zluda log & install improvements

---
 installer.py               |  1 -
 modules/zluda_hijacks.py   |  2 --
 modules/zluda_installer.py | 13 +++++++++----
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/installer.py b/installer.py
index 4f384f2b8..289b2d5b3 100644
--- a/installer.py
+++ b/installer.py
@@ -652,7 +652,6 @@ def install_rocm_zluda():
                 zluda_installer.make_copy()
                 zluda_installer.load()
                 torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0 torchvision --index-url https://download.pytorch.org/whl/cu118')
-                log.info(f'Using ZLUDA in {zluda_installer.path}')
             except Exception as e:
                 error = e
                 log.warning(f'Failed to load ZLUDA: {e}')
diff --git a/modules/zluda_hijacks.py b/modules/zluda_hijacks.py
index 4f1224923..ea60c846a 100644
--- a/modules/zluda_hijacks.py
+++ b/modules/zluda_hijacks.py
@@ -42,8 +42,6 @@ def torch__C__cuda_getCurrentRawStream(device):
 
 
 def do_hijack():
-    torch.version.hip = rocm.version
-
     if zluda.default_agent is not None:
         DeviceProperties.PROPERTIES_OVERRIDE["gcnArchName"] = zluda.default_agent.name
     torch.cuda._get_device_properties = torch_cuda__get_device_properties # pylint: disable=protected-access
diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py
index 236ca85f5..54a0f9234 100644
--- a/modules/zluda_installer.py
+++ b/modules/zluda_installer.py
@@ -6,6 +6,7 @@ import shutil
 import zipfile
 import urllib.request
 from typing import Union
+from installer import args, log
 from modules import rocm
 
 
@@ -25,8 +26,6 @@ path = os.path.abspath(os.environ.get('ZLUDA', '.zluda'))
 default_agent: Union[rocm.Agent, None] = None
 hipBLASLt_enabled = False
 
-nightly = os.environ.get("ZLUDA_NIGHTLY", "0") == "1"
-
 
 class ZLUDAResult(ctypes.Structure):
     _fields_ = [
@@ -101,7 +100,10 @@ def install() -> None:
 
     platform = "windows"
     commit = os.environ.get("ZLUDA_HASH", "dba64c0966df2c71e82255e942c96e2e1cea3a2d")
-    if nightly:
+    if os.environ.get("ZLUDA_NIGHTLY", "0") == "1":
+        log.warning("Environment variable 'ZLUDA_NIGHTLY' will be removed. Please use command-line argument '--use-nightly' instead.")
+        args.use_nightly = True
+    if args.use_nightly:
         platform = "nightly-" + platform
     urllib.request.urlretrieve(f'https://github.com/lshqqytiger/ZLUDA/releases/download/rel.{commit}/ZLUDA-{platform}-rocm{rocm.version[0]}-amd64.zip', '_zluda')
     with zipfile.ZipFile('_zluda', 'r') as archive:
@@ -150,6 +152,8 @@ def make_copy() -> None:
 
 
 def load() -> None:
+    log.info(f"ZLUDA load: path='{path}' nightly={bool(core.get_nightly_flag())}")
+
     os.environ["ZLUDA_COMGR_LOG_LEVEL"] = "1"
     os.environ["ZLUDA_NVRTC_LIB"] = os.path.join([v for v in site.getsitepackages() if v.endswith("site-packages")][0], "torch", "lib", "nvrtc64_112_0.dll")
 
@@ -171,7 +175,8 @@ def load() -> None:
         ctypes.windll.LoadLibrary(os.path.join(path, 'cudnn64_9.dll'))
 
     def conceal():
-        import torch # pylint: disable=unused-import
+        import torch
+        torch.version.hip = rocm.version
         platform = sys.platform
         sys.platform = ""
         from torch.utils import cpp_extension