Compute: WorkgroupSize decorations take precedence over LocalSize

src/Pipeline/ComputeProgram.cpp

@@ -49,15 +49,17 @@ namespace sw
 
 		auto &modes = shader->getModes();
 
+		int localSize[3] = {modes.WorkgroupSizeX, modes.WorkgroupSizeY, modes.WorkgroupSizeZ};
+
+		const int subgroupSize = SIMD::Width;
 
 		// Total number of invocations required to execute this workgroup.
-		const int numInvocations = modes.LocalSizeX * modes.LocalSizeY * modes.LocalSizeZ;
+		int numInvocations = localSize[0] * localSize[1] * localSize[2];
 
 		Int4 numWorkgroups = *Pointer<Int4>(data + OFFSET(Data, numWorkgroups));
 		Int4 workgroupID = *Pointer<Int4>(data + OFFSET(Data, workgroupID));
-		Int4 workgroupSize = Int4(modes.LocalSizeX, modes.LocalSizeY, modes.LocalSizeZ, 0);
-		const int subgroupSize = SIMD::Width;
-		const int numSubgroups = (numInvocations + subgroupSize - 1) / subgroupSize;
+		Int4 workgroupSize = Int4(localSize[0], localSize[1], localSize[2], 0);
+		Int numSubgroups = (numInvocations + subgroupSize - 1) / subgroupSize;
 
 		setInputBuiltin(spv::BuiltInNumWorkgroups, [&](const SpirvShader::BuiltinMapping& builtin, Array<Float4>& value)
 		{
@@ -107,10 +109,10 @@ namespace sw
 			Int4 localInvocationID[3];
 			{
 				Int4 idx = localInvocationIndex;
-				localInvocationID[ZZZZ] = idx / Int4(modes.LocalSizeX * modes.LocalSizeY);
-				idx -= localInvocationID[ZZZZ] * Int4(modes.LocalSizeX * modes.LocalSizeY); // modulo
-				localInvocationID[YYYY] = idx / Int4(modes.LocalSizeX);
-				idx -= localInvocationID[YYYY] * Int4(modes.LocalSizeX); // modulo
+				localInvocationID[ZZZZ] = idx / Int4(localSize[0] * localSize[1]);
+				idx -= localInvocationID[ZZZZ] * Int4(localSize[0] * localSize[1]); // modulo
+				localInvocationID[YYYY] = idx / Int4(localSize[0]);
+				idx -= localInvocationID[YYYY] * Int4(localSize[0]); // modulo
 				localInvocationID[XXXX] = idx;
 			}
 

src/Pipeline/SpirvShader.cpp

@@ -236,6 +236,25 @@ namespace sw
 					for (auto j = 0u; j < constituentTy.sizeInComponents; j++)
 						object.constantValue[offset++] = constituent.constantValue[j];
 				}
+
+				auto objectId = Object::ID(insn.word(2));
+				auto decorationsIt = decorations.find(objectId);
+				if (decorationsIt != decorations.end() &&
+					decorationsIt->second.BuiltIn == spv::BuiltInWorkgroupSize)
+				{
+					// https://www.khronos.org/registry/vulkan/specs/1.1/html/vkspec.html#interfaces-builtin-variables :
+					// Decorating an object with the WorkgroupSize built-in
+					// decoration will make that object contain the dimensions
+					// of a local workgroup. If an object is decorated with the
+					// WorkgroupSize decoration, this must take precedence over
+					// any execution mode set for LocalSize.
+					// The object decorated with WorkgroupSize must be declared
+					// as a three-component vector of 32-bit integers.
+					ASSERT(getType(object.type).sizeInComponents == 3);
+					modes.WorkgroupSizeX = object.constantValue[0];
+					modes.WorkgroupSizeY = object.constantValue[1];
+					modes.WorkgroupSizeZ = object.constantValue[2];
+				}
 				break;
 			}
 
@@ -555,9 +574,9 @@ namespace sw
 			modes.DepthUnchanged = true;
 			break;
 		case spv::ExecutionModeLocalSize:
-			modes.LocalSizeX = insn.word(3);
-			modes.LocalSizeZ = insn.word(5);
-			modes.LocalSizeY = insn.word(4);
+			modes.WorkgroupSizeX = insn.word(3);
+			modes.WorkgroupSizeY = insn.word(4);
+			modes.WorkgroupSizeZ = insn.word(5);
 			break;
 		case spv::ExecutionModeOriginUpperLeft:
 			// This is always the case for a Vulkan shader. Do nothing.

src/Pipeline/SpirvShader.hpp

@@ -289,7 +289,7 @@ namespace sw
 			bool NeedsCentroid : 1;
 
 			// Compute workgroup dimensions
-			int LocalSizeX, LocalSizeY, LocalSizeZ;
+			int WorkgroupSizeX = 1, WorkgroupSizeY = 1, WorkgroupSizeZ = 1;
 		};
 
 		Modes const &getModes() const

src/Vulkan/VkPipeline.cpp

@@ -20,6 +20,8 @@
 
 #include "spirv-tools/optimizer.hpp"
 
+#include <iostream>
+
 namespace
 {
 
@@ -230,6 +232,17 @@ std::vector<uint32_t> preprocessSpirv(
 
 	std::vector<uint32_t> optimized;
 	opt.Run(code.data(), code.size(), &optimized);
+
+	if (false) {
+		spvtools::SpirvTools core(SPV_ENV_VULKAN_1_1);
+		std::string preOpt;
+		core.Disassemble(code, &preOpt);
+		std::string postOpt;
+		core.Disassemble(optimized, &postOpt);
+		std::cout << "PRE-OPT: " << preOpt << std::endl
+		 		<< "POST-OPT: " << postOpt << std::endl;
+	}
+
 	return optimized;
 }