SpirvShader: Implement workgroup storage.

src/Pipeline/ComputeProgram.cpp

@@ -51,6 +51,7 @@ namespace sw
 		routine.descriptorDynamicOffsets = data + OFFSET(Data, descriptorDynamicOffsets);
 		routine.pushConstants = data + OFFSET(Data, pushConstants);
 		routine.constants = *Pointer<Pointer<Byte>>(data + OFFSET(Data, constants));
+		routine.workgroupMemory = *Pointer<Pointer<Byte>>(data + OFFSET(Data, workgroupMemory));
 
 		auto &modes = shader->getModes();
 
@@ -175,7 +176,7 @@ namespace sw
 	}
 
 	void ComputeProgram::run(
-		Routine *routine,
+		Routine *routine, SpirvShader const *shader,
 		vk::DescriptorSet::Bindings const &descriptorSets,
 		vk::DescriptorSet::DynamicOffsets const &descriptorDynamicOffsets,
 		PushConstantStorage const &pushConstants,
@@ -183,6 +184,11 @@ namespace sw
 	{
 		auto runWorkgroup = (void(*)(void*))(routine->getEntry());
 
+		// We're sharing a buffer here across all workgroups.
+		// We can only do this because we know workgroups are executed
+		// serially.
+		std::vector<uint8_t> workgroupMemory(shader->workgroupMemory.size());
+
 		Data data;
 		data.descriptorSets = descriptorSets;
 		data.descriptorDynamicOffsets = descriptorDynamicOffsets;
@@ -192,6 +198,7 @@ namespace sw
 		data.numWorkgroups[3] = 0;
 		data.pushConstants = pushConstants;
 		data.constants = &sw::constants;
+		data.workgroupMemory = workgroupMemory.data();
 
 		// TODO(bclayton): Split work across threads.
 		for (uint32_t groupZ = 0; groupZ < groupCountZ; groupZ++)

src/Pipeline/ComputeProgram.hpp

@@ -50,7 +50,7 @@ namespace sw
 		// run executes the compute shader routine for all workgroups.
 		// TODO(bclayton): This probably does not belong here. Consider moving.
 		static void run(
-			Routine *routine,
+			Routine *routine, SpirvShader const *shader,
 			vk::DescriptorSet::Bindings const &descriptorSetBindings,
 			vk::DescriptorSet::DynamicOffsets const &descriptorDynamicOffsets,
 			PushConstantStorage const &pushConstants,
@@ -71,6 +71,7 @@ namespace sw
 			uint4 workgroupID;
 			PushConstantStorage pushConstants;
 			const Constants *constants;
+			uint8_t* workgroupMemory;
 		};
 
 		SpirvRoutine routine;

src/Pipeline/SpirvShader.cpp

@@ -540,6 +540,7 @@ namespace sw
 				object.definition = insn;
 				object.type = typeId;
 
+				ASSERT(getType(typeId).definition.opcode() == spv::OpTypePointer);
 				ASSERT(getType(typeId).storageClass == storageClass);
 
 				switch (storageClass)
@@ -561,6 +562,13 @@ namespace sw
 					break; // Correctly handled.
 
 				case spv::StorageClassWorkgroup:
+				{
+					auto &elTy = getType(getType(typeId).element);
+					auto sizeInBytes = elTy.sizeInComponents * sizeof(float);
+					workgroupMemory.allocate(resultId, sizeInBytes);
+					object.kind = Object::Kind::Pointer;
+					break;
+				}
 				case spv::StorageClassAtomicCounter:
 				case spv::StorageClassImage:
 					UNIMPLEMENTED("StorageClass %d not yet implemented", (int)storageClass);
@@ -1147,6 +1155,7 @@ namespace sw
 		case spv::StorageClassUniform:
 		case spv::StorageClassStorageBuffer:
 		case spv::StorageClassPushConstant:
+		case spv::StorageClassWorkgroup:
 			return false;
 		default:
 			return true;
@@ -2479,6 +2488,14 @@ namespace sw
 			routine->createPointer(resultId, SIMD::Pointer(base, size));
 			break;
 		}
+		case spv::StorageClassWorkgroup:
+		{
+			ASSERT(objectTy.opcode() == spv::OpTypePointer);
+			auto base = &routine->workgroupMemory[0];
+			auto size = workgroupMemory.size();
+			routine->createPointer(resultId, SIMD::Pointer(base, size, workgroupMemory.offsetOf(resultId)));
+			break;
+		}
 		case spv::StorageClassInput:
 		{
 			if (object.kind == Object::Kind::InterfaceVariable)

src/Pipeline/SpirvShader.hpp

@@ -656,6 +656,30 @@ namespace sw
 			uint32_t SizeInComponents;
 		};
 
+		struct WorkgroupMemory
+		{
+			// allocates a new variable of size bytes with the given identifier.
+			inline void allocate(Object::ID id, uint32_t size)
+			{
+				uint32_t offset = totalSize;
+				auto it = offsets.emplace(id, offset);
+				ASSERT_MSG(it.second, "WorkgroupMemory already has an allocation for object %d", int(id.value()));
+				totalSize += size;
+			}
+			// returns the byte offset of the variable with the given identifier.
+			inline uint32_t offsetOf(Object::ID id) const
+			{
+				auto it = offsets.find(id);
+				ASSERT_MSG(it != offsets.end(), "WorkgroupMemory has no allocation for object %d", int(id.value()));
+				return it->second;
+			}
+			// returns the total allocated size in bytes.
+			inline uint32_t size() const { return totalSize; }
+		private:
+			uint32_t totalSize = 0; // in bytes
+			std::unordered_map<Object::ID, uint32_t> offsets; // in bytes
+		};
+
 		std::vector<InterfaceComponent> inputs;
 		std::vector<InterfaceComponent> outputs;
 
@@ -666,6 +690,7 @@ namespace sw
 		using BuiltInHash = std::hash<std::underlying_type<spv::BuiltIn>::type>;
 		std::unordered_map<spv::BuiltIn, BuiltinMapping, BuiltInHash> inputBuiltins;
 		std::unordered_map<spv::BuiltIn, BuiltinMapping, BuiltInHash> outputBuiltins;
+		WorkgroupMemory workgroupMemory;
 
 		Type const &getType(Type::ID id) const
 		{
@@ -956,6 +981,7 @@ namespace sw
 		Variable inputs = Variable{MAX_INTERFACE_COMPONENTS};
 		Variable outputs = Variable{MAX_INTERFACE_COMPONENTS};
 
+		Pointer<Byte> workgroupMemory;
 		Pointer<Pointer<Byte>> descriptorSets;
 		Pointer<Int> descriptorDynamicOffsets;
 		Pointer<Byte> pushConstants;

src/Vulkan/VkPipeline.cpp

@@ -568,7 +568,8 @@ void ComputePipeline::run(uint32_t groupCountX, uint32_t groupCountY, uint32_t g
 {
 	ASSERT_OR_RETURN(routine != nullptr);
 	sw::ComputeProgram::run(
-		routine, descriptorSets, descriptorDynamicOffsets, pushConstants,
+		routine, shader,
+		descriptorSets, descriptorDynamicOffsets, pushConstants,
 		groupCountX, groupCountY, groupCountZ);
 }