zluda update

modules/zluda.py

@@ -1,25 +1,23 @@
-import os
 import sys
 from typing import Union
 import torch
 from torch._prims_common import DeviceLikeType
 import onnxruntime as ort
 from modules import shared, devices
 from modules.onnx_impl.execution_providers import available_execution_providers, ExecutionProvider
+from modules.zluda_hijacks import do_hijack
 
 
 PLATFORM = sys.platform
 do_nothing = lambda _: None # pylint: disable=unnecessary-lambda-assignment
 
 
-def _join_rocm_home(*paths) -> str:
-    from torch.utils.cpp_extension import ROCM_HOME
-    return os.path.join(ROCM_HOME, *paths)
-
-
 def is_zluda(device: DeviceLikeType):
-    device = torch.device(device)
-    return torch.cuda.get_device_name(device).endswith("[ZLUDA]")
+    try:
+        device = torch.device(device)
+        return torch.cuda.get_device_name(device).endswith("[ZLUDA]")
+    except Exception:
+        return False
 
 
 def test(device: DeviceLikeType) -> Union[Exception, None]:
@@ -36,19 +34,12 @@ def test(device: DeviceLikeType) -> Union[Exception, None]:
 
 def initialize_zluda():
     device = devices.get_optimal_device()
-    if not torch.cuda.is_available() or not is_zluda(device):
+    if not devices.cuda_ok or not is_zluda(device):
         return
 
-    torch.version.hip = "5.7"
-    sys.platform = ""
-    from torch.utils import cpp_extension
-    sys.platform = PLATFORM
-    cpp_extension.IS_WINDOWS = PLATFORM == "win32"
-    cpp_extension.IS_MACOS = False
-    cpp_extension.IS_LINUX = sys.platform.startswith('linux')
-    cpp_extension._join_rocm_home = _join_rocm_home # pylint: disable=protected-access
+    do_hijack()
 
-    if cpp_extension.IS_WINDOWS:
+    if PLATFORM == "win32":
         torch.backends.cudnn.enabled = False
         torch.backends.cuda.enable_flash_sdp(False)
         torch.backends.cuda.enable_flash_sdp = do_nothing

modules/zluda_hijacks.py

@@ -0,0 +1,28 @@
+import os
+import sys
+import torch
+
+
+_topk = torch.topk
+def topk(tensor: torch.Tensor, *args, **kwargs):
+    device = tensor.device
+    values, indices = _topk(tensor.cpu(), *args, **kwargs)
+    return torch.return_types.topk((values.to(device), indices.to(device),))
+
+
+def _join_rocm_home(*paths) -> str:
+    from torch.utils.cpp_extension import ROCM_HOME
+    return os.path.join(ROCM_HOME, *paths)
+
+
+def do_hijack():
+    torch.version.hip = "5.7"
+    torch.topk = topk
+    platform = sys.platform
+    sys.platform = ""
+    from torch.utils import cpp_extension
+    sys.platform = platform
+    cpp_extension.IS_WINDOWS = platform == "win32"
+    cpp_extension.IS_MACOS = False
+    cpp_extension.IS_LINUX = platform.startswith('linux')
+    cpp_extension._join_rocm_home = _join_rocm_home # pylint: disable=protected-access