hlsl_generator: add missing invocation instructions

alichraghi · alichraghi · commit 08de8a520943 · 2024-09-09T00:45:29.000+03:30
Signed-off-by: Ali Cheraghi &lt;alichraghi@proton.me&gt;
diff --git a/tools/hlsl_generator/gen.py b/tools/hlsl_generator/gen.py
@@ -185,6 +185,7 @@ def processInst(writer: io.TextIOWrapper, instruction, options: InstOptions):
     op_name = instruction["opname"]
     fn_name = op_name[2].lower() + op_name[3:]
     result_types = []
+    exts = instruction["extensions"] if "extensions" in instruction else []
 
     if "capabilities" in instruction and len(instruction["capabilities"]) > 0:
         for cap in instruction["capabilities"]:
@@ -223,56 +224,55 @@ def processInst(writer: io.TextIOWrapper, instruction, options: InstOptions):
         case "Bit": 
             if len(result_types) == 0: conds.append("(is_signed_v<T> || is_unsigned_v<T>)")
 
-    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:
-            overload_caps = caps.copy()
-            match rt:
-                case "uint16_t" | "int16_t": overload_caps.append("Int16")
-                case "uint64_t" | "int64_t": overload_caps.append("Int64")
-                case "float16_t": overload_caps.append("Float16")
-                case "float64_t": overload_caps.append("Float64")
-
-            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")
+    if "operands" in instruction and instruction["operands"][0]["kind"] == "IdResultType":
+        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:
+        overload_caps = caps.copy()
+        match rt:
+            case "uint16_t" | "int16_t": overload_caps.append("Int16")
+            case "uint64_t" | "int64_t": overload_caps.append("Int64")
+            case "float16_t": overload_caps.append("Float16")
+            case "float64_t": overload_caps.append("Float64")
+
+        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
                 
-                op_ty = "T"
-                if options.op_ty != None:
-                    op_ty = options.op_ty
-                elif rt != "void":
-                    op_ty = rt
-
-                args = []
-                for operand in operands:
+                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 = []
+            if "operands" in instruction:
+                for operand in instruction["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 "IdResult" | "IdResultType": continue
                         case "IdRef":
                             match operand["name"]:
                                 case "'Pointer'":
@@ -295,34 +295,38 @@ def processInst(writer: io.TextIOWrapper, instruction, options: InstOptions):
                                 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
+                                    else: return ignore(op_name) # TODO
+                                case _: return ignore(op_name) # 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, exts, op_name, final_fn_name, conds, rt, args + ["[[vk::ext_literal]] uint32_t memoryAccess"])
+                                writeInst(writer, final_templates, cap, exts, 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, exts, 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 ignore(op_name) # TODO
 
-                writeInst(writer, final_templates, cap, op_name, final_fn_name, conds, rt, args)
+            writeInst(writer, final_templates, cap, exts, op_name, final_fn_name, conds, rt, args)
 
 
-def writeInst(writer: io.TextIOWrapper, templates, cap, op_name, fn_name, conds, result_type, args):
+def writeInst(writer: io.TextIOWrapper, templates, cap, exts, op_name, fn_name, conds, result_type, args):
     if len(templates) > 0:
         writer.write("template<" + ", ".join(templates) + ">\n")
-    if (cap != None):
+    if cap != None:
         writer.write("[[vk::ext_capability(spv::Capability" + cap + ")]]\n")
+    for ext in exts:
+        writer.write("[[vk::ext_extension(\"" + ext + "\")]]\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")
 
+def ignore(op_name):
+    print("\033[93mWARNING\033[0m: instruction " + op_name + " ignored")
 
 if __name__ == "__main__":
     script_dir_path = os.path.abspath(os.path.dirname(__file__))
diff --git a/tools/hlsl_generator/out.hlsl b/tools/hlsl_generator/out.hlsl
@@ -1365,9 +1365,40 @@ T groupNonUniformQuadAnyKHR(bool predicate);
 
 template<typename T>
 [[vk::ext_capability(spv::CapabilityGroupNonUniformPartitionedNV)]]
+[[vk::ext_extension("SPV_NV_shader_subgroup_partitioned")]]
 [[vk::ext_instruction(spv::OpGroupNonUniformPartitionNV)]]
 T groupNonUniformPartitionNV(T value);
 
+[[vk::ext_capability(spv::CapabilityFragmentShaderSampleInterlockEXT)]]
+[[vk::ext_extension("SPV_EXT_fragment_shader_interlock")]]
+[[vk::ext_instruction(spv::OpBeginInvocationInterlockEXT)]]
+void beginInvocationInterlockEXT_FragmentShaderSampleInterlockEXT();
+
+[[vk::ext_capability(spv::CapabilityFragmentShaderPixelInterlockEXT)]]
+[[vk::ext_extension("SPV_EXT_fragment_shader_interlock")]]
+[[vk::ext_instruction(spv::OpBeginInvocationInterlockEXT)]]
+void beginInvocationInterlockEXT_FragmentShaderPixelInterlockEXT();
+
+[[vk::ext_capability(spv::CapabilityFragmentShaderShadingRateInterlockEXT)]]
+[[vk::ext_extension("SPV_EXT_fragment_shader_interlock")]]
+[[vk::ext_instruction(spv::OpBeginInvocationInterlockEXT)]]
+void beginInvocationInterlockEXT_FragmentShaderShadingRateInterlockEXT();
+
+[[vk::ext_capability(spv::CapabilityFragmentShaderSampleInterlockEXT)]]
+[[vk::ext_extension("SPV_EXT_fragment_shader_interlock")]]
+[[vk::ext_instruction(spv::OpEndInvocationInterlockEXT)]]
+void endInvocationInterlockEXT_FragmentShaderSampleInterlockEXT();
+
+[[vk::ext_capability(spv::CapabilityFragmentShaderPixelInterlockEXT)]]
+[[vk::ext_extension("SPV_EXT_fragment_shader_interlock")]]
+[[vk::ext_instruction(spv::OpEndInvocationInterlockEXT)]]
+void endInvocationInterlockEXT_FragmentShaderPixelInterlockEXT();
+
+[[vk::ext_capability(spv::CapabilityFragmentShaderShadingRateInterlockEXT)]]
+[[vk::ext_extension("SPV_EXT_fragment_shader_interlock")]]
+[[vk::ext_instruction(spv::OpEndInvocationInterlockEXT)]]
+void endInvocationInterlockEXT_FragmentShaderShadingRateInterlockEXT();
+
 [[vk::ext_capability(spv::CapabilityAtomicFloat16MinMaxEXT)]]
 [[vk::ext_instruction(spv::OpAtomicFMinEXT)]]
 float16_t atomicMinEXT_AtomicFloat16MinMaxEXT([[vk::ext_reference]] float16_t pointer, uint32_t memoryScope,  uint32_t semantics, float16_t value);
@@ -1461,48 +1492,58 @@ template<typename P>
 enable_if_t<is_spirv_type_v<P>, float64_t> atomicMaxEXT_Float64(P pointer, uint32_t memoryScope,  uint32_t semantics, float64_t value);
 
 [[vk::ext_capability(spv::CapabilityAtomicFloat16AddEXT)]]
+[[vk::ext_extension("SPV_EXT_shader_atomic_float_add")]]
 [[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_extension("SPV_EXT_shader_atomic_float_add")]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 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_extension("SPV_EXT_shader_atomic_float_add")]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 float32_t atomicAddEXT_AtomicFloat32AddEXT([[vk::ext_reference]] float32_t pointer, uint32_t memoryScope,  uint32_t semantics, float32_t value);
 
 [[vk::ext_capability(spv::CapabilityAtomicFloat64AddEXT)]]
+[[vk::ext_extension("SPV_EXT_shader_atomic_float_add")]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 float64_t atomicAddEXT_AtomicFloat64AddEXT([[vk::ext_reference]] float64_t pointer, uint32_t memoryScope,  uint32_t semantics, float64_t value);
 
 [[vk::ext_capability(spv::CapabilityFloat64)]]
+[[vk::ext_extension("SPV_EXT_shader_atomic_float_add")]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 float64_t atomicAddEXT_Float64([[vk::ext_reference]] float64_t pointer, uint32_t memoryScope,  uint32_t semantics, float64_t value);
 
 template<typename P>
 [[vk::ext_capability(spv::CapabilityAtomicFloat16AddEXT)]]
+[[vk::ext_extension("SPV_EXT_shader_atomic_float_add")]]
 [[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, typename N>
 [[vk::ext_capability(spv::CapabilityAtomicFloat16VectorNV)]]
+[[vk::ext_extension("SPV_EXT_shader_atomic_float_add")]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 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)]]
+[[vk::ext_extension("SPV_EXT_shader_atomic_float_add")]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 enable_if_t<is_spirv_type_v<P>, float32_t> atomicAddEXT_AtomicFloat32AddEXT(P pointer, uint32_t memoryScope,  uint32_t semantics, float32_t value);
 
 template<typename P>
 [[vk::ext_capability(spv::CapabilityAtomicFloat64AddEXT)]]
+[[vk::ext_extension("SPV_EXT_shader_atomic_float_add")]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 enable_if_t<is_spirv_type_v<P>, float64_t> atomicAddEXT_AtomicFloat64AddEXT(P pointer, uint32_t memoryScope,  uint32_t semantics, float64_t value);
 
 template<typename P>
 [[vk::ext_capability(spv::CapabilityFloat64)]]
+[[vk::ext_extension("SPV_EXT_shader_atomic_float_add")]]
 [[vk::ext_instruction(spv::OpAtomicFAddEXT)]]
 enable_if_t<is_spirv_type_v<P>, float64_t> atomicAddEXT_Float64(P pointer, uint32_t memoryScope,  uint32_t semantics, float64_t value);
 
