From 632994116b084212d9519635307a0b7dcc9878dd Mon Sep 17 00:00:00 2001
From: Stephen Jia <ssjia@meta.com>
Date: Mon, 10 Mar 2025 11:16:18 -0700
Subject: [PATCH] [ET-VK][ez] Retry with no optimization during shader
 compilation

## Context

In some cases, a shader can be compiled successfully with `glslc`, but if an optimization flag such as `-O` or `-Os` is used, the compilation will fail with a message similar to

```
internal error: compilation succeeded but failed to optimize: Invalid use of 8- or 16-bit result\n  %133 = OpExtInst %half %1 Pow %125 %132
```

Which will cause the shader compilation script to fail.

This diff introduces a check that if `glslc` complains about not being able to optimize the SPIR-V, then retry compilation without optimization arguments.

Differential Revision: [D70906097](https://our.internmc.facebook.com/intern/diff/D70906097/)

[ghstack-poisoned]
---
 backends/vulkan/runtime/gen_vulkan_spv.py | 56 ++++++++++++++---------
 1 file changed, 34 insertions(+), 22 deletions(-)

diff --git a/backends/vulkan/runtime/gen_vulkan_spv.py b/backends/vulkan/runtime/gen_vulkan_spv.py
index 1db75b7edc9..a3d214f5ae8 100644
--- a/backends/vulkan/runtime/gen_vulkan_spv.py
+++ b/backends/vulkan/runtime/gen_vulkan_spv.py
@@ -549,7 +549,12 @@ def __init__(
 
         self.env = env
         self.glslc_path = glslc_path
-        self.glslc_flags = glslc_flags
+        self.glslc_flags = glslc_flags.split()
+        self.glslc_flags_no_opt = self.glslc_flags.copy()
+        if "-O" in self.glslc_flags_no_opt:
+            self.glslc_flags_no_opt.remove("-O")
+        if "-Os" in self.glslc_flags_no_opt:
+            self.glslc_flags_no_opt.remove("-Os")
         self.replace_u16vecn = replace_u16vecn
 
         self.glsl_src_files: Dict[str, str] = {}
@@ -751,30 +756,37 @@ def process_shader(shader_paths_pair):
             if self.glslc_path is not None:
                 spv_out_path = os.path.join(output_dir, f"{shader_name}.spv")
 
-                cmd = (
-                    [
-                        self.glslc_path,
-                        "-fshader-stage=compute",
-                        glsl_out_path,
-                        "-o",
-                        spv_out_path,
-                        "--target-env=vulkan1.1",
-                        "-Werror",
-                    ]
-                    + [
-                        arg
-                        for src_dir_path in self.src_dir_paths
-                        for arg in ["-I", src_dir_path]
-                    ]
-                    + self.glslc_flags.split()
-                )
+                cmd_base = [
+                    self.glslc_path,
+                    "-fshader-stage=compute",
+                    glsl_out_path,
+                    "-o",
+                    spv_out_path,
+                    "--target-env=vulkan1.1",
+                    "-Werror",
+                ] + [
+                    arg
+                    for src_dir_path in self.src_dir_paths
+                    for arg in ["-I", src_dir_path]
+                ]
+                cmd = cmd_base + self.glslc_flags
 
                 try:
-                    subprocess.check_call(cmd)
+                    subprocess.run(cmd, check=True, capture_output=True, text=True)
                 except subprocess.CalledProcessError as e:
-                    raise RuntimeError(
-                        f"Failed to compile {os.getcwd()}/{glsl_out_path}"
-                    ) from e
+                    opt_fail = "compilation succeeded but failed to optimize"
+                    err_msg_base = f"Failed to compile {os.getcwd()}/{glsl_out_path}: "
+                    if opt_fail in e.stderr or opt_fail in e.stdout:
+                        cmd_no_opt = cmd_base + self.glslc_flags_no_opt
+                        try:
+                            subprocess.run(cmd_no_opt, check=True, capture_output=True)
+                        except subprocess.CalledProcessError as e_no_opt:
+                            raise RuntimeError(
+                                f"{err_msg_base} {e_no_opt.stderr}"
+                            ) from e_no_opt
+
+                    else:
+                        raise RuntimeError(f"{err_msg_base} {e.stderr}") from e
 
                 return (spv_out_path, glsl_out_path)