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

Enable 64bit for models

Browse Source
This commit is contained in:
Anush Elangovan
2022-06-20 20:22:28 +00:00
parent 7198f0db0e
commit c66e285e2b
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
shark/iree_utils.py
View File

@@ -148,6 +148,11 @@ def get_iree_frontend_args(frontend):
# Frontend not found.
return []
def get_iree_common_args():
return [
"--iree-stream-resource-index-bits=64",
"--iree-vm-target-index-bits=64"
]
def compile_module_to_flatbuffer(module, device, frontend, func_name,
model_config_path):
@@ -155,6 +160,7 @@ def compile_module_to_flatbuffer(module, device, frontend, func_name,
input_type = ""
args = get_iree_frontend_args(frontend)
args += get_iree_device_args(device)
args += get_iree_common_args()
if frontend in ["tensorflow", "tf"]:
input_type = "mhlo"