github/AMD-SHARK-Studio
1
1
Fork 0
mirror of https://github.com/nod-ai/AMD-SHARK-Studio.git synced 2026-04-03 03:00:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

Add more NVIDIA targets (#640)

Browse Source
This commit is contained in:
nirvedhmeshram
2022-12-15 11:24:38 -06:00
committed by GitHub
parent 24a16a4cfe
commit 2928179331
1 changed files with 23 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
shark/iree_utils/vulkan_utils.py
View File

@@ -50,12 +50,14 @@ def get_vulkan_triple_flag(extra_args=[]):
return None
system_os = get_os_name()
vulkan_device = get_vulkan_device_name()
# Apple Targets
if all(x in vulkan_device for x in ("Apple", "M1")):
print(f"Found {vulkan_device} Device. Using m1-moltenvk-macos")
return "-iree-vulkan-target-triple=m1-moltenvk-macos"
elif all(x in vulkan_device for x in ("Apple", "M2")):
print("Found Apple M2 Device. Using m1-moltenvk-macos")
return "-iree-vulkan-target-triple=m1-moltenvk-macos"
# Nvidia Targets
elif all(x in vulkan_device for x in ("A100", "SXM4")):
print(
f"Found {vulkan_device} Device. Using ampere-rtx3080-{system_os}"
@@ -71,6 +73,27 @@ def get_vulkan_triple_flag(extra_args=[]):
f"Found {vulkan_device} Device. Using ampere-rtx3090-{system_os}"
)
return f"-iree-vulkan-target-triple=ampere-rtx3090-{system_os}"
elif all(x in vulkan_device for x in ("RTX", "4000")):
print(
f"Found {vulkan_device} Device. Using turing-rtx4000-{system_os}"
)
return f"-iree-vulkan-target-triple=turing-rtx4000-{system_os}"
elif all(x in vulkan_device for x in ("RTX", "5000")):
print(
f"Found {vulkan_device} Device. Using turing-rtx5000-{system_os}"
)
return f"-iree-vulkan-target-triple=turing-rtx5000-{system_os}"
elif all(x in vulkan_device for x in ("RTX", "6000")):
print(
f"Found {vulkan_device} Device. Using turing-rtx6000-{system_os}"
)
return f"-iree-vulkan-target-triple=turing-rtx6000-{system_os}"
elif all(x in vulkan_device for x in ("RTX", "8000")):
print(
f"Found {vulkan_device} Device. Using turing-rtx8000-{system_os}"
)
return f"-iree-vulkan-target-triple=turing-rtx8000-{system_os}"
# Amd Targets
elif all(x in vulkan_device for x in ("AMD", "7900")):
print(f"Found {vulkan_device} Device. Using rdna3-7900-{system_os}")
return f"-iree-vulkan-target-triple=rdna3-7900-{system_os}"