From 74687dbc442b49f928e54fc6aad1ba8ceb7648c1 Mon Sep 17 00:00:00 2001
From: Seunghoon Lee <lshqqytiger@naver.com>
Date: Mon, 13 May 2024 17:44:44 +0900
Subject: [PATCH] improve zluda operation test

---
 cli/zluda-python.py | 22 ++++++++++++++++++++++
 modules/zluda.py    | 15 +++++++++------
 2 files changed, 31 insertions(+), 6 deletions(-)
 create mode 100644 cli/zluda-python.py

diff --git a/cli/zluda-python.py b/cli/zluda-python.py
new file mode 100644
index 000000000..e5a63b716
--- /dev/null
+++ b/cli/zluda-python.py
@@ -0,0 +1,22 @@
+import os
+import sys
+
+
+if __name__ == '__main__':
+    sys.path.append(os.getcwd())
+
+    from modules.zluda_installer import find, load
+    load(find())
+
+    import torch
+    print(f'Python with ZLUDA {sys.version}')
+    print('Type "help", "copyright", "credits" or "license" for more information.')
+
+    while True:
+        print('>>> ', end='')
+        try:
+            exec(input(), {
+                'torch': torch,
+            })
+        except Exception as e:
+            print(f'{e.__class__.__name__}: {e}')
diff --git a/modules/zluda.py b/modules/zluda.py
index 945f434d8..c679feeaa 100644
--- a/modules/zluda.py
+++ b/modules/zluda.py
@@ -1,3 +1,4 @@
+from typing import Union
 import platform
 import torch
 from torch._prims_common import DeviceLikeType
@@ -20,15 +21,15 @@ def is_zluda(device: DeviceLikeType):
     return torch.cuda.get_device_name(device).endswith("[ZLUDA]")
 
 
-def test(device: DeviceLikeType):
+def test(device: DeviceLikeType) -> Union[Exception, None]:
     device = torch.device(device)
     try:
         ten1 = torch.randn((2, 4,), device=device)
         ten2 = torch.randn((4, 8,), device=device)
         out = torch.mm(ten1, ten2)
-        return out.sum().is_nonzero()
-    except Exception:
-        return False
+        assert out.sum().is_nonzero()
+    except Exception as e:
+        return e
 
 
 def initialize_zluda():
@@ -49,8 +50,10 @@ def initialize_zluda():
             torch.nn.functional.conv2d = conv2d_cudnn_disabled
         devices.device_codeformer = devices.cpu
 
-        if not test(device):
-            shared.log.error(f'ZLUDA device failed to pass basic operation test: index={device.index}, device_name={torch.cuda.get_device_name(device)}')
+        result = test(device)
+        if result is not None:
+            shared.log.warning(f'ZLUDA device failed to pass basic operation test: index={device.index}, device_name={torch.cuda.get_device_name(device)}')
+            shared.log.error(result)
             torch.cuda.is_available = lambda: False
             devices.cuda_ok = False
             devices.backend = 'cpu'