From c3c57fe6b0827a4f0e249ac947c44c7cb0ca15fe Mon Sep 17 00:00:00 2001
From: Stephen Jia <ssjia@meta.com>
Date: Thu, 25 Jul 2024 11:22:30 -0700
Subject: [PATCH] Enable SPIR-V compiler optimization (#4402)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/4402

Call the SPIR-V compiler with the `-O` flag, which enables optimizations when compiling GLSL to SPIR-V. The `-Os` flag (which tries to minimize SPIR-V size) was tested as well, but resulted in (very) slightly worse performance.

Reviewed By: jorgep31415

Differential Revision: D60193514
---
 backends/vulkan/runtime/gen_vulkan_spv.py | 16 +++++++++++++---
 backends/vulkan/targets.bzl               |  2 ++
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/backends/vulkan/runtime/gen_vulkan_spv.py b/backends/vulkan/runtime/gen_vulkan_spv.py
index c9e3aaa31ea..9b4ca6eb35d 100644
--- a/backends/vulkan/runtime/gen_vulkan_spv.py
+++ b/backends/vulkan/runtime/gen_vulkan_spv.py
@@ -482,6 +482,7 @@ def __init__(
         src_dir_paths: Union[str, List[str]],
         env: Dict[Any, Any],
         glslc_path: Optional[str],
+        glslc_flags: str = "",
     ) -> None:
         if isinstance(src_dir_paths, str):
             self.src_dir_paths = [src_dir_paths]
@@ -490,6 +491,7 @@ def __init__(
 
         self.env = env
         self.glslc_path = glslc_path
+        self.glslc_flags = glslc_flags
 
         self.glsl_src_files: Dict[str, str] = {}
         self.template_yaml_files: List[str] = []
@@ -680,7 +682,7 @@ def process_shader(shader_paths_pair):
                     arg
                     for src_dir_path in self.src_dir_paths
                     for arg in ["-I", src_dir_path]
-                ]
+                ] + self.glslc_flags.split()
 
                 subprocess.check_call(cmd)
 
@@ -966,6 +968,8 @@ def main(argv: List[str]) -> int:
     parser.add_argument("-c", "--glslc-path", required=True, help="")
     parser.add_argument("-t", "--tmp-dir-path", required=True, help="/tmp")
     parser.add_argument("-o", "--output-path", required=True, help="")
+    parser.add_argument("--optimize_size", action="store_true", help="")
+    parser.add_argument("--optimize", action="store_true", help="")
     parser.add_argument(
         "--env", metavar="KEY=VALUE", nargs="*", help="Set a number of key-value pairs"
     )
@@ -983,8 +987,14 @@ def main(argv: List[str]) -> int:
 
     if not os.path.exists(options.tmp_dir_path):
         os.makedirs(options.tmp_dir_path)
-
-    shader_generator = SPVGenerator(options.glsl_paths, env, options.glslc_path)
+    
+    glslc_flags = ""
+    if options.optimize_size:
+        glslc_flags += "-Os"
+    elif options.optimize:
+        glslc_flags += "-O"
+
+    shader_generator = SPVGenerator(options.glsl_paths, env, options.glslc_path, glslc_flags)
     output_spv_files = shader_generator.generateSPV(options.tmp_dir_path)
 
     genCppFiles(
diff --git a/backends/vulkan/targets.bzl b/backends/vulkan/targets.bzl
index 981552f17ab..0b824f6f481 100644
--- a/backends/vulkan/targets.bzl
+++ b/backends/vulkan/targets.bzl
@@ -22,6 +22,8 @@ def vulkan_spv_shader_lib(name, spv_filegroups, is_fbcode = False):
         "--glslc-path=$(exe {})".format(glslc_path),
         "--tmp-dir-path=$OUT",
     ]
+    if not is_fbcode:
+        genrule_cmd += ["--optimize"]
 
     genrule_name = "gen_{}_cpp".format(name)
     runtime.genrule(