[Vulkan] Enable including GLSL files from custom locations in gen_vulkan_spv (#91913)

@bypass-github-export-checks

To include custom locations when building with buck, use a ```-c gen_vulkan_spv.additional_glsl_paths="..."``` flag where ... is a list of filegroups and source directory paths separated by spaces,

ex. to include the sources added in D41413913, you would use

```
buck build ... -c gen_vulkan_spv.additional_glsl_paths="//xplat/caffe2:test_glsl_src_path_a test_src/a //xplat/caffe2:test_glsl_src_path_b test_src/b"
```

(as shown in the test plan)

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

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D41413914/)!
Pull Request resolved: https://github.com/pytorch/pytorch/pull/91913
Approved by: https://github.com/mcr229

buckbuild.bzl

@@ -110,6 +110,35 @@ def get_static_dispatch_backend():
         return []
     return static_dispatch_backend.split(";")
 
+def get_glsl_paths():
+    paths = [
+        "//xplat/caffe2:aten_vulkan_glsl_src_path",
+        "aten/src/ATen/native/vulkan/glsl",
+    ] + [
+        p
+        for p in read_config("gen_vulkan_spv", "additional_glsl_paths", "").split(" ")
+        if p
+    ]
+
+    if len(paths) % 2 != 0:
+        fail(
+            "gen_vulkan_spv.additional_glsl_paths must contain an even number of elements"
+        )
+
+    return " ".join(
+        [
+            "$(location {})/{}".format(
+                paths[i],
+                paths[i + 1],
+            )
+            for i in range(
+                0,
+                len(paths),
+                2,
+            )
+        ]
+    )
+
 # @lint-ignore BUCKRESTRICTEDSYNTAX
 IS_OSS = read_config("pt", "is_oss", "0") == "1"  # True for OSS BUCK build, and False for internal BUCK build
 

tools/gen_vulkan_spv.py

@@ -227,26 +227,29 @@ def genGLSLFromGLSLT(src_dir_path: str, tmp_dir_path: str) -> None:
 def genCppH(
     hFilePath: str,
     cppFilePath: str,
-    srcDirPath: str,
+    srcDirPaths: str,
     glslcPath: str,
     tmpDirPath: str,
     env: Dict[Any, Any],
 ) -> None:
     print(
-        "hFilePath:{} cppFilePath:{} srcDirPath:{} glslcPath:{} tmpDirPath:{}".format(
-            hFilePath, cppFilePath, srcDirPath, glslcPath, tmpDirPath
+        "hFilePath:{} cppFilePath:{} srcDirPaths:{} glslcPath:{} tmpDirPath:{}".format(
+            hFilePath, cppFilePath, srcDirPaths, glslcPath, tmpDirPath
         )
     )
 
-    vexs = glob.glob(os.path.join(srcDirPath, '**', '*.glsl'), recursive=True)
     templateSrcPaths = []
-    for f in vexs:
-        if len(f) > 1:
-            templateSrcPaths.append(f)
-            templateSrcPaths.sort()
 
-    # Now add glsl files that are generated from templates
-    genGLSLFromGLSLT(srcDirPath, tmpDirPath)
+    for srcDirPath in srcDirPaths:
+        vexs = glob.glob(os.path.join(srcDirPath, '**', '*.glsl'), recursive=True)
+        for f in vexs:
+            if len(f) > 1:
+                templateSrcPaths.append(f)
+                templateSrcPaths.sort()
+
+        # Now add glsl files that are generated from templates
+        genGLSLFromGLSLT(srcDirPath, tmpDirPath)
+
     vexs = glob.glob(os.path.join(tmpDirPath, '**', '*.glsl'), recursive=True)
     for f in vexs:
         if len(f) > 1:
@@ -273,9 +276,8 @@ def genCppH(
             glslcPath, "-fshader-stage=compute",
             srcPath, "-o", spvPath,
             "--target-env=vulkan1.0",
-            "-I", srcDirPath,
             "-Werror"
-        ]
+        ] + [arg for srcDirPath in srcDirPaths for arg in ["-I", srcDirPath]]
 
         print("\nglslc cmd:", cmd)
 
@@ -373,9 +375,11 @@ def main(argv: List[str]) -> int:
     parser = argparse.ArgumentParser(description='')
     parser.add_argument(
         '-i',
-        '--glsl-path',
-        help='',
-        default='.')
+        '--glsl-paths',
+        nargs='+',
+        help='List of paths to look for GLSL source files, separated by spaces. Ex: --glsl-paths "path1 path2 path3"',
+        default=['.'],
+    )
     parser.add_argument(
         '-c',
         '--glslc-path',
@@ -410,7 +414,7 @@ def main(argv: List[str]) -> int:
     genCppH(
         hFilePath=options.output_path + "/spv.h",
         cppFilePath=options.output_path + "/spv.cpp",
-        srcDirPath=options.glsl_path,
+        srcDirPaths=options.glsl_paths,
         glslcPath=options.glslc_path,
         tmpDirPath=options.tmp_dir_path,
         env=env)