executorch 1.1.0__cp313-cp313-win_amd64.whl

executorch/backends/vulkan/runtime/gen_vulkan_spv.py

@@ -0,0 +1,1517 @@
+#!/usr/bin/env python3
+# 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.
+
+# pyre-unsafe
+
+import argparse
+import array
+import codecs
+import copy
+import glob
+import hashlib
+import io
+import os
+import re
+import shutil
+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
+import textwrap
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+
+import yaml
+from yaml.constructor import ConstructorError
+from yaml.nodes import MappingNode
+
+try:
+    from yaml import CLoader as Loader
+except ImportError:
+    from yaml import Loader  # type: ignore[assignment, misc]
+
+CPP_H_NAME = "spv.h"
+CPP_SRC_NAME = "spv.cpp"
+
+# Basic configuration settings for shaders
+DEFAULT_ENV: Dict[str, Any] = {
+    "PRECISION": "highp",
+    # B is shorthand for "binding". This is used to automatically increment the
+    # layout binding index when declaring layout bindings. Note that a container
+    # type is used because integers are immutable in Python.
+    "B": [0],
+    # C is shorthand for "constant_id". This is used to automatically increment the
+    # constant_id index for specialization constants.
+    # Note that it starts at 3, as 0-2 are reserved for local workgroup size ids.
+    "C": [3],
+}
+
+# Establishes relationships between different tensor types and different GLSL types
+TYPE_MAPPINGS: Dict[str, Any] = {
+    "IMAGE_T": {
+        3: {
+            "double": "image3D",
+            "float": "image3D",
+            "half": "image3D",
+            # integer dtypes
+            "int8": "iimage3D",
+            "uint8": "uimage3D",
+            "int16": "iimage3D",
+            "uint16": "uimage3D",
+            "int32": "iimage3D",
+            "uint32": "uimage3D",
+            "int64": "iimage3D",
+            "uint64": "uimage3D",
+            # common dtype aliases
+            "bool": "uimage3D",
+            "int": "iimage3D",
+            "uint": "uimage3D",
+        },
+        2: {
+            "double": "image2D",
+            "float": "image2D",
+            "half": "image2D",
+            # integer dtypes
+            "int8": "iimage2D",
+            "uint8": "uimage2D",
+            "int16": "iimage2D",
+            "uint16": "uimage2D",
+            "int32": "iimage2D",
+            "uint32": "uimage2D",
+            "int64": "iimage2D",
+            "uint64": "uimage2D",
+            # common dtype aliases
+            "bool": "uimage2D",
+            "int": "iimage2D",
+            "uint": "uimage2D",
+        },
+    },
+    "SAMPLER_T": {
+        3: {
+            "double": "sampler3D",
+            "float": "sampler3D",
+            "half": "sampler3D",
+            # integer dtypes
+            "int8": "isampler3D",
+            "uint8": "usampler3D",
+            "int16": "isampler3D",
+            "uint16": "usampler3D",
+            "int32": "isampler3D",
+            "uint32": "usampler3D",
+            "int64": "isampler3D",
+            "uint64": "usampler3D",
+            # common dtype aliases
+            "bool": "usampler3D",
+            "int": "isampler3D",
+            "uint": "usampler3D",
+        },
+        2: {
+            "double": "sampler2D",
+            "float": "sampler2D",
+            "half": "sampler2D",
+            # integer dtypes
+            "int8": "isampler2D",
+            "uint8": "usampler2D",
+            "int16": "isampler2D",
+            "uint16": "usampler2D",
+            "int32": "isampler2D",
+            "uint32": "usampler2D",
+            "int64": "isampler2D",
+            "uint64": "usampler2D",
+            # common dtype aliases
+            "bool": "usampler2D",
+            "int": "isampler2D",
+            "uint": "usampler2D",
+        },
+    },
+    "IMAGE_FORMAT": {
+        "double": "rgba32f",
+        "float": "rgba32f",
+        "half": "rgba16f",
+        # integer dtypes
+        "int8": "rgba8i",
+        "uint8": "rgba8ui",
+        "int16": "rgba16i",
+        "uint16": "rgba16ui",
+        "int32": "rgba32i",
+        "uint32": "rgba32ui",
+        "int64": "rgba32i",
+        "uint64": "rgba32ui",
+        # common dtype aliases
+        "bool": "rgba8ui",
+        "int": "rgba32i",
+        "uint": "rgba32ui",
+    },
+}
+
+
+def define_variable(name: str) -> str:
+    if name in locals():
+        return f"#define {name} {locals()[name]}"
+    elif name in globals():
+        return f"#define {name} {globals()[name]}"
+    else:
+        raise RuntimeError(f"{name} is not defined")
+
+
+def buffer_scalar_type(dtype: str) -> str:
+    if dtype == "half":
+        return "float16_t"
+    elif dtype == "float":
+        return "float"
+    elif dtype == "double":
+        return "float64_t"
+    # integer dtype alias conversion
+    elif dtype == "bool":
+        return "uint8_t"
+    # we don't want to append _t for int32 or uint32 as int is already 32bit
+    elif dtype == "int32" or dtype == "uint32":
+        return "int" if dtype == "int32" else "uint"
+    elif dtype[-1].isdigit():
+        return dtype + "_t"
+    return dtype
+
+
+def buffer_gvec_type(dtype: str, n: int) -> str:
+    if n == 1:
+        return buffer_scalar_type(dtype)
+
+    dtype_map = {
+        "half": f"f16vec{n}",
+        "float": f"vec{n}",
+        "double": f"vec{n}",  # No 64bit image format support in GLSL
+        "int8": f"i8vec{n}",
+        "uint8": f"u8vec{n}",
+        "int16": f"i16vec{n}",
+        "uint16": f"u16vec{n}",
+        "int32": f"ivec{n}",
+        "int": f"ivec{n}",
+        "uint32": f"uvec{n}",
+        "uint": f"uvec{n}",
+        "int64": f"ivec{n}",  # No 64bit image format support in GLSL
+        "uint64": f"uvec{n}",  # No 64bit image format support in GLSL
+        "bool": f"u8vec{n}",
+    }
+
+    vector_type = dtype_map.get(dtype)
+    if vector_type is None:
+        raise AssertionError(f"Invalid dtype: {dtype}")
+
+    return vector_type
+
+
+def texel_type(dtype: str) -> str:
+    image_format = TYPE_MAPPINGS["IMAGE_FORMAT"][dtype]
+    if image_format[-1:] == "f":
+        return "vec4"
+    elif image_format[-2:] == "ui":
+        return "uvec4"
+    elif image_format[-1:] == "i":
+        return "ivec4"
+    raise AssertionError(f"Invalid image format: {image_format}")
+
+
+def gvec_type(dtype: str, n: int) -> str:
+    gvec4_type = texel_type(dtype)
+    return gvec4_type[:-1] + str(n)
+
+
+def texel_component_type(dtype: str) -> str:
+    vec4_type = texel_type(dtype)
+    if vec4_type[:3] == "vec":
+        return "float"
+    elif vec4_type[:4] == "ivec":
+        return "int"
+    elif vec4_type[:4] == "uvec":
+        return "uint"
+    raise AssertionError(f"Invalid vec4 type: {vec4_type}")
+
+
+def accum_vec_type(dtype: str) -> str:
+    return texel_type(dtype)
+
+
+def accum_scalar_type(dtype: str) -> str:
+    return texel_component_type(dtype)
+
+
+def texel_load_type(dtype: str, storage_type: str) -> str:
+    if storage_type.lower() == "buffer":
+        return buffer_gvec_type(dtype, 4)
+    else:
+        return texel_type(dtype)
+
+
+def texel_load_component_type(dtype: str, storage_type: str) -> str:
+    if storage_type.lower() == "buffer":
+        return buffer_scalar_type(dtype)
+    else:
+        return texel_component_type(dtype)
+
+
+def get_access_qualifier(access_type: Optional[str]) -> str:
+    if access_type is None:
+        return ""
+    if access_type.lower() == "r":
+        return "readonly"
+    if access_type.lower() == "w":
+        return "writeonly"
+    if access_type.lower() == "rw":
+        return ""
+
+    raise AssertionError(f"Invalid access type: {access_type}")
+
+
+def get_slot_val(slot: Union[int, List[int]]) -> int:
+    if isinstance(slot, list):
+        return slot[0]
+    return slot
+
+
+def layout_declare_buffer(
+    slot: Union[int, List[int]],
+    access_type: str,
+    var_name: str,
+    dtype: str,
+    precision: str = "PRECISION",
+    is_scalar_array: bool = True,
+) -> str:
+    array_type = buffer_gvec_type(dtype, 4)
+    if is_scalar_array:
+        array_type = buffer_scalar_type(dtype)
+
+    out_str = f"""
+layout(set = 0, binding = {get_slot_val(slot)}) buffer {precision} restrict {get_access_qualifier(access_type)} {var_name}Buffer {{
+    {array_type} {var_name}[];
+}};
+"""
+
+    if isinstance(slot, list):
+        slot[0] = slot[0] + 1
+    return out_str
+
+
+def layout_declare_image(
+    slot: Union[int, List[int]],
+    access_type: str,
+    var_name: str,
+    dtype: str,
+    precision: str = "PRECISION",
+    image_ndim: int = 3,
+) -> str:
+    image_format = TYPE_MAPPINGS["IMAGE_FORMAT"][dtype]
+    image_type = TYPE_MAPPINGS["IMAGE_T"][image_ndim][dtype]
+
+    ret_str = f"layout(set = 0, binding = {get_slot_val(slot)}, {image_format}) uniform {precision} restrict {get_access_qualifier(access_type)} {image_type} {var_name};"
+
+    if isinstance(slot, list):
+        slot[0] = slot[0] + 1
+    return ret_str
+
+
+def layout_declare_sampler(
+    slot: Union[int, List[int]],
+    access_type: str,
+    var_name: str,
+    dtype: str,
+    precision: str = "PRECISION",
+    access_qualifier: Optional[str] = None,
+    image_ndim: int = 3,
+) -> str:
+    sampler_type = TYPE_MAPPINGS["SAMPLER_T"][image_ndim][dtype]
+
+    ret_str = f"layout(set = 0, binding = {get_slot_val(slot)}) uniform {precision} {sampler_type} {var_name};"
+
+    if isinstance(slot, list):
+        slot[0] = slot[0] + 1
+    return ret_str
+
+
+def layout_declare_tensor(
+    slot: Union[int, List[int]],
+    access_type: str,
+    var_name: str,
+    dtype: str,
+    storage_type: str,
+    is_scalar_array: bool = True,
+    precision: str = "PRECISION",
+) -> str:
+    assert storage_type.lower() in ["buffer", "texture3d", "texture2d"]
+
+    image_ndim = 3
+    if storage_type.lower() == "texture2d":
+        image_ndim = 2
+
+    # Create buffer binding
+    if storage_type.lower() == "buffer":
+        return layout_declare_buffer(
+            slot,
+            access_type,
+            var_name,
+            dtype,
+            precision,
+            is_scalar_array=is_scalar_array,
+        )
+
+    # Create image/sampler binding
+    if access_type.lower() == "r":
+        return layout_declare_sampler(
+            slot, access_type, var_name, dtype, precision, image_ndim=image_ndim
+        )
+    else:
+        return layout_declare_image(
+            slot, access_type, var_name, dtype, precision, image_ndim=image_ndim
+        )
+
+
+def layout_declare_ubo(
+    slot: Union[int, List[int]], *args, precision: str = "PRECISION"
+) -> str:
+    assert len(args) % 2 == 0
+
+    var_list = list(zip(args[::2], args[1::2]))
+
+    ubo_name = ""
+    for _, var_name in var_list:
+        ubo_name += var_name + "_"
+
+    out_str = f"""
+layout(set = 0, binding = {get_slot_val(slot)}) uniform {precision} restrict readonly {ubo_name}UBO {{
+"""
+    for type_name, var_name in var_list:
+        out_str += f"  {type_name} {var_name};\n"
+    out_str += "};"
+
+    if isinstance(slot, list):
+        slot[0] = slot[0] + 1
+    return out_str
+
+
+def layout_declare_spec_const(
+    slot: Union[int, List[int]],
+    type_name: str,
+    var_name: str,
+    initial_val: Optional[str] = None,
+) -> str:
+    assert type_name in ["int", "uint", "float", "bool"]
+
+    out_str = f"layout(constant_id = {get_slot_val(slot)}) const {type_name} {var_name}"
+    if initial_val is not None:
+        out_str += f" = {initial_val}"
+    out_str += ";"
+
+    if isinstance(slot, list):
+        slot[0] = slot[0] + 1
+    return out_str
+
+
+def define_active_storage_type(storage_type: str):
+    if storage_type.lower() == "buffer":
+        return "#define USING_BUFFER"
+    elif storage_type.lower() == "texture3d":
+        return "#define USING_TEXTURE3D"
+    elif storage_type.lower() == "texture2d":
+        return "#define USING_TEXTURE2D"
+    else:
+        raise AssertionError(f"Invalid storage type: {storage_type}")
+
+
+def define_required_extensions(dtypes: Union[str, List[str]]):
+    out_str = "\n"
+    dtype_list = dtypes if isinstance(dtypes, list) else [dtypes]
+
+    for dtype in dtype_list:
+        nbit = None
+        glsl_type = None
+        if dtype == "half":
+            nbit = "16bit"
+            glsl_type = "float16"
+        elif dtype == "double":
+            # We only need to allow float64_t type usage
+            glsl_type = "float64"
+        elif dtype in ["int8", "uint8", "bool"]:
+            nbit = "8bit"
+            glsl_type = "int8"
+        elif dtype in ["int16", "uint16"]:
+            nbit = "16bit"
+            glsl_type = "int16"
+        elif dtype in ["int64", "uint64"]:
+            # We only need to allow int64_t and uint64_t type usage
+            glsl_type = "int64"
+
+        if nbit is not None:
+            out_str += f"#extension GL_EXT_shader_{nbit}_storage : require\n"
+        if glsl_type is not None:
+            out_str += f"#extension GL_EXT_shader_explicit_arithmetic_types_{glsl_type} : require\n"
+
+    return out_str
+
+
+UTILITY_FNS: Dict[str, Any] = {
+    "macro_define": define_variable,
+    "get_pos": {
+        3: lambda pos: pos,
+        2: lambda pos: f"{pos}.xy",
+    },
+    "buffer_scalar_type": buffer_scalar_type,
+    "buffer_gvec_type": buffer_gvec_type,
+    "texel_type": texel_type,
+    "gvec_type": gvec_type,
+    "accum_vec_type": accum_vec_type,
+    "accum_scalar_type": accum_scalar_type,
+    "texel_component_type": texel_component_type,
+    "texel_load_type": texel_load_type,
+    "texel_load_component_type": texel_load_component_type,
+    "layout_declare_buffer": layout_declare_buffer,
+    "layout_declare_image": layout_declare_image,
+    "layout_declare_sampler": layout_declare_sampler,
+    "layout_declare_tensor": layout_declare_tensor,
+    "layout_declare_ubo": layout_declare_ubo,
+    "layout_declare_spec_const": layout_declare_spec_const,
+    "define_active_storage_type": define_active_storage_type,
+    "define_required_extensions": define_required_extensions,
+}
+
+
+def extract_filename(path: str, keep_ext: bool = True) -> Any:
+    if keep_ext:
+        return os.path.basename(path)
+    else:
+        return os.path.basename(path).split(".")[0]
+
+
+def extract_extension(path: str) -> str:
+    return os.path.splitext(extract_filename(path))[1][1:]
+
+
+############################
+#  SPIR-V Code Generation  #
+############################
+
+
+# https://gist.github.com/pypt/94d747fe5180851196eb
+class UniqueKeyLoader(Loader):
+    def construct_mapping(self, node, deep=False):  # type: ignore[no-untyped-def]
+        if not isinstance(node, MappingNode):
+            raise ConstructorError(
+                None,
+                None,
+                f"expected a mapping node, but found {node.id}",
+                node.start_mark,
+            )
+        mapping = {}
+        for key_node, value_node in node.value:
+            key = self.construct_object(key_node, deep=deep)  # type: ignore[no-untyped-call]
+            try:
+                hash(key)
+            except TypeError as e:
+                raise ConstructorError(
+                    "while constructing a mapping",
+                    node.start_mark,
+                    "found unacceptable key ",
+                    key_node.start_mark,
+                ) from e
+            # check for duplicate keys
+            if key in mapping:
+                raise ConstructorError(
+                    "while constructing a mapping",
+                    node.start_mark,
+                    "found duplicate key",
+                    key_node.start_mark,
+                )
+            value = self.construct_object(value_node, deep=deep)  # type: ignore[no-untyped-call]
+            mapping[key] = value
+        return mapping
+
+
+# https://github.com/google/XNNPACK/blob/master/tools/xngen.py
+def extract_leading_whitespace(line: str) -> str:
+    match = re.match(r"\s*", line)
+    return match.group(0) if match else ""
+
+
+# https://github.com/google/XNNPACK/blob/master/tools/xngen.py
+def escape(line: str) -> str:
+    output_parts = []
+    while "${" in line:
+        start_pos = line.index("${")
+        end_pos = line.index("}", start_pos + 2)
+        if start_pos != 0:
+            output_parts.append('"' + line[:start_pos].replace('"', '\\"') + '"')
+        output_parts.append("str(" + line[start_pos + 2 : end_pos] + ")")
+        line = line[end_pos + 1 :]
+    if line:
+        output_parts.append('"' + line.replace('"', '\\"') + '"')
+    return " + ".join(output_parts)
+
+
+# https://github.com/google/XNNPACK/blob/master/tools/xngen.py
+def preprocess(
+    input_text: str, variables: Dict[str, Any], input_path: str = "codegen"
+) -> str:
+    # Workaround to handle source files using \ to extend mecros to a new line
+    input_text = re.sub(r"\\$", r"\\\\", input_text, flags=re.MULTILINE)
+
+    input_lines = input_text.splitlines()
+    python_lines = []
+
+    blank_lines = 0
+
+    last_indent = ""
+
+    # List of tuples (total_index, python_indent)
+    indent_stack = [("", "")]
+
+    # Indicates whether this is the first line inside Python
+    # code block (i.e. for, while, if, elif, else)
+    python_block_start = True
+    for input_line in input_lines:
+        if input_line == "":
+            blank_lines += 1
+            continue
+        # Skip lint markers.
+        if "LINT" in input_line:
+            continue
+
+        input_indent = extract_leading_whitespace(input_line)
+        if python_block_start:
+            assert input_indent.startswith(last_indent)
+            extra_python_indent = input_indent[len(last_indent) :]
+            python_indent = indent_stack[-1][1] + extra_python_indent
+            indent_stack.append((input_indent, python_indent))
+            assert input_indent.startswith(indent_stack[-1][0])
+        else:
+            while not input_indent.startswith(indent_stack[-1][0]):
+                del indent_stack[-1]
+        python_block_start = False
+
+        python_indent = indent_stack[-1][1]
+        stripped_input_line = input_line.strip()
+        if stripped_input_line.startswith("$") and not stripped_input_line.startswith(
+            "${"
+        ):
+            if stripped_input_line.endswith(":"):
+                python_block_start = True
+            while blank_lines != 0:
+                python_lines.append(python_indent + "print(file=OUT_STREAM)")
+                blank_lines -= 1
+            python_lines.append(python_indent + stripped_input_line.replace("$", ""))
+        else:
+            assert input_line.startswith(python_indent)
+            while blank_lines != 0:
+                python_lines.append(python_indent + "print(file=OUT_STREAM)")
+                blank_lines -= 1
+            python_lines.append(
+                python_indent
+                + "print(%s, file=OUT_STREAM)"
+                % escape(input_line[len(python_indent) :])
+            )
+        last_indent = input_indent
+
+    while blank_lines != 0:
+        python_lines.append(python_indent + "print(file=OUT_STREAM)")
+        blank_lines -= 1
+
+    exec_globals = dict(variables)
+    output_stream = io.StringIO()
+    exec_globals["OUT_STREAM"] = output_stream
+
+    python_bytecode = compile("\n".join(python_lines), input_path, "exec")
+    exec(python_bytecode, exec_globals)
+
+    return output_stream.getvalue()
+
+
+class SPVGenerator:
+    def __init__(
+        self,
+        src_dir_paths: Union[str, List[str]],
+        env: Dict[Any, Any],
+        glslc_path: Optional[str],
+        glslc_flags: str = "",
+        replace_u16vecn: bool = False,
+    ) -> None:
+        if isinstance(src_dir_paths, str):
+            self.src_dir_paths = [src_dir_paths]
+        else:
+            self.src_dir_paths = src_dir_paths
+
+        self.env = env
+        self.glslc_path = glslc_path
+        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.src_files: Dict[str, str] = {}
+        self.template_yaml_files: List[str] = []
+
+        self.addSrcAndYamlFiles(self.src_dir_paths)
+        self.shader_template_params: Dict[Any, Any] = {}
+        for yaml_file in self.template_yaml_files:
+            self.parseTemplateYaml(yaml_file)
+
+        self.output_file_map: Dict[str, Tuple[str, Dict[str, str]]] = {}
+        self.constructOutputMap()
+
+    def addSrcAndYamlFiles(self, src_dir_paths: List[str]) -> None:
+        for src_path in src_dir_paths:
+            # 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:
+                if len(file) > 1:
+                    self.src_files[extract_filename(file, keep_ext=False)] = file
+            # Collect template yaml files
+            yaml_files = glob.glob(
+                os.path.join(src_path, "**", "*.yaml"), recursive=True
+            )
+            for file in yaml_files:
+                if len(file) > 1:
+                    self.template_yaml_files.append(file)
+
+    def generateVariantCombinations(  # noqa: C901
+        self,
+        iterated_params: Dict[str, Any],
+        exclude_params: Optional[Set[str]] = None,
+    ) -> List[Any]:
+        if exclude_params is None:
+            exclude_params = set()
+        all_iterated_params = []
+        for param_name, value_list in iterated_params.items():
+            if re.match(r"^combination\d*$", param_name):
+                param_values = []
+                param_names = value_list["parameter_names"]
+                combos = value_list["combos"]
+                for combo in combos:
+                    parameter_values = combo["parameter_values"]
+                    if "suffix" in combo:
+                        suffix = combo["suffix"]
+                    else:
+                        suffix = ""
+                        for param_value in parameter_values:
+                            if len(str(param_value)) > 0:
+                                suffix += "_" + str(param_value)
+                        suffix = suffix[1:]
+                    param_values.append((param_names, suffix, parameter_values))
+
+                all_iterated_params.append(param_values)
+
+            elif param_name not in exclude_params:
+                param_values = []
+                for value in value_list:
+                    if "RANGE" in value:
+                        value_range = value["RANGE"]
+                        suffix = value.get("SUFFIX", "")
+                        if isinstance(value_range, list) and len(value_range) == 2:
+                            for i in range(value_range[0], value_range[1] + 1):
+                                curr_suffix = (
+                                    suffix + "_" + str(i) if suffix else str(i)
+                                )
+                                param_values.append((param_name, curr_suffix, i))
+                        else:
+                            raise ValueError(
+                                f"{value['RANGE']} is not a valid range. Must be in format [start, end] (inclusive)."
+                            )
+
+                    elif "VALUE" in value:
+                        suffix = value.get("SUFFIX", value["VALUE"])
+                        if value["VALUE"] in ["int", "uint"]:
+                            raise ValueError(
+                                f"Use int32 or uint32 instead of {value['VALUE']}"
+                            )
+                        param_values.append((param_name, suffix, value["VALUE"]))
+
+                    else:
+                        raise KeyError(
+                            "Parameter must be 'VALUE: string' or 'RANGE: [a, b]'"
+                        )
+
+                all_iterated_params.append(param_values)
+
+        return list(product(*all_iterated_params))
+
+    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")
+
+                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] = []
+
+                default_iterated_params = params_dict.get(
+                    "generate_variant_forall", None
+                )
+
+                for variant in params_dict["shader_variants"]:
+                    default_iterated_params_names = set(
+                        default_iterated_params.keys()
+                        if default_iterated_params is not None
+                        else {}
+                    )
+                    variant_params_names = set(variant.keys())
+
+                    invalid_keys = (
+                        variant_params_names
+                        - default_iterated_params_names
+                        - params_names
+                        - {"generate_variant_forall"}
+                    )
+                    assert len(invalid_keys) == 0
+
+                    iterated_params = variant.get(
+                        "generate_variant_forall", default_iterated_params
+                    )
+
+                    if iterated_params is not None:
+                        variant_combinations = self.generateVariantCombinations(
+                            iterated_params, variant_params_names
+                        )
+
+                        for combination in variant_combinations:
+                            default_params_copy = copy.deepcopy(default_params)
+                            for key in variant:
+                                if key != "generate_variant_forall":
+                                    default_params_copy[key] = variant[key]
+
+                            variant_name = variant["NAME"]
+
+                            for setting in combination:
+                                param_names = setting[0]
+                                suffix = setting[1]
+                                param_values = setting[2]
+                                if isinstance(param_names, list):
+                                    for param_name, param_value in zip(
+                                        param_names, param_values
+                                    ):
+                                        default_params_copy[param_name] = param_value
+                                else:
+                                    default_params_copy[param_names] = param_values
+
+                                if len(str(suffix)) > 0:
+                                    variant_name = f"{variant_name}_{suffix}"
+
+                            default_params_copy["NAME"] = variant_name
+                            default_params_copy["VARIANT_NAME"] = variant["NAME"]
+
+                            self.shader_template_params[template_name].append(
+                                default_params_copy
+                            )
+                    else:
+                        default_params_copy = copy.deepcopy(default_params)
+                        for key in variant:
+                            default_params_copy[key] = variant[key]
+
+                        self.shader_template_params[template_name].append(
+                            default_params_copy
+                        )
+
+    def create_shader_params(
+        self, variant_params: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, str]:
+        if variant_params is None:
+            variant_params = {}
+        shader_params = copy.deepcopy(self.env)
+        for key, value in variant_params.items():
+            shader_params[key] = value
+
+        return shader_params
+
+    def constructOutputMap(self) -> None:
+        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"]] = (
+                    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] = (
+                    src_file_fullpath,
+                    self.create_shader_params(),
+                )
+
+    def maybe_replace_u16vecn(self, input_text: str) -> str:
+        """
+        There is a latency benefit to using u16vecn variables to store texture position
+        variables instead of ivecn, likely due to reduced register pressure. However,
+        SwiftShader does not support 16 bit integer types in shaders, so this is a crude
+        way to fallback to using ivecn to store texture positions so that testing with
+        SwiftShader is still possible.
+        """
+        if not self.replace_u16vecn:
+            return input_text
+        if "codegen-nosub" in input_text:
+            return input_text
+
+        # Remove extension requirement so that generated ShaderInfo does not mark it
+        input_text = input_text.replace(
+            "#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require", ""
+        )
+        input_text = input_text.replace("u16vec", "ivec")
+        input_text = input_text.replace("uint16_t", "int")
+        return input_text
+
+    def get_md5_checksum(self, file_path: str) -> str:
+        # Use a reasonably sized buffer for better performance with large files
+        BUF_SIZE = 65536  # 64kb chunks
+
+        md5 = hashlib.md5()
+
+        with open(file_path, "rb") as f:
+            while True:
+                data = f.read(BUF_SIZE)
+                if not data:
+                    break
+                md5.update(data)
+
+        # Get the hexadecimal digest and compare
+        file_md5 = md5.hexdigest()
+        return file_md5
+
+    def generateSPV(  # noqa: C901
+        self,
+        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:
+        #
+        # 1. A bool indicationg if the file has changed since the last compilation; this
+        #    is determined by comparing against the cached version.
+        # 2. List of other source files included by the generated file.
+        gen_file_meta: Dict[str, Tuple[bool, List[str], str]] = {}
+
+        # Return value of the function mapping the abspath of compiled SPIR-V binaries
+        # to the abspath of the generated GLSL file they were compiled from.
+        spv_to_glsl_map: Dict[str, str] = {}
+
+        # Convert output_dir to absolute path
+        assert os.path.exists(output_dir)
+        output_dir = os.path.abspath(output_dir)
+
+        if cache_dir is not None:
+            assert os.path.exists(cache_dir)
+
+        def get_glsl_includes(glsl_text):
+            """
+            Parse GLSL text content and return a list of included files.
+
+            Args:
+                glsl_text: String containing the GLSL file content to analyze
+
+            Returns:
+                List of included file names (e.g., ["random.h"])
+            """
+            includes = []
+            for line in glsl_text.splitlines():
+                # Look for #include directives with quoted filenames
+                # Matches: #include "filename.h" or #include <filename.h>
+                include_match = re.match(
+                    r'^\s*#include\s+[<"]([^>"]+)[>"]', line.strip()
+                )
+                if include_match:
+                    includes.append(include_match.group(1))
+
+            return includes
+
+        def file_has_changed(gen_file_path, cached_file_path):
+            # If the file does not exist in the cache, then return True
+            if not os.path.exists(cached_file_path):
+                return True
+            current_checksum = self.get_md5_checksum(gen_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):
+            """
+            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.
+            """
+            gen_file_changed, includes_list = gen_file_meta[gen_file_path]
+            any_changed = gen_file_changed
+            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
+                )
+
+            return any_changed
+
+        def generate_src_file(shader_paths_pair) -> Tuple[bool, List[str]]:
+            """
+            Given an input tuple containing the following items:
+            (src_file_name, (template_file_path, codegen_params))
+
+            This function generates src_file_name by processing
+            template_file_path with the Python preprocessor using the
+            parameters specified by codegen_params.
+
+            Then, it returns a tuple containing:
+            1. The path of the generated source file
+            2. A bool indicating if the generated source file has changed since the last
+               compilation.
+            3. A list of files included by the generated source file
+            """
+            # name of .glsl, .glslh, or .h file to be generated
+            src_file_name = shader_paths_pair[0]
+            # path of template file used for codegen
+            template_file_path = shader_paths_pair[1][0]
+            # args to be used for codegen
+            codegen_params = shader_paths_pair[1][1]
+
+            # Assume that generated files will have the same file extension as the
+            # source template file.
+            out_file_ext = extract_extension(template_file_path)
+
+            # Construct generated file name
+            gen_out_path = os.path.join(output_dir, f"{src_file_name}.{out_file_ext}")
+            # Construct path of cached generated file
+            cached_gen_out_path = os.path.join(
+                cache_dir, f"{src_file_name}.{out_file_ext}"
+            )
+
+            # Execute codegen to generate the output file
+            with codecs.open(template_file_path, "r", encoding="utf-8") as input_file:
+                input_text = input_file.read()
+                input_text = self.maybe_replace_u16vecn(input_text)
+                output_text = preprocess(input_text, codegen_params)
+
+            included_files = get_glsl_includes(output_text)
+
+            with codecs.open(gen_out_path, "w", encoding="utf-8") as output_file:
+                output_file.write(output_text)
+
+            file_changed = (
+                file_has_changed(gen_out_path, cached_gen_out_path) or force_rebuild
+            )
+
+            # Save the generated file to cache so it can be used for future checks
+            if cache_dir is not None and file_changed:
+                shutil.copyfile(gen_out_path, cached_gen_out_path)
+
+            return gen_out_path, file_changed, included_files
+
+        def compile_spirv(shader_paths_pair) -> Tuple[str, str]:
+            """
+            Given an input tuple containing the following items:
+            (src_file_name, (template_file_path, codegen_params))
+
+            Infer the path of the GLSL source file generated by generate_src_file and
+            compile a SPIR-V binary from it. Returns the path of the compiled SPIR-V
+            binary and the path of the source file used to compile it.
+
+            This function also utilizes a caching mechanism; if generate_src_file
+            reported that the source file was unchanged since the last successful
+            compilation, AND if the SPIR-V from the last successful compilation was
+            stored in the cache, then directly use the cached SPIR-V without triggering
+            a re-compilation.
+            """
+            # name of generated .glsl, .glslh, or .h from generate_src_file
+            src_file_name = shader_paths_pair[0]
+            # path of template file used for codegen
+            template_file_path = shader_paths_pair[1][0]
+            # args used for codegen
+            codegen_params = shader_paths_pair[1][1]
+
+            # Assume that generated files will have the same file extension as the
+            # source template file.
+            out_file_ext = extract_extension(template_file_path)
+
+            # Infer name of generated file (created by generate_src_file)
+            gen_out_path = os.path.join(output_dir, f"{src_file_name}.{out_file_ext}")
+
+            # Only proceed if GLSL -> SPIR-V compilation is required for this file
+            if out_file_ext != "glsl":
+                return (None, gen_out_path)
+
+            # Validate that the source file actually exists
+            assert os.path.exists(gen_out_path) and gen_out_path in gen_file_meta
+
+            # Construct name of SPIR-V file to be compiled
+            spv_out_path = os.path.join(output_dir, f"{src_file_name}.spv")
+
+            if cache_dir is not None:
+                # Construct the file names of cached SPIR-V file to check if they exist
+                # 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)
+                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)
+
+            vk_version = codegen_params.get("VK_VERSION", "1.1")
+            # Only proceed if a GLSL compiler was specified
+            if self.glslc_path is not None:
+                cmd_base = [
+                    self.glslc_path,
+                    "-fshader-stage=compute",
+                    gen_out_path,
+                    "-o",
+                    spv_out_path,
+                    "--target-env=vulkan{}".format(vk_version),
+                    "-Werror",
+                    "-I",
+                    output_dir,
+                ]
+                cmd = cmd_base + self.glslc_flags
+
+                try:
+                    subprocess.run(cmd, check=True, capture_output=True, text=True)
+                except subprocess.CalledProcessError as e:
+                    opt_fail = "compilation succeeded but failed to optimize"
+                    err_msg_base = f"Failed to compile {os.getcwd()}/{gen_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:
+                            # Delete any existing cached SPIR-V file if it exists
+                            if os.path.exists(cached_spv_out_path):
+                                os.remove(cached_spv_out_path)
+
+                            raise RuntimeError(
+                                f"{err_msg_base} {e_no_opt.stderr}"
+                            ) from e_no_opt
+
+                    else:
+                        # Delete any existing cached SPIR-V file if it exists
+                        if os.path.exists(cached_spv_out_path):
+                            os.remove(cached_spv_out_path)
+
+                        raise RuntimeError(f"{err_msg_base} {e.stderr}") from e
+
+                # If compilation was successful, store the compiled SPIR-V file in the
+                # cache for future use.
+                if cache_dir is not None:
+                    shutil.copyfile(spv_out_path, cached_spv_out_path)
+
+            return (spv_out_path, gen_out_path)
+
+        # Run codegen serially to ensure that all .glsl, .glslh, and .h files are up to
+        # date before compilation
+        for generated_file_tuple in self.output_file_map.items():
+            gen_out_path, file_changed, include_list = generate_src_file(
+                generated_file_tuple
+            )
+            gen_file_meta[gen_out_path] = (file_changed, include_list)
+
+        # Parallelize SPIR-V compilation to optimize build time
+        # 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
+
+
+##############################################
+#  Shader Info and Shader Registry Handling  #
+##############################################
+
+
+@dataclass
+class ShaderInfo:
+    tile_size: List[int]
+    layouts: List[str]
+    weight_storage_type: str = ""
+    bias_storage_type: str = ""
+    register_for: Optional[Tuple[str, List[str]]] = None
+    requires_shader_int16_ext: bool = False
+    requires_16bit_storage_ext: bool = False
+    requires_8bit_storage_ext: bool = False
+    requires_integer_dot_product_ext: bool = False
+    requires_shader_int64_ext: bool = False
+    requires_shader_float64_ext: bool = False
+
+
+def getName(filePath: str) -> str:
+    return os.path.basename(filePath).replace("/", "_").replace(".", "_")
+
+
+def isDescriptorLine(lineStr: str) -> bool:
+    descriptorLineId = r"^layout\(set"
+    return re.search(descriptorLineId, lineStr) is not None
+
+
+def isTileSizeLine(lineStr: str) -> bool:
+    tile_size_id = r"^ \* TILE_SIZE = \("
+    return re.search(tile_size_id, lineStr) is not None
+
+
+def findTileSizes(lineStr: str) -> List[int]:
+    tile_size_id = r"^ \* TILE_SIZE = \(([0-9]+), ([0-9]+), ([0-9]+)\)"
+    matches = re.search(tile_size_id, lineStr)
+    if matches is None:
+        raise AssertionError("matches is None in findTileSizes")
+    return [int(matches.group(1)), int(matches.group(2)), int(matches.group(3))]
+
+
+def isWeightStorageTypeLine(lineStr: str) -> bool:
+    weight_storage_id = r"^ \* WEIGHT_STORAGE = "
+    return re.search(weight_storage_id, lineStr) is not None
+
+
+def getWeightStorageType(lineStr: str) -> str:
+    weight_storage_id = r"^ \* WEIGHT_STORAGE = ([a-zA-Z]+_\dD)"
+    matches = re.search(weight_storage_id, lineStr)
+    if matches is None:
+        raise AssertionError("matches is None in getWeightStorageType")
+    return matches.group(1)
+
+
+def isBiasStorageTypeLine(lineStr: str) -> bool:
+    weight_storage_id = r"^ \* BIAS_STORAGE = "
+    return re.search(weight_storage_id, lineStr) is not None
+
+
+def getBiasStorageType(lineStr: str) -> str:
+    weight_storage_id = r"^ \* BIAS_STORAGE = ([a-zA-Z]+_\dD)"
+    matches = re.search(weight_storage_id, lineStr)
+    if matches is None:
+        raise AssertionError("matches is None in getBiasStorageType")
+    return matches.group(1)
+
+
+def isRegisterForLine(lineStr: str) -> bool:
+    # Check for Shader Name and a list of at least one Registry Key
+    register_for_id = (
+        r"^ \* REGISTER_FOR = \('([A-Za-z0-9_]+)'\s*,\s*\['([A-Za-z0-9_]+)'.*\]\)"
+    )
+    return re.search(register_for_id, lineStr) is not None
+
+
+def findRegisterFor(lineStr: str) -> Tuple[str, List[str]]:
+    register_for_pattern = r"'([A-Za-z0-9_]+)'"
+    matches = re.findall(register_for_pattern, lineStr)
+    if matches is None:
+        raise AssertionError("matches is None in getBiasStorageType")
+    matches_list = list(matches)
+    return (matches_list[0], matches_list[1:])
+
+
+def isExtensionRequireLine(lineStr: str) -> bool:
+    extension_require_id = r"^#extension ([A-Za-z0-9_]+)\s*:\s*require"
+    return re.search(extension_require_id, lineStr) is not None
+
+
+typeIdMapping = {
+    r"image[123]D\b": "VK_DESCRIPTOR_TYPE_STORAGE_IMAGE",
+    r"sampler[123]D\b": "VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER",
+    r"\bbuffer\b": "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER",
+    r"\buniform\b": "VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER",
+}
+
+
+def determineDescriptorType(lineStr: str) -> str:
+    for identifier, typeNum in typeIdMapping.items():
+        if re.search(identifier, lineStr):
+            return typeNum
+    raise AssertionError(
+        "No matching descriptor type for " + lineStr + " in determineDescriptorType"
+    )
+
+
+def getShaderInfo(srcFilePath: str) -> ShaderInfo:  # noqa: C901
+    shader_info = ShaderInfo([], [], "")
+    with open(srcFilePath) as srcFile:
+        for line in srcFile:
+            if isDescriptorLine(line):
+                shader_info.layouts.append(determineDescriptorType(line))
+            if isTileSizeLine(line):
+                shader_info.tile_size = findTileSizes(line)
+            if isWeightStorageTypeLine(line):
+                shader_info.weight_storage_type = getWeightStorageType(line)
+            if isBiasStorageTypeLine(line):
+                shader_info.bias_storage_type = getBiasStorageType(line)
+            if isRegisterForLine(line):
+                shader_info.register_for = findRegisterFor(line)
+            if isExtensionRequireLine(line):
+                if "GL_EXT_shader_explicit_arithmetic_types_int16" in line:
+                    shader_info.requires_shader_int16_ext = True
+                if "GL_EXT_shader_16bit_storage" in line:
+                    shader_info.requires_16bit_storage_ext = True
+                if "GL_EXT_shader_8bit_storage" in line:
+                    shader_info.requires_8bit_storage_ext = True
+                if "GL_EXT_integer_dot_product" in line:
+                    shader_info.requires_integer_dot_product_ext = True
+                if "GL_EXT_shader_explicit_arithmetic_types_int64" in line:
+                    shader_info.requires_shader_int64_ext = True
+                if "GL_EXT_shader_explicit_arithmetic_types_float64" in line:
+                    shader_info.requires_shader_float64_ext = True
+
+    return shader_info
+
+
+##########################
+#  C++ File Generation  #
+#########################
+
+cpp_template = """
+#include <executorch/backends/vulkan/runtime/api/ShaderRegistry.h>
+#include <stdint.h>
+#include <vector>
+
+using namespace vkcompute;
+
+namespace at {{
+namespace native {{
+namespace vulkan {{
+
+namespace {{
+
+{spv_bin_arrays}
+
+}}
+
+static void register_fn() {{
+
+{register_shader_infos}
+
+{shader_info_registry}
+
+}}
+
+static const api::ShaderRegisterInit register_shaders(&register_fn);
+
+}}
+}}
+}}
+
+"""
+
+
+def generateSpvBinStr(spvPath: str, name: str) -> Tuple[int, str]:
+    with open(spvPath, "rb") as fr:
+        next_bin = array.array("I", fr.read())
+        sizeBytes = 4 * len(next_bin)
+        spv_bin_str = "const uint32_t {}_bin[] = {{\n{}\n}};".format(
+            name,
+            textwrap.indent(",\n".join(str(x) for x in next_bin), "  "),
+        )
+
+    return sizeBytes, spv_bin_str
+
+
+def generateShaderInfoStr(shader_info: ShaderInfo, name: str, sizeBytes: int) -> str:
+    tile_size = (
+        f"{{{', '.join(str(x) for x in shader_info.tile_size)}}}"
+        if (len(shader_info.tile_size) > 0)
+        else "{1, 1, 1}"
+    )
+
+    shader_info_layouts = "{{{}}}".format(",\n ".join(shader_info.layouts))
+
+    def to_cpp_str(val: bool):
+        return "true" if val else "false"
+
+    shader_info_args = [
+        f'"{name}"',
+        f"{name}_bin",
+        str(sizeBytes),
+        shader_info_layouts,
+        tile_size,
+        to_cpp_str(shader_info.requires_shader_int16_ext),
+        to_cpp_str(shader_info.requires_16bit_storage_ext),
+        to_cpp_str(shader_info.requires_8bit_storage_ext),
+        to_cpp_str(shader_info.requires_integer_dot_product_ext),
+        to_cpp_str(shader_info.requires_shader_int64_ext),
+        to_cpp_str(shader_info.requires_shader_float64_ext),
+    ]
+
+    shader_info_str = textwrap.indent(
+        "api::shader_registry().register_shader(\n  vkapi::ShaderInfo(\n{args}));\n".format(
+            args=textwrap.indent(",\n".join(shader_info_args), "     "),
+        ),
+        "    ",
+    )
+
+    return shader_info_str
+
+
+def generateShaderDispatchStr(shader_info: ShaderInfo, name: str) -> str:
+    if shader_info.register_for is None:
+        return ""
+
+    (op_name, registry_keys) = shader_info.register_for
+    shader_dispatch_str = ""
+    for registry_key in registry_keys:
+        shader_dispatch_str = textwrap.indent(
+            f'api::shader_registry().register_op_dispatch("{op_name}", api::DispatchKey::{registry_key.upper()}, "{name}");',
+            "    ",
+        )
+
+    return shader_dispatch_str
+
+
+def genCppFiles(
+    spv_files: Dict[str, str], cpp_header_path: str, cpp_src_file_path: str
+) -> None:
+    spv_bin_strs = []
+    register_shader_info_strs = []
+    shader_registry_strs = []
+
+    for spvPath, srcPath in spv_files.items():
+        if spvPath is None:
+            continue
+
+        name = getName(spvPath).replace("_spv", "")
+
+        sizeBytes, spv_bin_str = generateSpvBinStr(spvPath, name)
+        spv_bin_strs.append(spv_bin_str)
+
+        shader_info = getShaderInfo(srcPath)
+
+        register_shader_info_strs.append(
+            generateShaderInfoStr(shader_info, name, sizeBytes)
+        )
+
+        if shader_info.register_for is not None:
+            shader_registry_strs.append(generateShaderDispatchStr(shader_info, name))
+
+    spv_bin_arrays = "\n".join(spv_bin_strs)
+    register_shader_infos = "\n".join(register_shader_info_strs)
+    shader_info_registry = "\n".join(shader_registry_strs)
+
+    cpp = cpp_template.format(
+        spv_bin_arrays=spv_bin_arrays,
+        register_shader_infos=register_shader_infos,
+        shader_info_registry=shader_info_registry,
+    )
+
+    with open(cpp_src_file_path, "w") as fw:
+        fw.write(cpp)
+
+
+##########
+#  Main  #
+##########
+
+
+def parse_arg_env(items: Dict[Any, Any]) -> Dict[Any, Any]:
+    d = {}
+    if items:
+        for item in items:
+            tokens = item.split("=")
+            key = tokens[0].strip()
+            value = tokens[1].strip()
+            d[key] = value
+    return d
+
+
+def main(argv: List[str]) -> int:
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument(
+        "-i",
+        "--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", required=True, help="")
+    parser.add_argument(
+        "-t", "--tmp-dir-path", required=True, help="/tmp/vulkan_shaders/"
+    )
+    parser.add_argument("-o", "--output-path", required=True, help="")
+    parser.add_argument("-f", "--force-rebuild", action="store_true", default=False)
+    parser.add_argument("--replace-u16vecn", action="store_true", default=False)
+    parser.add_argument("--optimize_size", action="store_true", help="")
+    parser.add_argument("--optimize", action="store_true", help="")
+    parser.add_argument("--spv_debug", action="store_true", default=False)
+    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
+    env.update(TYPE_MAPPINGS)
+    env.update(UTILITY_FNS)
+
+    for key, value in parse_arg_env(options.env).items():
+        env[key] = value
+
+    if not os.path.exists(options.output_path):
+        os.makedirs(options.output_path)
+
+    if not os.path.exists(options.tmp_dir_path):
+        os.makedirs(options.tmp_dir_path)
+
+    glslc_flags = []
+    if options.optimize_size:
+        glslc_flags.append("-Os")
+    elif options.optimize:
+        glslc_flags.append("-O")
+
+    if options.spv_debug:
+        glslc_flags.append("-g")
+
+    glslc_flags_str = " ".join(glslc_flags)
+
+    shader_generator = SPVGenerator(
+        options.glsl_paths,
+        env,
+        options.glslc_path,
+        glslc_flags=glslc_flags_str,
+        replace_u16vecn=options.replace_u16vecn,
+    )
+    output_spv_files = shader_generator.generateSPV(
+        options.output_path,
+        options.tmp_dir_path,
+        options.force_rebuild,
+        options.nthreads,
+    )
+
+    genCppFiles(
+        output_spv_files,
+        f"{options.output_path}/{CPP_H_NAME}",
+        f"{options.output_path}/{CPP_SRC_NAME}",
+    )
+
+    return 0
+
+
+def invoke_main() -> None:
+    sys.exit(main(sys.argv))
+
+
+if __name__ == "__main__":
+    invoke_main()  # pragma: no cover