Add ROCM parameters (huggingface#335)

Author: erman-gurses

shark/iree_utils/_common.py

@@ -76,6 +76,11 @@ def check_device_drivers(device):
             return True
     elif device == "cpu":
         return False
+    elif device == "rocm":
+        try:
+            subprocess.check_output("rocminfo")
+        except Exception:
+            return True
     # Unknown device.
     else:
         return True
@@ -89,5 +94,7 @@ def device_driver_info(device):
         return "nvidia-smi not found, please install the required drivers from https://www.nvidia.in/Download/index.aspx?lang=en-in"
     elif device in ["metal", "vulkan"]:
         return "vulkaninfo not found, Install from https://vulkan.lunarg.com/sdk/home or your distribution"
+    elif device == "rocm":
+        return "rocm info not found. Please install rocm"
     else:
         return f"{device} is not supported."

shark/iree_utils/compile_utils.py

@@ -31,6 +31,10 @@ def get_iree_device_args(device):
         from shark.iree_utils.vulkan_utils import get_iree_vulkan_args
 
         return get_iree_vulkan_args()
+    if device == "rocm":
+        from shark.iree_utils.gpu_utils import get_iree_rocm_args
+
+        return get_iree_rocm_args()
     return []
 
 

shark/iree_utils/gpu_utils.py

@@ -35,6 +35,18 @@ def get_iree_gpu_args():
         return ["--iree-hal-cuda-disable-loop-nounroll-wa"]
 
 
+# Get the default gpu args given the architecture.
+def get_iree_rocm_args():
+    ireert.flags.FUNCTION_INPUT_VALIDATION = False
+    # TODO: find a way to get arch from code.
+    rocm_arch = "gfx908"
+    return [
+        f"--iree-rocm-target-chip={rocm_arch}",
+        "--iree-rocm-link-bc=true",
+        "--iree-rocm-bc-dir=/opt/rocm/amdgcn/bitcode",
+    ]
+
+
 # Some constants taken from cuda.h
 CUDA_SUCCESS = 0
 CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT = 16