gen_vulkan_spv.py

#!/usr/bin/env python3

import argparse
import array
import copy
import glob
import os
import re
import sys

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import subprocess
import textwrap
from collections import OrderedDict
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

import yaml
from torchgen.code_template import CodeTemplate
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]

H_NAME = "spv.h"
CPP_NAME = "spv.cpp"
DEFAULT_ENV = {"precision": "highp", "format": "rgba32f"}

# 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


class VulkanShaderGenerator:
    standard_header = """
#version 450 core
#define PRECISION $precision
#define FORMAT $format

"""

    def __init__(self: "VulkanShaderGenerator") -> None:
        self.ops_template_params: Dict[Any, Any] = {}

    def add_params_yaml(self, parameters_yaml_file):  # type: ignore[no-untyped-def]
        all_template_params = OrderedDict()
        with open(parameters_yaml_file) as f:
            contents = yaml.load(f, Loader=UniqueKeyLoader)
            for key in contents:
                all_template_params[key] = contents[key]
        self.validate_and_construct_op_params(all_template_params)  # type: ignore[no-untyped-call]

    def validate_and_construct_op_params(self, all_template_params):  # type: ignore[no-untyped-def]
        for op in all_template_params:
            if op in self.ops_template_params:
                raise KeyError(f"{op} params file has already been parsed")
            op_params_default_vals = all_template_params[op][
                "parameter_names_with_default_values"
            ]
            template_params_set = set(op_params_default_vals.keys())
            self.ops_template_params[op] = []
            self.ops_template_params[op].append(op_params_default_vals)
            op_template_params_values = all_template_params[op]["parameter_values"]
            for param_vals in op_template_params_values:
                param_vals_set = set(param_vals.keys())
                missing_keys = template_params_set - param_vals_set
                invalid_keys = param_vals_set - template_params_set
                if (len(invalid_keys)) > 0:
                    raise KeyError(f"Invalid keys {invalid_keys} are found")
                param_vals_copy = copy.deepcopy(op_params_default_vals)
                for key in param_vals:
                    param_vals_copy[key] = param_vals[key]
                self.ops_template_params[op].append(param_vals_copy)

    def generate(self, glsl_template_in, out_dir):  # type: ignore[no-untyped-def]
        glsl_template_name = os.path.basename(glsl_template_in)
        op_name, extension_name = glsl_template_name.split(".")
        if extension_name != "glslt":
            raise TypeError(f"invalid file type for glsl template {extension_name}")
        if op_name not in self.ops_template_params:
            raise KeyError(f"{op_name} params have not been populated")
        code_template = CodeTemplate.from_file(glsl_template_in)
        for template_params in self.ops_template_params[op_name]:
            content = VulkanShaderGenerator.standard_header
            output_file_name = template_params["NAME"] + ".glsl"
            content += code_template.substitute(template_params)
            output_file = os.path.join(out_dir, output_file_name)
            with open(output_file, "w") as f:
                f.write(content)


@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


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:])


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",
}

storageTypeToEnum = {
    "TEXTURE_2D": "api::StorageType::TEXTURE_2D",
    "TEXTURE_3D": "api::StorageType::TEXTURE_3D",
    "BUFFER": "api::StorageType::BUFFER",
    "": "api::StorageType::UNKNOWN",
}


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:
    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)

    return shader_info


def genGLSLFromGLSLT(src_dir_path: str, tmp_dir_path: str) -> None:
    template_dir_path = os.path.join(src_dir_path, "templates")
    vexs = glob.glob(os.path.join(template_dir_path, "**", "*.yaml"), recursive=True)
    parameter_yaml_files = []
    for f in vexs:
        if len(f) > 1:
            parameter_yaml_files.append(f)
    generator = VulkanShaderGenerator()
    for params_yaml in parameter_yaml_files:
        generator.add_params_yaml(params_yaml)  # type: ignore[no-untyped-call]

    vexs = glob.glob(os.path.join(src_dir_path, "**", "*.glslt"), recursive=True)
    templateSrcPaths = []
    for f in vexs:
        if len(f) > 1:
            templateSrcPaths.append(f)
            templateSrcPaths.sort()
    for glslt in templateSrcPaths:
        generator.generate(glslt, tmp_dir_path)  # type: ignore[no-untyped-call]


def genCppH(
    hFilePath: str,
    cppFilePath: str,
    srcDirPaths: str,
    glslcPath: str,
    tmpDirPath: str,
    env: Dict[Any, Any],
) -> None:
    print(
        "hFilePath:{} cppFilePath:{} srcDirPaths:{} glslcPath:{} tmpDirPath:{}".format(
            hFilePath, cppFilePath, srcDirPaths, glslcPath, tmpDirPath
        )
    )

    templateSrcPaths = []

    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:
            templateSrcPaths.append(f)
            templateSrcPaths.sort()
    print(f"templateSrcPaths:{templateSrcPaths}")

    spvPaths = {}
    for templateSrcPath in templateSrcPaths:
        print(f"templateSrcPath {templateSrcPath}")
        name = getName(templateSrcPath).replace("_glsl", "")
        print(f"name {name}")

        codeTemplate = CodeTemplate.from_file(templateSrcPath)
        srcPath = tmpDirPath + "/" + name + ".glsl"
        content = codeTemplate.substitute(env)
        with open(srcPath, "w") as fw:
            fw.write(content)

        spvPath = tmpDirPath + "/" + name + ".spv"
        print(f"spvPath {spvPath}")

        cmd = [
            glslcPath,
            "-fshader-stage=compute",
            srcPath,
            "-o",
            spvPath,
            "--target-env=vulkan1.0",
            "-Werror",
        ] + [arg for srcDirPath in srcDirPaths for arg in ["-I", srcDirPath]]

        print("\nglslc cmd:", cmd)

        subprocess.check_call(cmd)
        spvPaths[spvPath] = templateSrcPath

    h = "#pragma once\n"
    h += "#include <ATen/native/vulkan/api/Types.h>\n"
    h += "#include <ATen/native/vulkan/api/vk_api.h>\n"
    h += "#include <c10/util/flat_hash_map.h>\n"
    h += "#include <string>\n"

    nsbegin = "namespace at {\nnamespace native {\nnamespace vulkan {\n"
    nsend = "} // namespace vulkan\n} // namespace native\n} // namespace at\n"

    anon_ns_begin = "namespace {\n"
    anon_ns_end = "} // namespace\n"

    h += nsbegin

    # Forward declaration of ShaderInfo
    h += "namespace api {\nstruct ShaderInfo;\n} // namespace api\n"
    h += "typedef ska::flat_hash_map<std::string, api::ShaderInfo> ShaderListing;\n"
    h += "typedef ska::flat_hash_map<std::string, std::string> RegistryKeyMap;\n"
    h += "typedef ska::flat_hash_map<std::string, RegistryKeyMap> ShaderRegistry;\n"
    h += "extern const ShaderListing shader_infos;\n"
    h += "extern ShaderRegistry shader_registry;\n"
    h += "inline const ShaderListing& get_shader_infos() {\n  return shader_infos;\n}\n"
    h += (
        "inline ShaderRegistry& get_shader_registry() {\n  return shader_registry;\n}\n"
    )

    h += nsend

    cpp = "#include <ATen/native/vulkan/api/Shader.h>\n"
    cpp += f"#include <ATen/native/vulkan/{H_NAME}>\n"
    cpp += "#include <stdint.h>\n"
    cpp += "#include <vector>\n"
    cpp += nsbegin

    shader_info_bin_code = []
    shader_info_cpp_code = []
    shader_info_registry_code = []

    for spvPath, srcPath in spvPaths.items():
        name = getName(spvPath).replace("_spv", "")

        print(f"spvPath:{spvPath}")
        with open(spvPath, "rb") as fr:
            next_bin = array.array("I", fr.read())
            sizeBytes = 4 * len(next_bin)
            shader_info_bin_code.append(
                "const uint32_t {}_bin[] = {{\n{}\n}};".format(
                    name,
                    textwrap.indent(",\n".join(str(x) for x in next_bin), "  "),
                ),
            )

        shader_info = getShaderInfo(srcPath)

        tile_size = (
            f"{{{', '.join(str(x) for x in shader_info.tile_size)}}}"
            if (len(shader_info.tile_size) > 0)
            else "std::vector<uint32_t>()"
        )

        shader_info_layouts = "{{{}}}".format(",\n ".join(shader_info.layouts))

        shader_info_args = [
            f'"vulkan.{name}"',
            f"{name}_bin",
            str(sizeBytes),
            shader_info_layouts,
            tile_size,
            storageTypeToEnum[shader_info.weight_storage_type],
            storageTypeToEnum[shader_info.bias_storage_type],
        ]

        shader_info_cpp_code.append(
            textwrap.indent(
                '{{"{}",\n api::ShaderInfo(\n{})}}'.format(
                    name,
                    textwrap.indent(",\n".join(shader_info_args), "     "),
                ),
                "    ",
            ),
        )

        if shader_info.register_for is not None:
            (op_name, registry_keys) = shader_info.register_for
            for registry_key in registry_keys:
                shader_info_registry_code.append(
                    textwrap.indent(
                        f'{{"{op_name}", {{{{"{registry_key}", "{name}"}}}}}}',
                        "        ",
                    ),
                )

    cpp += anon_ns_begin
    cpp += "\n".join(shader_info_bin_code) + "\n"
    cpp += anon_ns_end

    cpp += "const ShaderListing shader_infos = {{\n{}}};\n".format(
        ",\n".join(shader_info_cpp_code),
    )
    cpp += "ShaderRegistry shader_registry = {{\n{}}};\n".format(
        ",\n".join(shader_info_registry_code),
    )
    cpp += nsend

    with open(hFilePath, "w") as fw:
        fw.write(h)
    with open(cppFilePath, "w") as fw:
        fw.write(cpp)


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")
    parser.add_argument("-o", "--output-path", required=True, help="")
    parser.add_argument(
        "--env", metavar="KEY=VALUE", nargs="*", help="Set a number of key-value pairs"
    )
    options = parser.parse_args()
    env = DEFAULT_ENV
    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)

    genCppH(
        hFilePath=options.output_path + "/spv.h",
        cppFilePath=options.output_path + "/spv.cpp",
        srcDirPaths=options.glsl_paths,
        glslcPath=options.glslc_path,
        tmpDirPath=options.tmp_dir_path,
        env=env,
    )

    return 0


def invoke_main() -> None:
    sys.exit(main(sys.argv))


if __name__ == "__main__":
    invoke_main()  # pragma: no cover