From 67b438eb9f6a25035ab85c037d87bf8cade3fcb1 Mon Sep 17 00:00:00 2001
From: Ean Garvey <ean.garvey@amd.com>
Date: Tue, 4 Jun 2024 01:46:48 -0500
Subject: [PATCH] take all ireert calls out of studio flow

---
 apps/shark_studio/api/utils.py | 125 +++++++++++++++++----------------
 1 file changed, 63 insertions(+), 62 deletions(-)

diff --git a/apps/shark_studio/api/utils.py b/apps/shark_studio/api/utils.py
index 8eebb97966..85a59ada35 100644
--- a/apps/shark_studio/api/utils.py
+++ b/apps/shark_studio/api/utils.py
@@ -65,6 +65,7 @@ def iree_target_map(device):
 
 
 def get_available_devices():
+    return ['rocm', 'cpu']
     def get_devices_by_name(driver_name):
 
         device_list = []
@@ -225,65 +226,65 @@ def get_all_devices(driver_name):
     return device_list_src
 
 
-def get_device_mapping(driver, key_combination=3):
-    """This method ensures consistent device ordering when choosing
-    specific devices for execution
-    Args:
-        driver (str): execution driver (vulkan, cuda, rocm, etc)
-        key_combination (int, optional): choice for mapping value for
-            device name.
-        1 : path
-        2 : name
-        3 : (name, path)
-        Defaults to 3.
-    Returns:
-        dict: map to possible device names user can input mapped to desired
-            combination of name/path.
-    """
-
-    driver = iree_device_map(driver)
-    device_list = get_all_devices(driver)
-    device_map = dict()
-
-    def get_output_value(dev_dict):
-        if key_combination == 1:
-            return f"{driver}://{dev_dict['path']}"
-        if key_combination == 2:
-            return dev_dict["name"]
-        if key_combination == 3:
-            return dev_dict["name"], f"{driver}://{dev_dict['path']}"
-
-    # mapping driver name to default device (driver://0)
-    device_map[f"{driver}"] = get_output_value(device_list[0])
-    for i, device in enumerate(device_list):
-        # mapping with index
-        device_map[f"{driver}://{i}"] = get_output_value(device)
-        # mapping with full path
-        device_map[f"{driver}://{device['path']}"] = get_output_value(device)
-    return device_map
-
-
-def get_opt_flags(model, precision="fp16"):
-    iree_flags = []
-    if len(cmd_opts.iree_vulkan_target_triple) > 0:
-        iree_flags.append(
-            f"-iree-vulkan-target-triple={cmd_opts.iree_vulkan_target_triple}"
-        )
-    if "rocm" in cmd_opts.device:
-        from shark.iree_utils.gpu_utils import get_iree_rocm_args
-
-        rocm_args = get_iree_rocm_args()
-        iree_flags.extend(rocm_args)
-    if cmd_opts.iree_constant_folding == False:
-        iree_flags.append("--iree-opt-const-expr-hoisting=False")
-        iree_flags.append(
-            "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807"
-        )
-    if cmd_opts.data_tiling == False:
-        iree_flags.append("--iree-opt-data-tiling=False")
-
-    if "vae" not in model:
-        # Due to lack of support for multi-reduce, we always collapse reduction
-        # dims before dispatch formation right now.
-        iree_flags += ["--iree-flow-collapse-reduction-dims"]
-    return iree_flags
\ No newline at end of file
+# def get_device_mapping(driver, key_combination=3):
+#     """This method ensures consistent device ordering when choosing
+#     specific devices for execution
+#     Args:
+#         driver (str): execution driver (vulkan, cuda, rocm, etc)
+#         key_combination (int, optional): choice for mapping value for
+#             device name.
+#         1 : path
+#         2 : name
+#         3 : (name, path)
+#         Defaults to 3.
+#     Returns:
+#         dict: map to possible device names user can input mapped to desired
+#             combination of name/path.
+#     """
+
+#     driver = iree_device_map(driver)
+#     device_list = get_all_devices(driver)
+#     device_map = dict()
+
+#     def get_output_value(dev_dict):
+#         if key_combination == 1:
+#             return f"{driver}://{dev_dict['path']}"
+#         if key_combination == 2:
+#             return dev_dict["name"]
+#         if key_combination == 3:
+#             return dev_dict["name"], f"{driver}://{dev_dict['path']}"
+
+#     # mapping driver name to default device (driver://0)
+#     device_map[f"{driver}"] = get_output_value(device_list[0])
+#     for i, device in enumerate(device_list):
+#         # mapping with index
+#         device_map[f"{driver}://{i}"] = get_output_value(device)
+#         # mapping with full path
+#         device_map[f"{driver}://{device['path']}"] = get_output_value(device)
+#     return device_map
+
+
+# def get_opt_flags(model, precision="fp16"):
+#     iree_flags = []
+#     if len(cmd_opts.iree_vulkan_target_triple) > 0:
+#         iree_flags.append(
+#             f"-iree-vulkan-target-triple={cmd_opts.iree_vulkan_target_triple}"
+#         )
+#     if "rocm" in cmd_opts.device:
+#         from shark.iree_utils.gpu_utils import get_iree_rocm_args
+
+#         rocm_args = get_iree_rocm_args()
+#         iree_flags.extend(rocm_args)
+#     if cmd_opts.iree_constant_folding == False:
+#         iree_flags.append("--iree-opt-const-expr-hoisting=False")
+#         iree_flags.append(
+#             "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807"
+#         )
+#     if cmd_opts.data_tiling == False:
+#         iree_flags.append("--iree-opt-data-tiling=False")
+
+#     if "vae" not in model:
+#         # Due to lack of support for multi-reduce, we always collapse reduction
+#         # dims before dispatch formation right now.
+#         iree_flags += ["--iree-flow-collapse-reduction-dims"]
+#     return iree_flags
\ No newline at end of file