From 92e66a0b3595e66a49ed3e6ae978927c1f1b7cf3 Mon Sep 17 00:00:00 2001
From: Stephen Jia <ssjia@meta.com>
Date: Fri, 19 Dec 2025 13:44:39 -0800
Subject: [PATCH] Add a setting to disable multithreading when compiler shaders
 during build (#16325)

Summary:
## Context

Currently, the `gen_vulkan_spv.py` script which orchestrates compute shader compilation for the Vulkan backend uses a thread pool to help parallelize shader compilation.

However, when Claude Code is instructed to build ExecuTorch via CMake, it runs into permission issues when trying to launch threads.

As a workaround, this PR adds a setting to control the number of threads that is used during shader compilation. 1 will result in serial compilation and -1 will result in querying the number of processors available.


Reviewed By: manuelcandales

Differential Revision: D89497641

Pulled By: SS-JIA
---
 backends/vulkan/cmake/ShaderLibrary.cmake |  7 ++++++
 backends/vulkan/runtime/gen_vulkan_spv.py | 30 +++++++++++++++++++----
 backends/vulkan/targets.bzl               |  3 +++
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/backends/vulkan/cmake/ShaderLibrary.cmake b/backends/vulkan/cmake/ShaderLibrary.cmake
index 16a60abf6f3..3d4db89ac14 100644
--- a/backends/vulkan/cmake/ShaderLibrary.cmake
+++ b/backends/vulkan/cmake/ShaderLibrary.cmake
@@ -53,6 +53,13 @@ function(gen_vulkan_shader_lib_cpp shaders_path)
     endif()
   endif()
 
+  # Add nthreads argument for shader compilation
+  if(DEFINED EXECUTORCH_VULKAN_SHADER_COMPILE_NTHREADS)
+    list(APPEND GEN_SPV_ARGS "--nthreads"
+         "${EXECUTORCH_VULKAN_SHADER_COMPILE_NTHREADS}"
+    )
+  endif()
+
   add_custom_command(
     COMMENT "Generating Vulkan Compute Shaders"
     OUTPUT ${VULKAN_SHADERGEN_OUT_PATH}/spv.cpp
diff --git a/backends/vulkan/runtime/gen_vulkan_spv.py b/backends/vulkan/runtime/gen_vulkan_spv.py
index ab709092351..238c9b1fbb2 100644
--- a/backends/vulkan/runtime/gen_vulkan_spv.py
+++ b/backends/vulkan/runtime/gen_vulkan_spv.py
@@ -892,6 +892,7 @@ def generateSPV(  # noqa: C901
         output_dir: str,
         cache_dir: Optional[str] = None,
         force_rebuild: bool = False,
+        nthreads: int = -1,
     ) -> Dict[str, str]:
         # The key of this dictionary is the full path to a generated source file. The
         # value is a tuple that contains 3 entries:
@@ -1118,11 +1119,21 @@ def compile_spirv(shader_paths_pair) -> Tuple[str, str]:
             gen_file_meta[gen_out_path] = (file_changed, include_list)
 
         # Parallelize SPIR-V compilation to optimize build time
-        with ThreadPool(os.cpu_count()) as pool:
-            for spv_out_path, glsl_out_path in pool.map(
-                compile_spirv, self.output_file_map.items()
-            ):
+        # Determine number of threads: -1 means use all CPU cores, 1 means sequential
+        num_processes = os.cpu_count() if nthreads == -1 else nthreads
+
+        if num_processes == 1:
+            # Sequential compilation (single-threaded)
+            for shader_pair in self.output_file_map.items():
+                spv_out_path, glsl_out_path = compile_spirv(shader_pair)
                 spv_to_glsl_map[spv_out_path] = glsl_out_path
+        else:
+            # Parallel compilation
+            with ThreadPool(num_processes) as pool:
+                for spv_out_path, glsl_out_path in pool.map(
+                    compile_spirv, self.output_file_map.items()
+                ):
+                    spv_to_glsl_map[spv_out_path] = glsl_out_path
 
         return spv_to_glsl_map
 
@@ -1443,6 +1454,12 @@ def main(argv: List[str]) -> int:
     parser.add_argument(
         "--env", metavar="KEY=VALUE", nargs="*", help="Set a number of key-value pairs"
     )
+    parser.add_argument(
+        "--nthreads",
+        type=int,
+        default=-1,
+        help="Number of threads for shader compilation. -1 (default) uses all available CPU cores, 1 uses sequential compilation.",
+    )
     options = parser.parse_args()
 
     env = DEFAULT_ENV
@@ -1477,7 +1494,10 @@ def main(argv: List[str]) -> int:
         replace_u16vecn=options.replace_u16vecn,
     )
     output_spv_files = shader_generator.generateSPV(
-        options.output_path, options.tmp_dir_path, options.force_rebuild
+        options.output_path,
+        options.tmp_dir_path,
+        options.force_rebuild,
+        options.nthreads,
     )
 
     genCppFiles(
diff --git a/backends/vulkan/targets.bzl b/backends/vulkan/targets.bzl
index 94c1f824633..64d6c672007 100644
--- a/backends/vulkan/targets.bzl
+++ b/backends/vulkan/targets.bzl
@@ -93,12 +93,15 @@ def vulkan_spv_shader_lib(name, spv_filegroups, is_fbcode = False, no_volk = Fal
     for target, subpath in spv_filegroups.items():
         glsl_paths.append("$(location {})/{}".format(target, subpath))
 
+    nthreads = read_config("etvk", "shader_compile_nthreads", "-1")
+
     genrule_cmd = (
         "$(exe {}) ".format(gen_vulkan_spv_target) +
         "--glsl-paths {} ".format(" ".join(glsl_paths)) +
         "--output-path $OUT " +
         "--glslc-path=$(exe {}) ".format(glslc_path) +
         "--tmp-dir-path=shader_cache " +
+        "--nthreads {} ".format(nthreads) +
         ("-f " if read_config("etvk", "force_shader_rebuild", "0") == "1" else " ") +
         select({
             "DEFAULT": "",