Merge pull request #434 from KhronosGroup/fix-429

Support gl_NumWorkgroups in HLSL.

main.cpp

@@ -887,6 +887,17 @@ static int main_inner(int argc, char *argv[])
 		}
 	}
 
+	if (args.hlsl)
+	{
+		auto *hlsl_compiler = static_cast<CompilerHLSL *>(compiler.get());
+		uint32_t new_builtin = hlsl_compiler->remap_num_workgroups_builtin();
+		if (new_builtin)
+		{
+			hlsl_compiler->set_decoration(new_builtin, DecorationDescriptorSet, 0);
+			hlsl_compiler->set_decoration(new_builtin, DecorationBinding, 0);
+		}
+	}
+
 	string glsl;
 	for (uint32_t i = 0; i < args.iterations; i++)
 	{

reference/opt/shaders-hlsl/comp/num-workgroups-alone.comp

@@ -0,0 +1,16 @@
+RWByteAddressBuffer _10 : register(u0);
+cbuffer SPIRV_Cross_NumWorkgroups : register(b0)
+{
+    uint3 SPIRV_Cross_NumWorkgroups_count : packoffset(c0);
+};
+
+void comp_main()
+{
+    _10.Store3(0, SPIRV_Cross_NumWorkgroups_count);
+}
+
+[numthreads(1, 1, 1)]
+void main()
+{
+    comp_main();
+}

reference/opt/shaders-hlsl/comp/num-workgroups-with-builtins.comp

@@ -0,0 +1,23 @@
+RWByteAddressBuffer _10 : register(u0);
+cbuffer SPIRV_Cross_NumWorkgroups : register(b0)
+{
+    uint3 SPIRV_Cross_NumWorkgroups_count : packoffset(c0);
+};
+
+static uint3 gl_WorkGroupID;
+struct SPIRV_Cross_Input
+{
+    uint3 gl_WorkGroupID : SV_GroupID;
+};
+
+void comp_main()
+{
+    _10.Store3(0, SPIRV_Cross_NumWorkgroups_count + gl_WorkGroupID);
+}
+
+[numthreads(1, 1, 1)]
+void main(SPIRV_Cross_Input stage_input)
+{
+    gl_WorkGroupID = stage_input.gl_WorkGroupID;
+    comp_main();
+}

reference/shaders-hlsl/comp/num-workgroups-alone.comp

@@ -0,0 +1,16 @@
+RWByteAddressBuffer _10 : register(u0);
+cbuffer SPIRV_Cross_NumWorkgroups : register(b0)
+{
+    uint3 SPIRV_Cross_NumWorkgroups_count : packoffset(c0);
+};
+
+void comp_main()
+{
+    _10.Store3(0, SPIRV_Cross_NumWorkgroups_count);
+}
+
+[numthreads(1, 1, 1)]
+void main()
+{
+    comp_main();
+}

reference/shaders-hlsl/comp/num-workgroups-with-builtins.comp

@@ -0,0 +1,23 @@
+RWByteAddressBuffer _10 : register(u0);
+cbuffer SPIRV_Cross_NumWorkgroups : register(b0)
+{
+    uint3 SPIRV_Cross_NumWorkgroups_count : packoffset(c0);
+};
+
+static uint3 gl_WorkGroupID;
+struct SPIRV_Cross_Input
+{
+    uint3 gl_WorkGroupID : SV_GroupID;
+};
+
+void comp_main()
+{
+    _10.Store3(0, SPIRV_Cross_NumWorkgroups_count + gl_WorkGroupID);
+}
+
+[numthreads(1, 1, 1)]
+void main(SPIRV_Cross_Input stage_input)
+{
+    gl_WorkGroupID = stage_input.gl_WorkGroupID;
+    comp_main();
+}

shaders-hlsl/comp/num-workgroups-alone.comp

@@ -0,0 +1,13 @@
+#version 310 es
+layout(local_size_x = 1) in;
+
+layout(std430, binding = 0) buffer SSBO
+{
+	uvec3 outdata;
+};
+
+void main()
+{
+	outdata = gl_NumWorkGroups;
+}
+

shaders-hlsl/comp/num-workgroups-with-builtins.comp

@@ -0,0 +1,13 @@
+#version 310 es
+layout(local_size_x = 1) in;
+
+layout(std430, binding = 0) buffer SSBO
+{
+	uvec3 outdata;
+};
+
+void main()
+{
+	outdata = gl_NumWorkGroups + gl_WorkGroupID;
+}
+

spirv_hlsl.cpp

@@ -588,6 +588,10 @@ void CompilerHLSL::emit_builtin_inputs_in_struct()
 			semantic = "SV_GroupID";
 			break;
 
+		case BuiltInNumWorkgroups:
+			// Handled specially.
+			break;
+
 		default:
 			SPIRV_CROSS_THROW("Unsupported builtin in HLSL.");
 			break;
@@ -772,6 +776,15 @@ std::string CompilerHLSL::builtin_to_glsl(spv::BuiltIn builtin, spv::StorageClas
 		return "gl_VertexID";
 	case BuiltInInstanceId:
 		return "gl_InstanceID";
+	case BuiltInNumWorkgroups:
+	{
+		if (!num_workgroups_builtin)
+			SPIRV_CROSS_THROW("NumWorkgroups builtin is used, but remap_num_workgroups_builtin() was not called. Cannot emit code for this builtin.");
+
+		auto &var = get<SPIRVariable>(num_workgroups_builtin);
+		auto &type = get<SPIRType>(var.basetype);
+		return sanitize_underscores(join(to_name(num_workgroups_builtin), "_", get_member_name(type.self, 0)));
+	}
 	default:
 		return CompilerGLSL::builtin_to_glsl(builtin, storage);
 	}
@@ -827,6 +840,10 @@ void CompilerHLSL::emit_builtin_variables()
 			type = "uint";
 			break;
 
+		case BuiltInNumWorkgroups:
+			// Handled specially.
+			break;
+
 		default:
 			SPIRV_CROSS_THROW(join("Unsupported builtin in HLSL: ", unsigned(builtin)));
 			break;
@@ -1106,7 +1123,7 @@ void CompilerHLSL::emit_resources()
 		return name1.compare(name2) < 0;
 	};
 
-	if (!input_variables.empty() || active_input_builtins)
+	if (!input_variables.empty() || (active_input_builtins & ~(1ull << BuiltInNumWorkgroups)))
 	{
 		require_input = true;
 		statement("struct SPIRV_Cross_Input");
@@ -1739,6 +1756,10 @@ void CompilerHLSL::emit_hlsl_entry_point()
 			statement(builtin, " = int(stage_input.", builtin, ");");
 			break;
 
+		case BuiltInNumWorkgroups:
+			// Handled specially.
+			break;
+
 		default:
 			statement(builtin, " = stage_input.", builtin, ";");
 			break;
@@ -3659,6 +3680,50 @@ string CompilerHLSL::compile(std::vector<HLSLVertexAttributeRemap> vertex_attrib
 	return compile();
 }
 
+uint32_t CompilerHLSL::remap_num_workgroups_builtin()
+{
+	update_active_builtins();
+	if ((active_input_builtins & (1ull << BuiltInNumWorkgroups)) == 0)
+		return 0;
+
+	// Create a new, fake UBO.
+	uint32_t offset = increase_bound_by(4);
+
+	uint32_t uint_type_id = offset;
+	uint32_t block_type_id = offset + 1;
+	uint32_t block_pointer_type_id = offset + 2;
+	uint32_t variable_id = offset + 3;
+
+	SPIRType uint_type;
+	uint_type.basetype = SPIRType::UInt;
+	uint_type.vecsize = 3;
+	uint_type.columns = 1;
+	set<SPIRType>(uint_type_id, uint_type);
+
+	SPIRType block_type;
+	block_type.basetype = SPIRType::Struct;
+	block_type.member_types.push_back(uint_type_id);
+	set<SPIRType>(block_type_id, block_type);
+	set_decoration(block_type_id, DecorationBlock);
+	set_member_name(block_type_id, 0, "count");
+	set_member_decoration(block_type_id, 0, DecorationOffset, 0);
+
+	SPIRType block_pointer_type = block_type;
+	block_pointer_type.pointer = true;
+	block_pointer_type.storage = StorageClassUniform;
+	block_pointer_type.parent_type = block_type_id;
+	auto &ptr_type = set<SPIRType>(block_pointer_type_id, block_pointer_type);
+
+	// Preserve self.
+	ptr_type.self = block_type_id;
+
+	set<SPIRVariable>(variable_id, block_pointer_type_id, StorageClassUniform);
+	meta[variable_id].decoration.alias = "SPIRV_Cross_NumWorkgroups";
+
+	num_workgroups_builtin = variable_id;
+	return variable_id;
+}
+
 string CompilerHLSL::compile()
 {
 	// Do not deal with ES-isms like precision, older extensions and such.

spirv_hlsl.hpp

@@ -68,6 +68,18 @@ public:
 	std::string compile(std::vector<HLSLVertexAttributeRemap> vertex_attributes);
 	std::string compile() override;
 
+	// This is a special HLSL workaround for the NumWorkGroups builtin.
+	// This does not exist in HLSL, so the calling application must create a dummy cbuffer in
+	// which the application will store this builtin.
+	// The cbuffer layout will be:
+	// cbuffer SPIRV_Cross_NumWorkgroups : register(b#, space#) { uint3 SPIRV_Cross_NumWorkgroups_count; };
+	// This must be called before compile().
+	// The function returns 0 if NumWorkGroups builtin is not statically used in the shader from the current entry point.
+	// If non-zero, this returns the variable ID of a cbuffer which corresponds to
+	// the cbuffer declared above. By default, no binding or descriptor set decoration is set,
+	// so the calling application should declare explicit bindings on this ID before calling compile().
+	uint32_t remap_num_workgroups_builtin();
+
 private:
 	std::string type_to_glsl(const SPIRType &type, uint32_t id = 0) override;
 	std::string image_type_hlsl(const SPIRType &type);
@@ -159,6 +171,8 @@ private:
 
 	void emit_io_block(const SPIRVariable &var);
 	std::string to_semantic(uint32_t vertex_location);
+
+	uint32_t num_workgroups_builtin = 0;
 };
 }