hlsl_generator: fix vector instructions type

Signed-off-by: Ali Cheraghi <[email protected]>

Author: alichraghi

tools/hlsl_generator/gen.py

@@ -195,6 +195,8 @@ def processInst(writer: io.TextIOWrapper, instruction, options: InstOptions):
     if options.shape == Shape.PTR_TEMPLATE:
         templates.append("typename P")
         conds.append("is_spirv_type_v<P>")
+    elif options.shape == Shape.BDA:
+        caps.append("PhysicalStorageBufferAddresses")
 
     # 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)]
@@ -242,74 +244,74 @@ def processInst(writer: io.TextIOWrapper, instruction, options: InstOptions):
                 case "float16_t": overload_caps.append("Float16")
                 case "float64_t": overload_caps.append("Float64")
 
-            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.BDA:    
-                                    if (not "typename T" in templates) and (rt == "T" or op_ty == "T"):
-                                        templates = ["typename T"] + templates
-                                    overload_caps.append("PhysicalStorageBufferAddresses")
-                                    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":
-                        if options.shape != Shape.BDA:
-                            writeInst(writer, templates, overload_caps, op_name, fn_name, conds, rt, args + ["[[vk::ext_literal]] uint32_t memoryAccess"])
-                            writeInst(writer, templates, overload_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"], overload_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, overload_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:
-            if (("Float16" in cap and result_type != "float16_t") or
-                ("Float32" in cap and result_type != "float32_t") or
-                ("Float64" in cap and result_type != "float64_t") or
-                ("Int16" in cap and result_type != "int16_t" and result_type != "uint16_t") or
-                ("Int64" in cap and result_type != "int64_t" and result_type != "uint64_t")): continue
-            
-            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):
+            for cap in overload_caps or [None]:
+                final_fn_name = fn_name + "_" + cap if (len(overload_caps) > 1) else fn_name
+                final_templates = templates.copy()
+                
+                if (not "typename T" in final_templates) and (rt == "T"):
+                    final_templates = ["typename T"] + final_templates
+
+                if len(overload_caps) > 0:
+                    if (("Float16" in cap and rt != "float16_t") or
+                        ("Float32" in cap and rt != "float32_t") or
+                        ("Float64" in cap and rt != "float64_t") or
+                        ("Int16" in cap and rt != "int16_t" and rt != "uint16_t") or
+                        ("Int64" in cap and rt != "int64_t" and rt != "uint64_t")): continue
+                    
+                    if "Vector" in cap:
+                        rt = "vector<" + rt + ", N> "
+                        final_templates.append("typename N")
+                
+                op_ty = "T"
+                if options.op_ty != None:
+                    op_ty = options.op_ty
+                elif rt != "void":
+                    op_ty = rt
+
+                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.BDA:    
+                                        if (not "typename T" in final_templates) and (rt == "T" or op_ty == "T"):
+                                            final_templates = ["typename T"] + final_templates
+                                        args.append("pointer_t<spv::StorageClassPhysicalStorageBuffer, " + op_ty + "> " + operand_name)
+                                    else:    
+                                        if (not "typename T" in final_templates) and (rt == "T" or op_ty == "T"):
+                                            final_templates = ["typename T"] + final_templates
+                                        args.append("[[vk::ext_reference]] " + op_ty + " " + operand_name)
+                                case "'Value'" | "'Object'" | "'Comparator'" | "'Base'" | "'Insert'":
+                                    if (not "typename T" in final_templates) and (rt == "T" or op_ty == "T"):
+                                        final_templates = ["typename T"] + final_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":
+                            assert len(overload_caps) <= 1
+                            if options.shape != Shape.BDA:
+                                writeInst(writer, final_templates, cap, op_name, final_fn_name, conds, rt, args + ["[[vk::ext_literal]] uint32_t memoryAccess"])
+                                writeInst(writer, final_templates, cap, op_name, final_fn_name, conds, rt, args + ["[[vk::ext_literal]] uint32_t memoryAccess, [[vk::ext_literal]] uint32_t memoryAccessParam"])
+                            writeInst(writer, final_templates + ["uint32_t alignment"], cap, op_name, final_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, final_templates, cap, op_name, final_fn_name, conds, rt, args)
+
+
+def writeInst(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):

tools/hlsl_generator/out.hlsl

@@ -1372,13 +1372,10 @@ T groupNonUniformPartitionNV(T value);
 [[vk::ext_instruction(spv::OpAtomicFMinEXT)]]
 float16_t atomicMinEXT_AtomicFloat16MinMaxEXT([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
 
+template<typename N>
 [[vk::ext_capability(spv::CapabilityAtomicFloat16VectorNV)]]
 [[vk::ext_instruction(spv::OpAtomicFMinEXT)]]
-float16_t atomicMinEXT_AtomicFloat16VectorNV([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
-
-[[vk::ext_capability(spv::CapabilityFloat16)]]
-[[vk::ext_instruction(spv::OpAtomicFMinEXT)]]
-float16_t atomicMinEXT_Float16([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
+vector<float16_t, N>  atomicMinEXT_AtomicFloat16VectorNV([[vk::ext_reference]] vector<float16_t, N>  pointer, uint32_t memoryScope,  uint32_t semantics, vector<float16_t, N>  value);
 
 [[vk::ext_capability(spv::CapabilityAtomicFloat32MinMaxEXT)]]
 [[vk::ext_instruction(spv::OpAtomicFMinEXT)]]
@@ -1397,15 +1394,10 @@ template<typename P>
 [[vk::ext_instruction(spv::OpAtomicFMinEXT)]]
 enable_if_t<is_spirv_type_v<P>, float16_t> atomicMinEXT_AtomicFloat16MinMaxEXT(P pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
 
-template<typename P>
+template<typename P, typename N>
 [[vk::ext_capability(spv::CapabilityAtomicFloat16VectorNV)]]
 [[vk::ext_instruction(spv::OpAtomicFMinEXT)]]
-enable_if_t<is_spirv_type_v<P>, float16_t> atomicMinEXT_AtomicFloat16VectorNV(P pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
-
-template<typename P>
-[[vk::ext_capability(spv::CapabilityFloat16)]]
-[[vk::ext_instruction(spv::OpAtomicFMinEXT)]]
-enable_if_t<is_spirv_type_v<P>, float16_t> atomicMinEXT_Float16(P pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
+enable_if_t<is_spirv_type_v<P>, vector<float16_t, N> > atomicMinEXT_AtomicFloat16VectorNV(P pointer, uint32_t memoryScope,  uint32_t semantics, vector<float16_t, N>  value);
 
 template<typename P>
 [[vk::ext_capability(spv::CapabilityAtomicFloat32MinMaxEXT)]]
@@ -1426,13 +1418,10 @@ enable_if_t<is_spirv_type_v<P>, float64_t> atomicMinEXT_Float64(P pointer, uint3
 [[vk::ext_instruction(spv::OpAtomicFMaxEXT)]]
 float16_t atomicMaxEXT_AtomicFloat16MinMaxEXT([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
 
+template<typename N>
 [[vk::ext_capability(spv::CapabilityAtomicFloat16VectorNV)]]
 [[vk::ext_instruction(spv::OpAtomicFMaxEXT)]]
-float16_t atomicMaxEXT_AtomicFloat16VectorNV([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
-
-[[vk::ext_capability(spv::CapabilityFloat16)]]
-[[vk::ext_instruction(spv::OpAtomicFMaxEXT)]]
-float16_t atomicMaxEXT_Float16([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
+vector<float16_t, N>  atomicMaxEXT_AtomicFloat16VectorNV([[vk::ext_reference]] vector<float16_t, N>  pointer, uint32_t memoryScope,  uint32_t semantics, vector<float16_t, N>  value);
 
 [[vk::ext_capability(spv::CapabilityAtomicFloat32MinMaxEXT)]]
 [[vk::ext_instruction(spv::OpAtomicFMaxEXT)]]
@@ -1451,15 +1440,10 @@ template<typename P>
 [[vk::ext_instruction(spv::OpAtomicFMaxEXT)]]
 enable_if_t<is_spirv_type_v<P>, float16_t> atomicMaxEXT_AtomicFloat16MinMaxEXT(P pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
 
-template<typename P>
+template<typename P, typename N>
 [[vk::ext_capability(spv::CapabilityAtomicFloat16VectorNV)]]
 [[vk::ext_instruction(spv::OpAtomicFMaxEXT)]]
-enable_if_t<is_spirv_type_v<P>, float16_t> atomicMaxEXT_AtomicFloat16VectorNV(P pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
-
-template<typename P>
-[[vk::ext_capability(spv::CapabilityFloat16)]]
-[[vk::ext_instruction(spv::OpAtomicFMaxEXT)]]
-enable_if_t<is_spirv_type_v<P>, float16_t> atomicMaxEXT_Float16(P pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
+enable_if_t<is_spirv_type_v<P>, vector<float16_t, N> > atomicMaxEXT_AtomicFloat16VectorNV(P pointer, uint32_t memoryScope,  uint32_t semantics, vector<float16_t, N>  value);
 
 template<typename P>
 [[vk::ext_capability(spv::CapabilityAtomicFloat32MinMaxEXT)]]
@@ -1480,13 +1464,10 @@ enable_if_t<is_spirv_type_v<P>, float64_t> atomicMaxEXT_Float64(P pointer, uint3
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 float16_t atomicAddEXT_AtomicFloat16AddEXT([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
 
+template<typename N>
 [[vk::ext_capability(spv::CapabilityAtomicFloat16VectorNV)]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
-float16_t atomicAddEXT_AtomicFloat16VectorNV([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
-
-[[vk::ext_capability(spv::CapabilityFloat16)]]
-[[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
-float16_t atomicAddEXT_Float16([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
+vector<float16_t, N>  atomicAddEXT_AtomicFloat16VectorNV([[vk::ext_reference]] vector<float16_t, N>  pointer, uint32_t memoryScope,  uint32_t semantics, vector<float16_t, N>  value);
 
 [[vk::ext_capability(spv::CapabilityAtomicFloat32AddEXT)]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
@@ -1505,15 +1486,10 @@ template<typename P>
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 enable_if_t<is_spirv_type_v<P>, float16_t> atomicAddEXT_AtomicFloat16AddEXT(P pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
 
-template<typename P>
+template<typename P, typename N>
 [[vk::ext_capability(spv::CapabilityAtomicFloat16VectorNV)]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
-enable_if_t<is_spirv_type_v<P>, float16_t> atomicAddEXT_AtomicFloat16VectorNV(P pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
-
-template<typename P>
-[[vk::ext_capability(spv::CapabilityFloat16)]]
-[[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
-enable_if_t<is_spirv_type_v<P>, float16_t> atomicAddEXT_Float16(P pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
+enable_if_t<is_spirv_type_v<P>, vector<float16_t, N> > atomicAddEXT_AtomicFloat16VectorNV(P pointer, uint32_t memoryScope,  uint32_t semantics, vector<float16_t, N>  value);
 
 template<typename P>
 [[vk::ext_capability(spv::CapabilityAtomicFloat32AddEXT)]]