gen.py

import json
import io
import os
import re
from enum import Enum
from argparse import ArgumentParser
from typing import NamedTuple
from typing import Optional

head = """// Copyright (C) 2023 - DevSH Graphics Programming Sp. z O.O.
// This file is part of the "Nabla Engine".
// For conditions of distribution and use, see copyright notice in nabla.h
#ifndef _NBL_BUILTIN_HLSL_SPIRV_INTRINSICS_CORE_INCLUDED_
#define _NBL_BUILTIN_HLSL_SPIRV_INTRINSICS_CORE_INCLUDED_

#ifdef __HLSL_VERSION
#include "spirv/unified1/spirv.hpp"
#include "spirv/unified1/GLSL.std.450.h"
#endif

#include "nbl/builtin/hlsl/type_traits.hlsl"

namespace nbl 
{
namespace hlsl
{
#ifdef __HLSL_VERSION
namespace spirv
{

//! General Decls
template<uint32_t StorageClass, typename T>
using pointer_t = vk::SpirvOpaqueType<spv::OpTypePointer, vk::Literal< vk::integral_constant<uint32_t, StorageClass> >, T>;

// The holy operation that makes addrof possible
template<uint32_t StorageClass, typename T>
[[vk::ext_instruction(spv::OpCopyObject)]]
pointer_t<StorageClass, T> copyObject([[vk::ext_reference]] T value);

//! Std 450 Extended set operations
template<typename SquareMatrix>
[[vk::ext_instruction(GLSLstd450MatrixInverse)]]
SquareMatrix matrixInverse(NBL_CONST_REF_ARG(SquareMatrix) mat);

// Add specializations if you need to emit a `ext_capability` (this means that the instruction needs to forward through an `impl::` struct and so on)
template<typename T, typename U>
[[vk::ext_capability(spv::CapabilityPhysicalStorageBufferAddresses)]]
[[vk::ext_instruction(spv::OpBitcast)]]
enable_if_t<is_spirv_type_v<T> && is_spirv_type_v<U>, T> bitcast(U);

template<typename T>
[[vk::ext_capability(spv::CapabilityPhysicalStorageBufferAddresses)]]
[[vk::ext_instruction(spv::OpBitcast)]]
uint64_t bitcast(pointer_t<spv::StorageClassPhysicalStorageBuffer,T>);

template<typename T>
[[vk::ext_capability(spv::CapabilityPhysicalStorageBufferAddresses)]]
[[vk::ext_instruction(spv::OpBitcast)]]
pointer_t<spv::StorageClassPhysicalStorageBuffer,T> bitcast(uint64_t);

template<class T, class U>
[[vk::ext_instruction(spv::OpBitcast)]]
T bitcast(U);
"""

foot = """}

#endif
}
}

#endif
"""

def gen(grammer_path, output_path):
    grammer_raw = open(grammer_path, "r").read()
    grammer = json.loads(grammer_raw)
    del grammer_raw
            
    output = open(output_path, "w", buffering=1024**2)

    builtins = [x for x in grammer["operand_kinds"] if x["kind"] == "BuiltIn"][0]["enumerants"]
    execution_modes = [x for x in grammer["operand_kinds"] if x["kind"] == "ExecutionMode"][0]["enumerants"]
    group_operations = [x for x in grammer["operand_kinds"] if x["kind"] == "GroupOperation"][0]["enumerants"]

    with output as writer:
        writer.write(head)

        writer.write("\n//! Builtins\nnamespace builtin\n{")
        for b in builtins:
            builtin_type = None
            is_output = False
            builtin_name = b["enumerant"]
            match builtin_name:
                case "HelperInvocation": builtin_type = "bool"
                case "VertexIndex": builtin_type = "uint32_t"
                case "InstanceIndex": builtin_type = "uint32_t"
                case "NumWorkgroups": builtin_type = "uint32_t3"
                case "WorkgroupId": builtin_type = "uint32_t3"
                case "LocalInvocationId": builtin_type = "uint32_t3"
                case "GlobalInvocationId": builtin_type = "uint32_t3"
                case "LocalInvocationIndex": builtin_type = "uint32_t"
                case "SubgroupEqMask": builtin_type = "uint32_t4"
                case "SubgroupGeMask": builtin_type = "uint32_t4"
                case "SubgroupGtMask": builtin_type = "uint32_t4"
                case "SubgroupLeMask": builtin_type = "uint32_t4"
                case "SubgroupLtMask": builtin_type = "uint32_t4"
                case "SubgroupSize": builtin_type = "uint32_t"
                case "NumSubgroups": builtin_type = "uint32_t"
                case "SubgroupId": builtin_type = "uint32_t"
                case "SubgroupLocalInvocationId": builtin_type = "uint32_t"
                case "Position":
                    builtin_type = "float32_t4"
                    is_output = True
                case _: continue
            if is_output:
                writer.write("[[vk::ext_builtin_output(spv::BuiltIn" + builtin_name + ")]]\n")
                writer.write("static " + builtin_type + " " + builtin_name + ";\n")
            else:
                writer.write("[[vk::ext_builtin_input(spv::BuiltIn" + builtin_name + ")]]\n")
                writer.write("static const " + builtin_type + " " + builtin_name + ";\n")
        writer.write("}\n")

        writer.write("\n//! Execution Modes\nnamespace execution_mode\n{")
        for em in execution_modes:
            name = em["enumerant"]
            name_l = name[0].lower() + name[1:]
            writer.write("\n\tvoid " + name_l + "()\n\t{\n\t\tvk::ext_execution_mode(spv::ExecutionMode" + name + ");\n\t}\n")
        writer.write("}\n")

        writer.write("\n//! Group Operations\nnamespace group_operation\n{\n")
        for go in group_operations:
            name = go["enumerant"]
            value = go["value"]
            writer.write("\tstatic const uint32_t " + name + " = " + str(value) + ";\n")
        writer.write("}\n")

        writer.write("\n//! Instructions\n")
        for instruction in grammer["instructions"]:
            match instruction["class"]:
                case "Atomic":
                    processInst(writer, instruction, InstOptions())
                    processInst(writer, instruction, InstOptions(shape=Shape.PTR_TEMPLATE))
                case "Memory":
                    processInst(writer, instruction, InstOptions(shape=Shape.PTR_TEMPLATE))
                    processInst(writer, instruction, InstOptions(shape=Shape.PSB_RT))
                case "Barrier" | "Bit":
                    processInst(writer, instruction, InstOptions())
                case "Reserved":
                    match instruction["opname"]:
                        case "OpBeginInvocationInterlockEXT" | "OpEndInvocationInterlockEXT":
                            processInst(writer, instruction, InstOptions())
                case "Non-Uniform":
                    match instruction["opname"]:
                        case "OpGroupNonUniformElect" | "OpGroupNonUniformAll" | "OpGroupNonUniformAny" | "OpGroupNonUniformAllEqual":
                            processInst(writer, instruction, InstOptions(result_ty="bool"))
                        case "OpGroupNonUniformBallot":
                            processInst(writer, instruction, InstOptions(result_ty="uint32_t4",op_ty="bool"))
                        case "OpGroupNonUniformInverseBallot" | "OpGroupNonUniformBallotBitExtract":
                            processInst(writer, instruction, InstOptions(result_ty="bool",op_ty="uint32_t4"))
                        case "OpGroupNonUniformBallotBitCount" | "OpGroupNonUniformBallotFindLSB" | "OpGroupNonUniformBallotFindMSB":
                            processInst(writer, instruction, InstOptions(result_ty="uint32_t",op_ty="uint32_t4"))
                        case _: processInst(writer, instruction, InstOptions())
                case _: continue # TODO

        writer.write(foot)

class Shape(Enum):
    DEFAULT = 0,
    PTR_TEMPLATE = 1, # TODO: this is a DXC Workaround
    PSB_RT = 2, # PhysicalStorageBuffer Result Type

class InstOptions(NamedTuple):
    shape: Shape = Shape.DEFAULT
    result_ty: Optional[str] = None
    op_ty: Optional[str] = None

def processInst(writer: io.TextIOWrapper, instruction, options: InstOptions):
    templates = []
    caps = []
    conds = []
    op_name = instruction["opname"]
    fn_name = op_name[2].lower() + op_name[3:]
    result_types = []

    if "capabilities" in instruction and len(instruction["capabilities"]) > 0:
        for cap in instruction["capabilities"]:
            if cap == "Shader" or cap == "Kernel": continue
            caps.append(cap)
    
    if options.shape == Shape.PTR_TEMPLATE:
        templates.append("typename P")
        conds.append("is_spirv_type_v<P>")
    
    # split upper case words
    matches = [(m.group(1), m.span(1)) for m in re.finditer(r'([A-Z])[A-Z][a-z]', fn_name)]

    for m in matches:
        match m[0]:
            case "I":
                conds.append("(is_signed_v<T> || is_unsigned_v<T>)")
                break
            case "U":
                fn_name = fn_name[0:m[1][0]] + fn_name[m[1][1]:]
                result_types = ["uint32_t", "uint64_t"]
                break
            case "S":
                fn_name = fn_name[0:m[1][0]] + fn_name[m[1][1]:]
                result_types = ["int32_t", "int64_t"]
                break
            case "F":
                fn_name = fn_name[0:m[1][0]] + fn_name[m[1][1]:]
                result_types = ["float"]
                break

    if "operands" in instruction:
        operands = instruction["operands"]
        if operands[0]["kind"] == "IdResultType":
            operands = operands[2:]
            if len(result_types) == 0:
                if options.result_ty == None:
                    result_types = ["T"]
                else:
                    result_types = [options.result_ty]
        else:
            assert len(result_types) == 0
            result_types = ["void"]

        for rt in result_types:
            op_ty = "T"
            if options.op_ty != None:
                op_ty = options.op_ty
            elif rt != "void":
                op_ty = rt
            
            if (not "typename T" in templates) and (rt == "T"):
                templates = ["typename T"] + templates

            args = []
            for operand in operands:
                operand_name = operand["name"].strip("'") if "name" in operand else None
                operand_name = operand_name[0].lower() + operand_name[1:] if (operand_name != None) else ""
                match operand["kind"]:
                    case "IdRef":
                        match operand["name"]:
                            case "'Pointer'":
                                if options.shape == Shape.PTR_TEMPLATE:
                                    args.append("P " + operand_name)
                                elif options.shape == Shape.PSB_RT:    
                                    if (not "typename T" in templates) and (rt == "T" or op_ty == "T"):
                                        templates = ["typename T"] + templates
                                    args.append("pointer_t<spv::StorageClassPhysicalStorageBuffer, " + op_ty + "> " + operand_name)
                                else:    
                                    if (not "typename T" in templates) and (rt == "T" or op_ty == "T"):
                                        templates = ["typename T"] + templates
                                    args.append("[[vk::ext_reference]] " + op_ty + " " + operand_name)
                            case "'Value'" | "'Object'" | "'Comparator'" | "'Base'" | "'Insert'":
                                if (not "typename T" in templates) and (rt == "T" or op_ty == "T"):
                                    templates = ["typename T"] + templates
                                args.append(op_ty + " " + operand_name)
                            case "'Offset'" | "'Count'" | "'Id'" | "'Index'" | "'Mask'" | "'Delta'":
                                args.append("uint32_t " + operand_name)
                            case "'Predicate'": args.append("bool " + operand_name)
                            case "'ClusterSize'":
                                if "quantifier" in operand and operand["quantifier"] == "?": continue # TODO: overload
                                else: return # TODO
                            case _: return # TODO
                    case "IdScope": args.append("uint32_t " + operand_name.lower() + "Scope")
                    case "IdMemorySemantics": args.append(" uint32_t " + operand_name)
                    case "GroupOperation": args.append("[[vk::ext_literal]] uint32_t " + operand_name)
                    case "MemoryAccess":
                        writeInst(writer, templates, caps, op_name, fn_name, conds, rt, args + ["[[vk::ext_literal]] uint32_t memoryAccess"])
                        writeInst(writer, templates, caps, op_name, fn_name, conds, rt, args + ["[[vk::ext_literal]] uint32_t memoryAccess, [[vk::ext_literal]] uint32_t memoryAccessParam"])
                        writeInst(writer, templates + ["uint32_t alignment"], caps, op_name, fn_name, conds, rt, args + ["[[vk::ext_literal]] uint32_t __aligned = /*Aligned*/0x00000002", "[[vk::ext_literal]] uint32_t __alignment = alignment"])
                    case _: return # TODO

            writeInst(writer, templates, caps, op_name, fn_name, conds, rt, args)


def writeInst(writer: io.TextIOWrapper, templates, caps, op_name, fn_name, conds, result_type, args):
    if len(caps) > 0: 
        for cap in caps:
            final_fn_name = fn_name
            if (len(caps) > 1): final_fn_name = fn_name + "_" + cap
            writeInstInner(writer, templates, cap, op_name, final_fn_name, conds, result_type, args)
    else:
        writeInstInner(writer, templates, None, op_name, fn_name, conds, result_type, args)

def writeInstInner(writer: io.TextIOWrapper, templates, cap, op_name, fn_name, conds, result_type, args):
    if len(templates) > 0:
        writer.write("template<" + ", ".join(templates) + ">\n")
    if (cap != None):
        writer.write("[[vk::ext_capability(spv::Capability" + cap + ")]]\n")
    writer.write("[[vk::ext_instruction(spv::" + op_name + ")]]\n")
    if len(conds) > 0:
        writer.write("enable_if_t<" + " && ".join(conds) + ", " + result_type + ">")
    else:
        writer.write(result_type)
    writer.write(" " + fn_name + "(" + ", ".join(args) + ");\n\n")


if __name__ == "__main__":
    script_dir_path = os.path.abspath(os.path.dirname(__file__))

    parser = ArgumentParser(description="Generate HLSL from SPIR-V instructions")
    parser.add_argument("output", type=str, help="HLSL output file")
    parser.add_argument("--grammer", required=False, type=str, help="Input SPIR-V grammer JSON file", default=os.path.join(script_dir_path, "../../include/spirv/unified1/spirv.core.grammar.json"))
    args = parser.parse_args()

    gen(args.grammer, args.output)