From 801ebdd08013c00f9b2af62fbc04f45e9dd8d137 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Thu, 24 Oct 2024 15:06:39 +0300
Subject: [PATCH] Treat Zluda as a different backend and auto disable BF16 for
 Zluda and ROCm on RDNA1-2

---
 modules/devices.py | 26 ++++++++++++++++++++++----
 modules/shared.py  |  8 ++++----
 modules/zluda.py   | 10 +---------
 3 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/modules/devices.py b/modules/devices.py
index c23f4a256..588873b1a 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -46,6 +46,16 @@ def has_xpu() -> bool:
     return bool(hasattr(torch, 'xpu') and torch.xpu.is_available())
 
 
+def has_zluda() -> bool:
+    if not cuda_ok:
+        return False
+    try:
+        device = torch.device("cuda")
+        return torch.cuda.get_device_name(device).endswith("[ZLUDA]")
+    except Exception:
+        return False
+
+
 def get_backend(shared_cmd_opts):
     global args # pylint: disable=global-statement
     args = shared_cmd_opts
@@ -55,6 +65,8 @@ def get_backend(shared_cmd_opts):
         name = 'directml'
     elif has_xpu():
         name = 'ipex'
+    elif has_zluda():
+        name = 'zluda'
     elif torch.cuda.is_available() and torch.version.cuda:
         name = 'cuda'
     elif torch.cuda.is_available() and torch.version.hip:
@@ -109,7 +121,7 @@ def get_gpu_info():
                     'device': f'{torch.xpu.get_device_name(torch.xpu.current_device())} n={torch.xpu.device_count()}',
                     'ipex': get_package_version('intel-extension-for-pytorch'),
                 }
-            elif backend == 'cuda':
+            elif backend == 'cuda' or backend == 'zluda':
                 return {
                     'device': f'{torch.cuda.get_device_name(torch.cuda.current_device())} n={torch.cuda.device_count()} arch={torch.cuda.get_arch_list()[-1]} capability={torch.cuda.get_device_capability(device)}',
                     'cuda': torch.version.cuda,
@@ -267,9 +279,15 @@ def test_bf16():
     global bf16_ok # pylint: disable=global-statement
     if bf16_ok is not None:
         return bf16_ok
-    if sys.platform == "darwin" or backend == 'openvino' or backend == 'directml': # override
-        bf16_ok = False
-        return bf16_ok
+    if opts.cuda_dtype != 'BF16': # don't override if the user sets it
+        if sys.platform == "darwin" or backend == 'openvino' or backend == 'directml' or backend == 'zluda': # override
+            bf16_ok = False
+            return bf16_ok
+        if backend == 'rocm':
+            gcn_arch = getattr(torch.cuda.get_device_properties(device), "gcnArchName", "gfx0000")[3:7]
+            if len(gcn_arch) == 4 and gcn_arch[0:2] == "10": # RDNA 1 and 2
+                bf16_ok = False
+                return bf16_ok
     try:
         import torch.nn.functional as F
         image = torch.randn(1, 4, 32, 32).to(device=device, dtype=torch.bfloat16)
diff --git a/modules/shared.py b/modules/shared.py
index b2ce7ebaa..bb1e13b9e 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -19,7 +19,6 @@ from modules import errors, devices, shared_items, shared_state, cmd_args, theme
 from modules.paths import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611
 from modules.dml import memory_providers, default_memory_provider, directml_do_hijack
 from modules.onnx_impl import initialize_onnx, execution_providers
-from modules.zluda import initialize_zluda
 from modules.memstats import memory_stats
 import modules.interrogate
 import modules.memmon
@@ -413,8 +412,8 @@ def get_default_modes():
 
     if devices.backend == "rocm":
         default_sdp_options =  ['Memory attention', 'Math attention']
-    #elif devices.backend == "zluda":
-    #    sdp_options_default =  ['Math attention']
+    elif devices.backend == "zluda":
+        default_sdp_options =  ['Math attention']
     else:
         default_sdp_options = ['Flash attention', 'Memory attention', 'Math attention']
     if (cmd_opts.lowvram or cmd_opts.medvram) and ('Flash attention' not in default_sdp_options):
@@ -1124,7 +1123,8 @@ mem_mon = modules.memmon.MemUsageMonitor("MemMon", devices.device)
 history = history.History()
 if devices.backend == "directml":
     directml_do_hijack()
-elif devices.backend == "cuda":
+elif devices.backend == "zluda":
+    from modules.zluda import initialize_zluda
     initialize_zluda()
 initialize_onnx()
 try:
diff --git a/modules/zluda.py b/modules/zluda.py
index d1b137cb6..e11e34a80 100644
--- a/modules/zluda.py
+++ b/modules/zluda.py
@@ -12,14 +12,6 @@ PLATFORM = sys.platform
 do_nothing = lambda _: None # pylint: disable=unnecessary-lambda-assignment
 
 
-def is_zluda(device: DeviceLikeType):
-    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]:
     device = torch.device(device)
     try:
@@ -35,7 +27,7 @@ def test(device: DeviceLikeType) -> Union[Exception, None]:
 def initialize_zluda():
     shared.cmd_opts.device_id = None
     device = devices.get_optimal_device()
-    if not devices.cuda_ok or not is_zluda(device):
+    if not devices.cuda_ok or not devices.has_zluda():
         return
 
     do_hijack()