Small fix to vulkan utils.

nod-ai · Jan 9, 2024 · e80bc9f · e80bc9f
1 parent 773e6eb
commit e80bc9f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 9 deletions.
diff --git a/shark/iree_utils/compile_utils.py b/shark/iree_utils/compile_utils.py
@@ -39,11 +39,7 @@ def get_iree_device_args(device, extra_args=[]):
         u_kernel_flag = ["--iree-llvmcpu-enable-ukernels"]
         stack_size_flag = ["--iree-llvmcpu-stack-allocation-limit=256000"]
 
-        return (
-            get_iree_cpu_args()
-            + u_kernel_flag
-            + stack_size_flag
-        )
+        return get_iree_cpu_args() + u_kernel_flag + stack_size_flag
     if device == "cuda":
         from shark.iree_utils.gpu_utils import get_iree_gpu_args
 

diff --git a/shark/iree_utils/vulkan_utils.py b/shark/iree_utils/vulkan_utils.py
@@ -143,12 +143,11 @@ def get_vulkan_target_triple(device_name):
         triple = f"rdna3-7600-{system_os}"
     elif "7700" in device_name:
         triple = f"rdna3-7700-{system_os}"
-    elif any(x in device_name for x in "AMD", "Radeon") and any(x in device_name for x in "6700", "6750"):
+    elif any(x in device_name for x in ("Radeon 6", "RX 6", "PRO W6")):
         triple = f"rdna2-unknown-{system_os}"
-    elif any(x in device_name for x in "AMD", "Radeon") and "7" in device_name:
-        triple = f"rdna3-unknown-{system_os}"
-    elif any(x in device_name for x in "AMD", "Radeon"):
+    elif any(x in device_name for x in ("AMD", "Radeon")):
         triple = f"rdna3-unknown-{system_os}"
+
     # Intel Targets
     elif any(x in device_name for x in ("A770", "A750")):
         triple = f"arc-770-{system_os}"