diff --git a/backends/vulkan/cmake/ShaderLibrary.cmake b/backends/vulkan/cmake/ShaderLibrary.cmake index d7e55cd3592..53d2839a45e 100644 --- a/backends/vulkan/cmake/ShaderLibrary.cmake +++ b/backends/vulkan/cmake/ShaderLibrary.cmake @@ -81,8 +81,14 @@ function(gen_vulkan_shader_lib_cpp shaders_path) endif() # Ninja cannot expand wildcards (*) in DEPENDS lists. - file(GLOB VULKAN_SHADERS "${shaders_path}/*.glsl" "${shaders_path}/*.glslh" - "${shaders_path}/*.yaml" "${shaders_path}/*.h" + file( + GLOB_RECURSE + VULKAN_SHADERS + CONFIGURE_DEPENDS + "${shaders_path}/*.glsl" + "${shaders_path}/*.glslh" + "${shaders_path}/*.yaml" + "${shaders_path}/*.h" ) add_custom_command( diff --git a/backends/vulkan/runtime/gen_vulkan_spv.py b/backends/vulkan/runtime/gen_vulkan_spv.py index c9b054f1e73..5016e9c48db 100644 --- a/backends/vulkan/runtime/gen_vulkan_spv.py +++ b/backends/vulkan/runtime/gen_vulkan_spv.py @@ -748,6 +748,7 @@ def __init__( self.src_files: Dict[str, str] = {} self.template_yaml_files: List[str] = [] + self.yaml_relative_dirs: Dict[str, str] = {} self.addSrcAndYamlFiles(self.src_dir_paths) self.shader_template_params: Dict[Any, Any] = {} @@ -759,20 +760,30 @@ def __init__( def addSrcAndYamlFiles(self, src_dir_paths: List[str]) -> None: for src_path in src_dir_paths: + src_path = os.path.abspath(src_path) # Collect glsl source files src_files_list = glob.glob( os.path.join(src_path, "**", "*.[gh]lsl*"), recursive=True ) + glob.glob(os.path.join(src_path, "**", "*.h"), recursive=True) - for file in src_files_list: + for file in sorted(src_files_list): if len(file) > 1: - self.src_files[extract_filename(file, keep_ext=False)] = file + relative_path = os.path.relpath(file, src_path) + src_name = os.path.join( + os.path.dirname(relative_path), + extract_filename(relative_path, keep_ext=False), + ) + if src_name in self.src_files: + raise KeyError(f"{relative_path} source file is defined twice") + self.src_files[src_name] = file # Collect template yaml files yaml_files = glob.glob( os.path.join(src_path, "**", "*.yaml"), recursive=True ) - for file in yaml_files: + for file in sorted(yaml_files): if len(file) > 1: self.template_yaml_files.append(file) + relative_path = os.path.relpath(file, src_path) + self.yaml_relative_dirs[file] = os.path.dirname(relative_path) def generateVariantCombinations( # noqa: C901 self, @@ -839,14 +850,17 @@ def parseTemplateYaml(self, yaml_file: str) -> None: # noqa: C901 with open(yaml_file) as f: contents = yaml.load(f, Loader=UniqueKeyLoader) for template_name, params_dict in contents.items(): - if template_name in self.shader_template_params: - raise KeyError(f"{template_name} params file is defined twice") + template_path = os.path.normpath( + os.path.join(self.yaml_relative_dirs[yaml_file], template_name) + ) + if template_path in self.shader_template_params: + raise KeyError(f"{template_path} params file is defined twice") default_params = params_dict["parameter_names_with_default_values"] default_params["YAML_SRC_FULLPATH"] = yaml_file params_names = set(default_params.keys()).union({"NAME"}) - self.shader_template_params[template_name] = [] + self.shader_template_params[template_path] = [] default_iterated_params = params_dict.get( "generate_variant_forall", None @@ -907,7 +921,7 @@ def parseTemplateYaml(self, yaml_file: str) -> None: # noqa: C901 default_params_copy["NAME"] = variant_name default_params_copy["VARIANT_NAME"] = variant["NAME"] - self.shader_template_params[template_name].append( + self.shader_template_params[template_path].append( default_params_copy ) else: @@ -916,7 +930,7 @@ def parseTemplateYaml(self, yaml_file: str) -> None: # noqa: C901 if key not in reserved_yaml_keys: default_params_copy[key] = variant[key] - self.shader_template_params[template_name].append( + self.shader_template_params[template_path].append( default_params_copy ) @@ -944,18 +958,36 @@ def create_shader_params( return shader_params def constructOutputMap(self) -> None: + shader_names = set() + + def add_output( + output_name: str, + src_file_fullpath: str, + params: Dict[str, str], + ) -> None: + if output_name in self.output_file_map: + raise KeyError(f"{output_name} output file is defined twice") + if extract_extension(src_file_fullpath) == "glsl": + shader_name = os.path.basename(output_name) + if shader_name in shader_names: + raise KeyError(f"{shader_name} shader is defined twice") + shader_names.add(shader_name) + self.output_file_map[output_name] = (src_file_fullpath, params) + for src_name, params in self.shader_template_params.items(): for variant in params: src_file_fullpath = self.src_files[src_name] - - self.output_file_map[variant["NAME"]] = ( + output_name = os.path.join(os.path.dirname(src_name), variant["NAME"]) + add_output( + output_name, src_file_fullpath, self.create_shader_params(variant), ) for src_name, src_file_fullpath in self.src_files.items(): if src_name not in self.shader_template_params: - self.output_file_map[src_name] = ( + add_output( + src_name, src_file_fullpath, self.create_shader_params(), ) @@ -1054,21 +1086,55 @@ def file_has_changed(gen_file_path, cached_file_path): cached_checksum = self.get_md5_checksum(cached_file_path) return current_checksum != cached_checksum - def any_sources_changed(gen_file_path, output_dir): + def get_include_dirs(gen_file_path): + """ + Return the directories searched for a generated file's includes: its + own directory, then each ancestor up to output_dir, nearest first. + """ + include_dir = os.path.dirname(gen_file_path) + include_dirs = [include_dir] + while include_dir != output_dir: + parent_dir = os.path.dirname(include_dir) + if parent_dir == include_dir: + raise ValueError( + f"{gen_file_path} is outside output directory {output_dir}" + ) + include_dir = parent_dir + include_dirs.append(include_dir) + return include_dirs + + def any_sources_changed(gen_file_path, visited=None): """ Given the path to a generated source file, check the gen_file_meta dict to determine if the ANY of the source files contributing to the compilation of this file were changed since the last successful compilation. """ + if visited is None: + visited = set() + if gen_file_path in visited: + return False + if gen_file_path not in gen_file_meta: + return True + visited.add(gen_file_path) + gen_file_changed, includes_list = gen_file_meta[gen_file_path] - any_changed = gen_file_changed + if gen_file_changed: + return True + include_dirs = get_include_dirs(gen_file_path) for included_file in includes_list: - included_file_path = os.path.join(output_dir, included_file) - any_changed = any_changed or any_sources_changed( - included_file_path, output_dir + candidate_paths = [ + os.path.normpath(os.path.join(include_dir, included_file)) + for include_dir in include_dirs + ] + included_file_path = next( + (path for path in candidate_paths if path in gen_file_meta), + # Untracked includes must disable the cache. + candidate_paths[-1], ) + if any_sources_changed(included_file_path, visited): + return True - return any_changed + return False def generate_src_file(shader_paths_pair) -> Tuple[bool, List[str]]: """ @@ -1102,6 +1168,8 @@ def generate_src_file(shader_paths_pair) -> Tuple[bool, List[str]]: cached_gen_out_path = os.path.join( cache_dir, f"{src_file_name}.{out_file_ext}" ) + os.makedirs(os.path.dirname(gen_out_path), exist_ok=True) + os.makedirs(os.path.dirname(cached_gen_out_path), exist_ok=True) # Execute codegen to generate the output file with codecs.open(template_file_path, "r", encoding="utf-8") as input_file: @@ -1191,7 +1259,7 @@ def compile_spirv(shader_paths_pair) -> Tuple[str, str]: # in the cache. cached_spv_out_path = os.path.join(cache_dir, f"{src_file_name}.spv") - can_use_cached = not any_sources_changed(gen_out_path, output_dir) + can_use_cached = not any_sources_changed(gen_out_path) if can_use_cached and os.path.exists(cached_spv_out_path): shutil.copyfile(cached_spv_out_path, spv_out_path) return (spv_out_path, gen_out_path) @@ -1207,9 +1275,9 @@ def compile_spirv(shader_paths_pair) -> Tuple[str, str]: spv_out_path, "--target-env=vulkan{}".format(vk_version), "-Werror", - "-I", - output_dir, ] + for include_dir in get_include_dirs(gen_out_path): + cmd_base += ["-I", include_dir] cmd = cmd_base + self.glslc_flags try: diff --git a/backends/vulkan/targets.bzl b/backends/vulkan/targets.bzl index d2dee601399..aaa96eccb63 100644 --- a/backends/vulkan/targets.bzl +++ b/backends/vulkan/targets.bzl @@ -184,7 +184,7 @@ def define_common_targets(is_fbcode = False): runtime.filegroup( name = "vulkan_graph_runtime_shaders", srcs = native.glob([ - "runtime/graph/ops/glsl/*", + "runtime/graph/ops/glsl/**/*", ]), ) diff --git a/backends/vulkan/test/custom_ops/targets.bzl b/backends/vulkan/test/custom_ops/targets.bzl index 1895180495a..e3631298119 100644 --- a/backends/vulkan/test/custom_ops/targets.bzl +++ b/backends/vulkan/test/custom_ops/targets.bzl @@ -34,8 +34,7 @@ def define_common_targets(is_fbcode = False): runtime.filegroup( name = "custom_ops_shaders", srcs = native.glob([ - "glsl/*.glsl", - "glsl/*.yaml", + "glsl/**/*", ]), visibility = ["PUBLIC"], ) diff --git a/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py b/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py new file mode 100644 index 00000000000..5a5d78ee5b5 --- /dev/null +++ b/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py @@ -0,0 +1,262 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import stat +import tempfile +import unittest +from pathlib import Path + +from runtime import gen_vulkan_spv + + +class TestNestedShaderSources(unittest.TestCase): + def _write_fake_compiler( + self, directory: Path, assertions: str = "" + ) -> tuple[Path, Path]: + compiler_log = directory / "compiler.log" + compiler = directory / "fake_glslc.py" + compiler.write_text( + "#!/usr/bin/env python3\n" + "import pathlib\n" + "import sys\n" + "log = pathlib.Path(__file__).with_name('compiler.log')\n" + "with log.open('a') as log_file:\n" + " log_file.write('compile\\n')\n" + "source = pathlib.Path(sys.argv[2])\n" + + assertions + + "output = pathlib.Path(sys.argv[sys.argv.index('-o') + 1])\n" + "output.write_bytes(b'\\x03\\x02\\x23\\x07')\n" + ) + compiler.chmod(compiler.stat().st_mode | stat.S_IXUSR) + return compiler, compiler_log + + def _write_shader(self, directory: Path, variant_name: str) -> None: + directory.mkdir(parents=True) + (directory / "common.glslh").write_text("const uint value = 1;\n") + (directory / "kernel.glsl").write_text( + '#version 450\n#include "common.glslh"\nvoid main() {}\n' + ) + (directory / "kernel.yaml").write_text( + "kernel:\n" + " parameter_names_with_default_values: {}\n" + " shader_variants:\n" + f" - NAME: {variant_name}\n" + ) + + def test_repeated_template_names_in_nested_directories(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + source_root = root / "src" + self._write_shader(source_root / "direct", "direct_kernel") + self._write_shader(source_root / "winograd", "winograd_kernel") + output_dir = root / "out" + cache_dir = root / "cache" + output_dir.mkdir() + cache_dir.mkdir() + + generator = gen_vulkan_spv.SPVGenerator( + str(source_root), {}, glslc_path=None + ) + generator.generateSPV(output_dir, cache_dir, nthreads=1) + + self.assertTrue((output_dir / "direct" / "direct_kernel.glsl").is_file()) + self.assertTrue((output_dir / "direct" / "common.glslh").is_file()) + self.assertTrue( + (output_dir / "winograd" / "winograd_kernel.glsl").is_file() + ) + self.assertTrue((output_dir / "winograd" / "common.glslh").is_file()) + + def test_multi_dot_source_preserves_legacy_template_name(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + source_root = root / "src" + source_root.mkdir() + (source_root / "kernel.special.glsl").write_text( + "#version 450\nvoid main() {}\n" + ) + (source_root / "kernel.yaml").write_text( + "kernel:\n" + " parameter_names_with_default_values: {}\n" + " shader_variants:\n" + " - NAME: generated_kernel\n" + ) + output_dir = root / "out" + cache_dir = root / "cache" + output_dir.mkdir() + cache_dir.mkdir() + + generator = gen_vulkan_spv.SPVGenerator( + str(source_root), {}, glslc_path=None + ) + generator.generateSPV(output_dir, cache_dir, nthreads=1) + + self.assertTrue((output_dir / "generated_kernel.glsl").is_file()) + + def test_rejects_generated_shader_outside_output_directory(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + source_root = root / "src" + self._write_shader(source_root, str(root / "escaped_kernel")) + output_dir = root / "out" + cache_dir = root / "cache" + output_dir.mkdir() + cache_dir.mkdir() + compiler, _ = self._write_fake_compiler(root) + + generator = gen_vulkan_spv.SPVGenerator( + str(source_root), {}, glslc_path=str(compiler) + ) + with self.assertRaisesRegex(ValueError, "outside output directory"): + generator.generateSPV(output_dir, cache_dir, nthreads=1) + + def test_local_include_change_recompiles_nested_shader(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + source_root = root / "src" + shader_dir = source_root / "nested" + self._write_shader(shader_dir, "nested_kernel") + (source_root / "root_common.glslh").write_text( + "const uint root_value = 1;\n" + ) + (shader_dir / "kernel.glsl").write_text( + '#version 450\n#include "common.glslh"\n' + '#include "root_common.glslh"\nvoid main() {}\n' + ) + output_dir = root / "out" + cache_dir = root / "cache" + output_dir.mkdir() + cache_dir.mkdir() + + compiler, compiler_log = self._write_fake_compiler( + root, + "include_indices = [i for i, arg in enumerate(sys.argv) if arg == '-I']\n" + "local_include_dir = pathlib.Path(sys.argv[include_indices[0] + 1])\n" + "root_include_dir = pathlib.Path(sys.argv[include_indices[1] + 1])\n" + "assert local_include_dir == source.parent\n" + "assert (local_include_dir / 'common.glslh').is_file()\n" + "assert root_include_dir == source.parent.parent\n" + "assert (root_include_dir / 'root_common.glslh').is_file()\n", + ) + + def generate() -> None: + generator = gen_vulkan_spv.SPVGenerator( + str(source_root), {}, glslc_path=str(compiler) + ) + generator.generateSPV(output_dir, cache_dir, nthreads=1) + + generate() + generate() + self.assertEqual(compiler_log.read_text().splitlines(), ["compile"]) + + (shader_dir / "common.glslh").write_text("const uint value = 2;\n") + generate() + + self.assertEqual( + compiler_log.read_text().splitlines(), ["compile", "compile"] + ) + + (source_root / "root_common.glslh").write_text( + "const uint root_value = 2;\n" + ) + generate() + + self.assertEqual( + compiler_log.read_text().splitlines(), + ["compile", "compile", "compile"], + ) + + def test_ancestor_include_change_recompiles_nested_shader(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + source_root = root / "src" + family_dir = source_root / "family" + shader_dir = family_dir / "route" + self._write_shader(shader_dir, "route_kernel") + (family_dir / "family_common.glslh").write_text( + "const uint family_value = 1;\n" + ) + (shader_dir / "kernel.glsl").write_text( + '#version 450\n#include "family_common.glslh"\nvoid main() {}\n' + ) + output_dir = root / "out" + cache_dir = root / "cache" + output_dir.mkdir() + cache_dir.mkdir() + + compiler, compiler_log = self._write_fake_compiler( + root, + "include_dirs = [pathlib.Path(sys.argv[i + 1]) " + "for i, arg in enumerate(sys.argv) if arg == '-I']\n" + "assert include_dirs == [source.parent, source.parent.parent, " + "source.parent.parent.parent]\n", + ) + + def generate() -> None: + generator = gen_vulkan_spv.SPVGenerator( + str(source_root), {}, glslc_path=str(compiler) + ) + generator.generateSPV(output_dir, cache_dir, nthreads=1) + + generate() + generate() + self.assertEqual(compiler_log.read_text().splitlines(), ["compile"]) + + (family_dir / "family_common.glslh").write_text( + "const uint family_value = 2;\n" + ) + generate() + + self.assertEqual( + compiler_log.read_text().splitlines(), ["compile", "compile"] + ) + + def test_untracked_include_disables_cache(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + source_root = root / "src" + self._write_shader(source_root, "kernel") + (source_root / "kernel.glsl").write_text( + '#version 450\n#include "external.glslh"\nvoid main() {}\n' + ) + output_dir = root / "out" + cache_dir = root / "cache" + output_dir.mkdir() + cache_dir.mkdir() + compiler, compiler_log = self._write_fake_compiler(root) + + generator = gen_vulkan_spv.SPVGenerator( + str(source_root), {}, glslc_path=str(compiler) + ) + generator.generateSPV(output_dir, cache_dir, nthreads=1) + generator.generateSPV(output_dir, cache_dir, nthreads=1) + + self.assertEqual( + compiler_log.read_text().splitlines(), ["compile", "compile"] + ) + + def test_cyclic_includes_use_cache(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + source_root = root / "src" + self._write_shader(source_root, "kernel") + (source_root / "kernel.glsl").write_text( + '#version 450\n#include "a.glslh"\nvoid main() {}\n' + ) + (source_root / "a.glslh").write_text('#include "b.glslh"\n') + (source_root / "b.glslh").write_text('#include "a.glslh"\n') + output_dir = root / "out" + cache_dir = root / "cache" + output_dir.mkdir() + cache_dir.mkdir() + compiler, compiler_log = self._write_fake_compiler(root) + + generator = gen_vulkan_spv.SPVGenerator( + str(source_root), {}, glslc_path=str(compiler) + ) + generator.generateSPV(output_dir, cache_dir, nthreads=1) + generator.generateSPV(output_dir, cache_dir, nthreads=1) + + self.assertEqual(compiler_log.read_text().splitlines(), ["compile"]) diff --git a/backends/vulkan/test/targets.bzl b/backends/vulkan/test/targets.bzl index 0051a74a530..e3b735137d5 100644 --- a/backends/vulkan/test/targets.bzl +++ b/backends/vulkan/test/targets.bzl @@ -98,6 +98,16 @@ def define_common_targets(is_fbcode = False): ], ) + python_unittest( + name = "test_gen_vulkan_spv", + srcs = [ + "spv_codegen/test_gen_vulkan_spv.py", + ], + deps = [ + "//executorch/backends/vulkan:gen_vulkan_spv_lib", + ], + ) + python_unittest( name = "test_serialization", srcs = [ @@ -164,7 +174,7 @@ def define_common_targets(is_fbcode = False): fb_native.filegroup( name = "test_shaders", srcs = glob([ - "glsl/*", + "glsl/**/*", ]), visibility = [ "PUBLIC", diff --git a/backends/vulkan/tools/gpuinfo/targets.bzl b/backends/vulkan/tools/gpuinfo/targets.bzl index 2bc9b3f1d04..be0a9680609 100644 --- a/backends/vulkan/tools/gpuinfo/targets.bzl +++ b/backends/vulkan/tools/gpuinfo/targets.bzl @@ -8,7 +8,7 @@ def define_common_targets(is_fbcode = False): buck_filegroup( name = "gpuinfo_shaders", srcs = native.glob([ - "glsl/*", + "glsl/**/*", ]), visibility = [ "PUBLIC",