Add AMD cpu device

nod-ai · Jun 23, 2023 · 28e0919 · 28e0919
1 parent 28f4d44
commit 28e0919
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/shark/iree_utils/_common.py b/shark/iree_utils/_common.py
@@ -63,6 +63,7 @@ def get_supported_device_list():
 _IREE_DEVICE_MAP = {
     "cpu": "local-task",
     "cpu-task": "local-task",
+    "AMD-AIE": "local-task",
     "cpu-sync": "local-sync",
     "cuda": "cuda",
     "vulkan": "vulkan",
@@ -81,6 +82,7 @@ def iree_target_map(device):
 _IREE_TARGET_MAP = {
     "cpu": "llvm-cpu",
     "cpu-task": "llvm-cpu",
+    "AMD-AIE": "llvm-cpu",
     "cpu-sync": "llvm-cpu",
     "cuda": "cuda",
     "vulkan": "vulkan",