From 8d65456b7a65901f19d4fc6cca3ec0040afdb103 Mon Sep 17 00:00:00 2001
From: PhaneeshB <b.phaneesh@gmail.com>
Date: Wed, 26 Jul 2023 20:15:42 +0530
Subject: [PATCH] Move vulkan runtime flags to shark_args

---
 .../stable_diffusion/src/utils/stable_args.py | 28 -------------------
 apps/stable_diffusion/src/utils/utils.py      |  7 ++---
 shark/iree_utils/vulkan_utils.py              | 10 +++++++
 shark/parser.py                               | 28 +++++++++++++++++++
 4 files changed, 40 insertions(+), 33 deletions(-)

diff --git a/apps/stable_diffusion/src/utils/stable_args.py b/apps/stable_diffusion/src/utils/stable_args.py
index b659081f..c2601fbb 100644
--- a/apps/stable_diffusion/src/utils/stable_args.py
+++ b/apps/stable_diffusion/src/utils/stable_args.py
@@ -425,34 +425,6 @@ p.add_argument(
     help="Specify target triple for metal.",
 )
 
-p.add_argument(
-    "--vulkan_debug_utils",
-    default=False,
-    action=argparse.BooleanOptionalAction,
-    help="Profiles vulkan device and collects the .rdc info.",
-)
-
-p.add_argument(
-    "--vulkan_large_heap_block_size",
-    default="2073741824",
-    help="Flag for setting VMA preferredLargeHeapBlockSize for "
-    "vulkan device, default is 4G.",
-)
-
-p.add_argument(
-    "--vulkan_validation_layers",
-    default=False,
-    action=argparse.BooleanOptionalAction,
-    help="Flag for disabling vulkan validation layers when benchmarking.",
-)
-
-p.add_argument(
-    "--vulkan_vma_allocator",
-    default=True,
-    action=argparse.BooleanOptionalAction,
-    help="Flag for enabling / disabling vulkan Vulkan VMA Allocator.",
-)
-
 ##############################################################################
 # Misc. Debug and Optimization flags
 ##############################################################################
diff --git a/apps/stable_diffusion/src/utils/utils.py b/apps/stable_diffusion/src/utils/utils.py
index db607c92..7e82b71e 100644
--- a/apps/stable_diffusion/src/utils/utils.py
+++ b/apps/stable_diffusion/src/utils/utils.py
@@ -22,6 +22,7 @@ from shark.shark_importer import import_with_fx
 from shark.iree_utils.vulkan_utils import (
     set_iree_vulkan_runtime_flags,
     get_vulkan_target_triple,
+    get_iree_vulkan_runtime_flags,
 )
 from shark.iree_utils.metal_utils import get_metal_target_triple
 from shark.iree_utils.gpu_utils import get_cuda_sm_cc
@@ -183,11 +184,7 @@ def compile_through_fx(
 
 
 def set_iree_runtime_flags():
-    vulkan_runtime_flags = [
-        f"--vulkan_large_heap_block_size={args.vulkan_large_heap_block_size}",
-        f"--vulkan_validation_layers={'true' if args.vulkan_validation_layers else 'false'}",
-        f"--vulkan_vma_allocator={'true' if args.vulkan_vma_allocator else 'false'}",
-    ]
+    vulkan_runtime_flags = get_iree_vulkan_runtime_flags()
     if args.enable_rgp:
         vulkan_runtime_flags += [
             f"--enable_rgp=true",
diff --git a/shark/iree_utils/vulkan_utils.py b/shark/iree_utils/vulkan_utils.py
index 06d3c1d7..3675a869 100644
--- a/shark/iree_utils/vulkan_utils.py
+++ b/shark/iree_utils/vulkan_utils.py
@@ -19,6 +19,7 @@ from shark.iree_utils._common import run_cmd
 import iree.runtime as ireert
 from sys import platform
 from shark.iree_utils.vulkan_target_env_utils import get_vulkan_target_env_flag
+from shark.parser import shark_args
 
 
 def get_vulkan_device_name(device_num=0):
@@ -171,6 +172,15 @@ def get_iree_vulkan_args(device_num=0, extra_args=[]):
     return res_vulkan_flag
 
 
+def get_iree_vulkan_runtime_flags():
+    vulkan_runtime_flags = [
+        f"--vulkan_large_heap_block_size={shark_args.vulkan_large_heap_block_size}",
+        f"--vulkan_validation_layers={'true' if shark_args.vulkan_validation_layers else 'false'}",
+        f"--vulkan_vma_allocator={'true' if shark_args.vulkan_vma_allocator else 'false'}",
+    ]
+    return vulkan_runtime_flags
+
+
 def set_iree_vulkan_runtime_flags(flags):
     for flag in flags:
         ireert.flags.parse_flags(flag)
diff --git a/shark/parser.py b/shark/parser.py
index 7c51290e..e7f3285a 100644
--- a/shark/parser.py
+++ b/shark/parser.py
@@ -126,4 +126,32 @@ parser.add_argument(
     help="passthrough flag for the iree flag of the same name. If None, defaults to cpu-count",
 )
 
+parser.add_argument(
+    "--vulkan_debug_utils",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="Profiles vulkan device and collects the .rdc info.",
+)
+
+parser.add_argument(
+    "--vulkan_large_heap_block_size",
+    default="2073741824",
+    help="Flag for setting VMA preferredLargeHeapBlockSize for "
+    "vulkan device, default is 4G.",
+)
+
+parser.add_argument(
+    "--vulkan_validation_layers",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="Flag for disabling vulkan validation layers when benchmarking.",
+)
+
+parser.add_argument(
+    "--vulkan_vma_allocator",
+    default=True,
+    action=argparse.BooleanOptionalAction,
+    help="Flag for enabling / disabling Vulkan VMA Allocator.",
+)
+
 shark_args, unknown = parser.parse_known_args()