Add flag for enabling rgp from the main.py SD script (huggingface#533)

Author: qedawkins

shark/examples/shark_inference/stable_diffusion/README.md

@@ -23,6 +23,12 @@ python main.py --precision="fp16" --device="vulkan" --iree-vulkan-target-triple=
 python shark/examples/shark_inference/stable_diffusion/main.py --precision=fp16 --device=vulkan --steps=50 --save_vmfb
 ```
 
+## Capture an RGP trace
+
+```shell
+python shark/examples/shark_inference/stable_diffusion/main.py --precision=fp16 --device=vulkan --steps=50 --save_vmfb --enable_rgp
+```
+
 ## Run the vae module with iree-benchmark-module (NCHW, fp16, vulkan, for example):
 
 ```shell

shark/examples/shark_inference/stable_diffusion/stable_args.py

@@ -118,4 +118,11 @@
     help="flag for setting VMA preferredLargeHeapBlockSize for vulkan device, default is 4G",
 )
 
+p.add_argument(
+    "--enable_rgp",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="flag for inserting debug frames between iterations for use with rgp.",
+)
+
 args = p.parse_args()

shark/examples/shark_inference/stable_diffusion/utils.py

@@ -70,6 +70,11 @@ def set_iree_runtime_flags():
     vulkan_runtime_flags = [
         f"--vulkan_large_heap_block_size={args.vulkan_large_heap_block_size}",
     ]
+    if args.enable_rgp:
+        vulkan_runtime_flags += [
+            f"--enable_rgp=true",
+            f"--vulkan_debug_utils=true",
+        ]
     if "vulkan" in args.device:
         set_iree_vulkan_runtime_flags(flags=vulkan_runtime_flags)