github/SHARK-Studio
1
1
Fork 0
mirror of https://github.com/nod-ai/SHARK-Studio.git synced 2026-01-09 13:57:54 -05:00
Code Issues Packages Projects Releases Wiki Activity

Fix local-task vs dylib naming changes

Browse Source
This commit is contained in:
Anush Elangovan
2022-06-12 07:21:48 +00:00
parent c0be20d835
commit cd8fe34baa
1 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
shark/iree_utils.py
View File

@@ -25,6 +25,15 @@ import re
import sys
IREE_DEVICE_MAP = {
"cpu": "local-task",
"gpu": "cuda",
"cuda": "cuda",
"vulkan": "vulkan",
"metal": "vulkan",
"rocm": "rocm"
}
IREE_TARGET_MAP = {
"cpu": "dylib",
"gpu": "cuda",
"cuda": "cuda",
@@ -33,6 +42,7 @@ IREE_DEVICE_MAP = {
"rocm": "rocm"
}
UNIT_TO_SECOND_MAP = {"ms": 0.001, "s": 1}
@@ -174,14 +184,14 @@ def compile_module_to_flatbuffer(module, device, frontend, func_name,
# Currently for MHLO/TOSA.
flatbuffer_blob = ireec.compile_str(
module,
target_backends=[IREE_DEVICE_MAP[device]],
target_backends=[IREE_TARGET_MAP[device]],
extra_args=args,
input_type=input_type)
else:
# Currently for Torch.
flatbuffer_blob = ireec.compile_str(
str(module),
target_backends=[IREE_DEVICE_MAP[device]],
target_backends=[IREE_TARGET_MAP[device]],
extra_args=args)
return flatbuffer_blob