Propagate descriptor decorations to access-chain and load results

Split descriptor decorations from other decorations since they're always
used separately and this makes propagating them more lightweight.

Bug b/129523279

Change-Id: I2f5acecc990bf15c7feb4ce81539bd2b5818bdf6
Reviewed-on: https://swiftshader-review.googlesource.com/c/SwiftShader/+/28445
Presubmit-Ready: Nicolas Capens <nicolascapens@google.com>
Kokoro-Presubmit: kokoro <noreply+kokoro@google.com>
Tested-by: Nicolas Capens <nicolas.capens@gmail.com>
Reviewed-by: Chris Forbes <chrisforbes@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
diff --git a/src/Pipeline/SpirvShader.cpp b/src/Pipeline/SpirvShader.cpp
index b7ba738..4aed978 100644
--- a/src/Pipeline/SpirvShader.cpp
+++ b/src/Pipeline/SpirvShader.cpp
@@ -12,9 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <spirv/unified1/spirv.hpp>
-#include <spirv/unified1/GLSL.std.450.h>
 #include "SpirvShader.hpp"
+
 #include "System/Math.hpp"
 #include "Vulkan/VkBuffer.hpp"
 #include "Vulkan/VkDebug.hpp"
@@ -22,6 +21,9 @@
 #include "Vulkan/VkPipelineLayout.hpp"
 #include "Device/Config.hpp"
 
+#include <spirv/unified1/spirv.hpp>
+#include <spirv/unified1/GLSL.std.450.h>
+
 #ifdef Bool
 #undef Bool // b/127920555
 #endif
@@ -173,9 +175,22 @@
 			{
 				TypeOrObjectID targetId = insn.word(1);
 				auto decoration = static_cast<spv::Decoration>(insn.word(2));
-				decorations[targetId].Apply(
-						decoration,
-						insn.wordCount() > 3 ? insn.word(3) : 0);
+				uint32_t value = insn.wordCount() > 3 ? insn.word(3) : 0;
+
+				decorations[targetId].Apply(decoration, value);
+
+				switch(decoration)
+				{
+				case spv::DecorationDescriptorSet:
+					descriptorDecorations[targetId].DescriptorSet = value;
+					break;
+				case spv::DecorationBinding:
+					descriptorDecorations[targetId].Binding = value;
+					break;
+				default:
+					// Only handling descriptor decorations here.
+					break;
+				}
 
 				if (decoration == spv::DecorationCentroid)
 					modes.NeedsCentroid = true;
@@ -186,13 +201,14 @@
 			{
 				Type::ID targetId = insn.word(1);
 				auto memberIndex = insn.word(2);
+				auto decoration = static_cast<spv::Decoration>(insn.word(3));
+				uint32_t value = insn.wordCount() > 4 ? insn.word(4) : 0;
+
 				auto &d = memberDecorations[targetId];
 				if (memberIndex >= d.size())
 					d.resize(memberIndex + 1);    // on demand; exact size would require another pass...
-				auto decoration = static_cast<spv::Decoration>(insn.word(3));
-				d[memberIndex].Apply(
-						decoration,
-						insn.wordCount() > 4 ? insn.word(4) : 0);
+				
+				d[memberIndex].Apply(decoration, value);
 
 				if (decoration == spv::DecorationCentroid)
 					modes.NeedsCentroid = true;
@@ -207,12 +223,17 @@
 
 			case spv::OpGroupDecorate:
 			{
-				auto const &srcDecorations = decorations[insn.word(1)];
+				uint32_t group = insn.word(1);
+				auto const &groupDecorations = decorations[group];
+				auto const &descriptorGroupDecorations = descriptorDecorations[group];
 				for (auto i = 2u; i < insn.wordCount(); i++)
 				{
-					// remaining operands are targets to apply the group to.
-					decorations[insn.word(i)].Apply(srcDecorations);
+					// Remaining operands are targets to apply the group to.
+					uint32_t target = insn.word(i);
+					decorations[target].Apply(groupDecorations);
+					descriptorDecorations[target].Apply(descriptorGroupDecorations);
 				}
+
 				break;
 			}
 
@@ -456,6 +477,21 @@
 			case spv::OpLoad:
 			case spv::OpAccessChain:
 			case spv::OpInBoundsAccessChain:
+				{
+					// Propagate the descriptor decorations to the result.
+					Object::ID resultId = insn.word(2);
+					Object::ID pointerId = insn.word(3);
+					const auto &d = descriptorDecorations.find(pointerId);
+
+					if(d != descriptorDecorations.end())
+					{
+						descriptorDecorations[resultId] = d->second;
+					}
+
+					DefineResult(insn);
+				}
+				break;
+
 			case spv::OpCompositeConstruct:
 			case spv::OpCompositeInsert:
 			case spv::OpCompositeExtract:
@@ -469,7 +505,8 @@
 			case spv::OpTranspose:
 			case spv::OpVectorExtractDynamic:
 			case spv::OpVectorInsertDynamic:
-			case spv::OpNot: // Unary ops
+			// Unary ops
+			case spv::OpNot:
 			case spv::OpBitFieldInsert:
 			case spv::OpBitFieldSExtract:
 			case spv::OpBitFieldUExtract:
@@ -478,7 +515,8 @@
 			case spv::OpSNegate:
 			case spv::OpFNegate:
 			case spv::OpLogicalNot:
-			case spv::OpIAdd: // Binary ops
+			// Binary ops
+			case spv::OpIAdd:
 			case spv::OpISub:
 			case spv::OpIMul:
 			case spv::OpSDiv:
@@ -934,10 +972,8 @@
 
 			case Object::Kind::DescriptorSet:
 			{
-				Decorations d = {};
-				ApplyDecorationsForId(&d, id);
-
-				ASSERT(d.DescriptorSet >= 0);
+				const auto &d = descriptorDecorations.at(id);
+				ASSERT(d.DescriptorSet >= 0 && d.DescriptorSet < vk::MAX_BOUND_DESCRIPTOR_SETS);
 				ASSERT(d.Binding >= 0);
 
 				auto set = routine->getPointer(id);
@@ -965,11 +1001,11 @@
 		}
 	}
 
-	SIMD::Pointer SpirvShader::WalkExplicitLayoutAccessChain(Object::ID id, uint32_t numIndexes, uint32_t const *indexIds, SpirvRoutine *routine) const
+	SIMD::Pointer SpirvShader::WalkExplicitLayoutAccessChain(Object::ID baseId, uint32_t numIndexes, uint32_t const *indexIds, SpirvRoutine *routine) const
 	{
 		// Produce a offset into external memory in sizeof(float) units
 
-		auto &baseObject = getObject(id);
+		auto &baseObject = getObject(baseId);
 		Type::ID typeId = getType(baseObject.type).element;
 		Decorations d = {};
 		ApplyDecorationsForId(&d, baseObject.type);
@@ -989,7 +1025,7 @@
 			}
 		}
 
-		auto ptr = GetPointerToData(id, arrayIndex, routine);
+		auto ptr = GetPointerToData(baseId, arrayIndex, routine);
 
 		int constantOffset = 0;
 
@@ -1054,21 +1090,21 @@
 		return ptr;
 	}
 
-	SIMD::Int SpirvShader::WalkAccessChain(Object::ID id, uint32_t numIndexes, uint32_t const *indexIds, SpirvRoutine *routine) const
+	SIMD::Int SpirvShader::WalkAccessChain(Object::ID baseId, uint32_t numIndexes, uint32_t const *indexIds, SpirvRoutine *routine) const
 	{
 		// TODO: avoid doing per-lane work in some cases if we can?
 		// Produce a *component* offset into location-oriented memory
 
 		int constantOffset = 0;
 		SIMD::Int dynamicOffset = SIMD::Int(0);
-		auto &baseObject = getObject(id);
+		auto &baseObject = getObject(baseId);
 		Type::ID typeId = getType(baseObject.type).element;
 
 		// The <base> operand is a divergent pointer itself.
 		// Start with its offset and build from there.
 		if (baseObject.kind == Object::Kind::DivergentPointer)
 		{
-			dynamicOffset += routine->getIntermediate(id).Int(0);
+			dynamicOffset += routine->getIntermediate(baseId).Int(0);
 		}
 
 		for (auto i = 0u; i < numIndexes; i++)
@@ -1166,14 +1202,6 @@
 			HasComponent = true;
 			Component = arg;
 			break;
-		case spv::DecorationDescriptorSet:
-			HasDescriptorSet = true;
-			DescriptorSet = arg;
-			break;
-		case spv::DecorationBinding:
-			HasBinding = true;
-			Binding = arg;
-			break;
 		case spv::DecorationBuiltIn:
 			HasBuiltIn = true;
 			BuiltIn = static_cast<spv::BuiltIn>(arg);
@@ -1235,18 +1263,6 @@
 			Component = src.Component;
 		}
 
-		if (src.HasDescriptorSet)
-		{
-			HasDescriptorSet = true;
-			DescriptorSet = src.DescriptorSet;
-		}
-
-		if (src.HasBinding)
-		{
-			HasBinding = true;
-			Binding = src.Binding;
-		}
-
 		if (src.HasOffset)
 		{
 			HasOffset = true;
@@ -1273,6 +1289,19 @@
 		RelaxedPrecision |= src.RelaxedPrecision;
 	}
 
+	void SpirvShader::DescriptorDecorations::Apply(const sw::SpirvShader::DescriptorDecorations &src)
+	{
+		if(src.DescriptorSet >= 0)
+		{
+			DescriptorSet = src.DescriptorSet;
+		}
+
+		if(src.Binding >= 0)
+		{
+			Binding = src.Binding;
+		}
+	}
+
 	void SpirvShader::ApplyDecorationsForId(Decorations *d, TypeOrObjectID id) const
 	{
 		auto it = decorations.find(id);
@@ -1289,6 +1318,17 @@
 		}
 	}
 
+	void SpirvShader::DefineResult(const InsnIterator &insn)
+	{
+		Type::ID typeId = insn.word(1);
+		Object::ID resultId = insn.word(2);
+		auto &object = defs[resultId];
+		object.type = typeId;
+		object.kind = (getType(typeId).opcode() == spv::OpTypePointer)
+			? Object::Kind::DivergentPointer : Object::Kind::Intermediate;
+		object.definition = insn;
+	}
+
 	uint32_t SpirvShader::GetConstantInt(Object::ID id) const
 	{
 		// Slightly hackish access to constants very early in translation.
@@ -1644,7 +1684,9 @@
 
 	SpirvShader::EmitResult SpirvShader::EmitInstruction(InsnIterator insn, EmitState *state) const
 	{
-		switch (insn.opcode())
+		auto opcode = insn.opcode();
+
+		switch (opcode)
 		{
 		case spv::OpTypeVoid:
 		case spv::OpTypeInt:
@@ -1862,7 +1904,7 @@
 			return EmitKill(insn, state);
 
 		default:
-			UNIMPLEMENTED("opcode: %s", OpcodeName(insn.opcode()).c_str());
+			UNIMPLEMENTED("opcode: %s", OpcodeName(opcode).c_str());
 			break;
 		}
 
@@ -1903,9 +1945,9 @@
 		case spv::StorageClassUniform:
 		case spv::StorageClassStorageBuffer:
 		{
-			Decorations d{};
-			ApplyDecorationsForId(&d, resultId);
-			ASSERT(d.DescriptorSet >= 0);
+			const auto &d = descriptorDecorations.at(resultId);
+			ASSERT(d.DescriptorSet >= 0 && d.DescriptorSet < vk::MAX_BOUND_DESCRIPTOR_SETS);
+
 			routine->createPointer(resultId, routine->descriptorSets[d.DescriptorSet]);
 			break;
 		}
@@ -1915,6 +1957,7 @@
 			break;
 		}
 		default:
+			UNIMPLEMENTED("Storage class %d", objectTy.storageClass);
 			break;
 		}
 
@@ -1933,6 +1976,10 @@
 		auto &pointerTy = getType(pointer.type);
 		std::memory_order memoryOrder = std::memory_order_relaxed;
 
+		ASSERT(getType(pointer.type).element == result.type);
+		ASSERT(Type::ID(insn.word(1)) == result.type);
+		ASSERT(!atomic || getType(getType(pointer.type).element).opcode() == spv::OpTypeInt);  // Vulkan 1.1: "Atomic instructions must declare a scalar 32-bit integer type, for the value pointed to by Pointer."
+
 		if(atomic)
 		{
 			Object::ID semanticsId = insn.word(5);
@@ -1940,10 +1987,6 @@
 			memoryOrder = MemoryOrder(memorySemantics);
 		}
 
-		ASSERT(getType(pointer.type).element == result.type);
-		ASSERT(Type::ID(insn.word(1)) == result.type);
-		ASSERT(!atomic || getType(getType(pointer.type).element).opcode() == spv::OpTypeInt);  // Vulkan 1.1: "Atomic instructions must declare a scalar 32-bit integer type, for the value pointed to by Pointer."
-
 		if (pointerTy.storageClass == spv::StorageClassImage)
 		{
 			UNIMPLEMENTED("StorageClassImage load not yet implemented");
diff --git a/src/Pipeline/SpirvShader.hpp b/src/Pipeline/SpirvShader.hpp
index f3e3919..28709da 100644
--- a/src/Pipeline/SpirvShader.hpp
+++ b/src/Pipeline/SpirvShader.hpp
@@ -282,8 +282,9 @@
 				// A pointer to a vk::DescriptorSet*.
 				// Pointer held by SpirvRoutine::pointers.
 				DescriptorSet,
+			};
 
-			} kind = Kind::Unknown;
+			Kind kind = Kind::Unknown;
 		};
 
 		// Block is an interval of SPIR-V instructions, starting with the
@@ -401,38 +402,36 @@
 
 		struct Decorations
 		{
-			int32_t Location;
-			int32_t Component;
-			int32_t DescriptorSet;
-			int32_t Binding;
-			spv::BuiltIn BuiltIn;
-			int32_t Offset;
-			int32_t ArrayStride;
-			int32_t MatrixStride;
+			int32_t Location = -1;
+			int32_t Component = 0;
+			spv::BuiltIn BuiltIn = static_cast<spv::BuiltIn>(-1);
+			int32_t Offset = -1;
+			int32_t ArrayStride = -1;
+			int32_t MatrixStride = 1;
+
 			bool HasLocation : 1;
 			bool HasComponent : 1;
-			bool HasDescriptorSet : 1;
-			bool HasBinding : 1;
 			bool HasBuiltIn : 1;
+			bool HasOffset : 1;
+			bool HasArrayStride : 1;
+			bool HasMatrixStride : 1;
+
 			bool Flat : 1;
 			bool Centroid : 1;
 			bool NoPerspective : 1;
 			bool Block : 1;
 			bool BufferBlock : 1;
-			bool HasOffset : 1;
-			bool HasArrayStride : 1;
-			bool HasMatrixStride : 1;
 			bool RelaxedPrecision : 1;
 
 			Decorations()
-					: Location{-1}, Component{0}, DescriptorSet{-1}, Binding{-1},
+					: Location{-1}, Component{0},
 					  BuiltIn{static_cast<spv::BuiltIn>(-1)},
 					  Offset{-1}, ArrayStride{-1}, MatrixStride{-1},
 					  HasLocation{false}, HasComponent{false},
-					  HasDescriptorSet{false}, HasBinding{false},
-					  HasBuiltIn{false}, Flat{false}, Centroid{false},
-					  NoPerspective{false}, Block{false}, BufferBlock{false},
-					  HasOffset{false}, HasArrayStride{false}, HasMatrixStride{false},
+					  HasBuiltIn{false}, HasOffset{false},
+					  HasArrayStride{false}, HasMatrixStride{false},
+					  Flat{false}, Centroid{false}, NoPerspective{false},
+					  Block{false}, BufferBlock{false},
 					  RelaxedPrecision{false}
 			{
 			}
@@ -447,6 +446,16 @@
 		std::unordered_map<TypeOrObjectID, Decorations, TypeOrObjectID::Hash> decorations;
 		std::unordered_map<Type::ID, std::vector<Decorations>> memberDecorations;
 
+		struct DescriptorDecorations
+		{
+			int32_t DescriptorSet = -1;
+			int32_t Binding = -1;
+
+			void Apply(DescriptorDecorations const &src); 
+		};
+
+		std::unordered_map<Object::ID, DescriptorDecorations> descriptorDecorations;
+
 		struct InterfaceComponent
 		{
 			AttribType Type;
@@ -525,6 +534,9 @@
 		void ApplyDecorationsForId(Decorations *d, TypeOrObjectID id) const;
 		void ApplyDecorationsForIdMember(Decorations *d, Type::ID id, uint32_t member) const;
 
+		// Creates an Object for the instruction's result in 'defs'.
+		void DefineResult(const InsnIterator &insn);
+
 		// Returns true if data in the given storage class is word-interleaved
 		// by each SIMD vector lane, otherwise data is stored linerally.
 		//