Merge remote-tracking branch 'origin/main' into tiktoken

Author: kirklandsign

backends/vulkan/CMakeLists.txt

@@ -61,6 +61,8 @@ list(APPEND VULKAN_CXX_FLAGS "-DUSE_VULKAN_VOLK")
 # vulkan API files
 
 file(GLOB_RECURSE vulkan_api_cpp ${RUNTIME_PATH}/api/*)
+file(GLOB_RECURSE vulkan_vkapi_cpp ${RUNTIME_PATH}/vk_api/*)
+list(APPEND vulkan_api_cpp ${vulkan_vkapi_cpp})
 list(APPEND vulkan_api_cpp ${VOLK_PATH}/volk.c)
 
 # vulkan ComputeGraph files

backends/vulkan/runtime/gen_vulkan_spv.py

@@ -15,6 +15,7 @@
 import re
 import sys
 from itertools import product
+from multiprocessing.pool import ThreadPool
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 import subprocess
@@ -620,9 +621,12 @@ def constructOutputMap(self) -> None:
 
     def generateSPV(self, output_dir: str) -> Dict[str, str]:
         output_file_map = {}
-        for shader_name in self.output_shader_map:
-            source_glsl = self.output_shader_map[shader_name][0]
-            shader_params = self.output_shader_map[shader_name][1]
+
+        def process_shader(shader_paths_pair):
+            shader_name = shader_paths_pair[0]
+
+            source_glsl = shader_paths_pair[1][0]
+            shader_params = shader_paths_pair[1][1]
 
             with codecs.open(source_glsl, "r", encoding="utf-8") as input_file:
                 input_text = input_file.read()
@@ -652,9 +656,15 @@ def generateSPV(self, output_dir: str) -> Dict[str, str]:
                 ]
 
                 print("glslc cmd:", cmd)
-                # pyre-ignore
                 subprocess.check_call(cmd)
 
+                return (spv_out_path, glsl_out_path)
+
+        # Parallelize shader compilation as much as possible to optimize build time.
+        with ThreadPool(os.cpu_count()) as pool:
+            for spv_out_path, glsl_out_path in pool.map(
+                process_shader, self.output_shader_map.items()
+            ):
                 output_file_map[spv_out_path] = glsl_out_path
 
         return output_file_map

exir/lowered_backend_module.py

@@ -139,14 +139,15 @@ def buffer(
         segment_alignment: int = 4096,
         constant_tensor_alignment: Optional[int] = None,
         delegate_alignment: Optional[int] = None,
+        memory_planning: MemoryPlanningPass = None,
     ) -> bytes:
         """
         Returns a buffer containing the serialized ExecuTorch binary.
         """
         # TODO(T181463742): avoid calling bytes(..) which incurs large copies.
         out = bytes(
             _serialize_pte_binary(
-                program=self.program(),
+                program=self.program(memory_planning=memory_planning),
                 extract_delegate_segments=extract_delegate_segments,
                 segment_alignment=segment_alignment,
                 constant_tensor_alignment=constant_tensor_alignment,
@@ -157,7 +158,11 @@ def buffer(
 
     # TODO(chenlai): re-consider recapture instead of manually constructing the program because
     # the meta data construction is done manually.
-    def program(self, emit_stacktrace: bool = False) -> Program:
+    def program(
+        self,
+        emit_stacktrace: bool = False,
+        memory_planning: MemoryPlanningPass = None,
+    ) -> Program:
         # Fix autodpes introuces cyclic dependencies:
         # program -> verifier -> lowered_backend_module -> program
         # @manual
@@ -319,9 +324,9 @@ def program(self, emit_stacktrace: bool = False) -> Program:
             example_inputs=None,
             verifier=lowered_exported_program.verifier,
         )
-        exported_program = _transform(
-            exported_program, SpecPropPass(), MemoryPlanningPass("greedy")
-        )
+        if memory_planning is None:
+            memory_planning = MemoryPlanningPass("greedy")
+        exported_program = _transform(exported_program, SpecPropPass(), memory_planning)
         emitted_program = emit_program(
             exported_program, emit_stacktrace=emit_stacktrace
         ).program