SpirvShader: Split instructions into blocks.

src/Pipeline/SpirvShader.cpp

@@ -36,12 +36,8 @@ namespace sw
 		// - There is exactly one entrypoint in the module, and it's the one we want
 		// - The only input/output OpVariables present are those used by the entrypoint
 
-		// TODO: Add real support for control flow. For now, track whether we've seen
-		// a label or a return already (if so, the shader does things we will mishandle).
-		// We expect there to be one of each in a simple shader -- the first and last instruction
-		// of the entrypoint function.
-		bool seenLabel = false;
-		bool seenReturn = false;
+		Block::ID currentBlock;
+		InsnIterator blockStart;
 
 		for (auto insn : *this)
 		{
@@ -114,16 +110,35 @@ namespace sw
 			}
 
 			case spv::OpLabel:
-				if (seenLabel)
-					UNIMPLEMENTED("Shader contains multiple labels, has control flow");
-				seenLabel = true;
+			{
+				ASSERT(currentBlock.value() == 0);
+				currentBlock = Block::ID(insn.word(1));
+				blockStart = insn;
 				break;
+			}
 
+			// Branch Instructions (subset of Termination Instructions):
+			case spv::OpBranch:
+			case spv::OpBranchConditional:
+			case spv::OpSwitch:
 			case spv::OpReturn:
-				if (seenReturn)
-					UNIMPLEMENTED("Shader contains multiple returns, has control flow");
-				seenReturn = true;
+			// fallthrough
+
+			// Termination instruction:
+			case spv::OpKill:
+			case spv::OpUnreachable:
+			{
+				ASSERT(currentBlock.value() != 0);
+				auto blockEnd = insn; blockEnd++;
+				blocks[currentBlock] = Block(blockStart, blockEnd);
+				currentBlock = Block::ID(0);
+
+				if (insn.opcode() == spv::OpKill)
+				{
+					modes.ContainsKill = true;
+				}
 				break;
+			}
 
 			case spv::OpTypeVoid:
 			case spv::OpTypeBool:
@@ -225,11 +240,31 @@ namespace sw
 			}
 
 			case spv::OpCapability:
-				// Various capabilities will be declared, but none affect our code generation at this point.
+				break; // Various capabilities will be declared, but none affect our code generation at this point.
 			case spv::OpMemoryModel:
-				// Memory model does not affect our code generation until we decide to do Vulkan Memory Model support.
+				break; // Memory model does not affect our code generation until we decide to do Vulkan Memory Model support.
+
 			case spv::OpEntryPoint:
+				break;
 			case spv::OpFunction:
+				ASSERT(mainBlockId.value() == 0); // Multiple functions found
+				// Scan forward to find the function's label.
+				for (auto it = insn; it != end() && mainBlockId.value() == 0; it++)
+				{
+					switch (it.opcode())
+					{
+					case spv::OpFunction:
+					case spv::OpFunctionParameter:
+						break;
+					case spv::OpLabel:
+						mainBlockId = Block::ID(it.word(1));
+						break;
+					default:
+						ERR("Unexpected opcode '%s' following OpFunction", OpcodeName(it.opcode()).c_str());
+					}
+				}
+				ASSERT(mainBlockId.value() != 0); // Function's OpLabel not found
+				break;
 			case spv::OpFunctionEnd:
 				// Due to preprocessing, the entrypoint and its function provide no value.
 				break;
@@ -363,10 +398,6 @@ namespace sw
 				// Don't need to do anything during analysis pass
 				break;
 
-			case spv::OpKill:
-				modes.ContainsKill = true;
-				break;
-
 			default:
 				UNIMPLEMENTED(OpcodeName(insn.opcode()).c_str());
 			}
@@ -530,7 +561,7 @@ namespace sw
 		}
 	}
 
-	uint32_t SpirvShader::ComputeTypeSize(sw::SpirvShader::InsnIterator insn)
+	uint32_t SpirvShader::ComputeTypeSize(InsnIterator insn)
 	{
 		// Types are always built from the bottom up (with the exception of forward ptrs, which
 		// don't appear in Vulkan shaders. Therefore, we can always assume our component parts have
@@ -959,8 +990,31 @@ namespace sw
 
 	void SpirvShader::emit(SpirvRoutine *routine) const
 	{
+		// Emit everything up to the first label
+		// TODO: Separate out dispatch of block from non-block instructions?
 		for (auto insn : *this)
 		{
+			if (insn.opcode() == spv::OpLabel)
+			{
+				break;
+			}
+			EmitInstruction(routine, insn);
+		}
+
+		// Emit the main function block
+		EmitBlock(routine, getBlock(mainBlockId));
+	}
+
+	void SpirvShader::EmitBlock(SpirvRoutine *routine, Block const &block) const
+	{
+		for (auto insn : block)
+		{
+			EmitInstruction(routine, insn);
+		}
+	}
+
+	void SpirvShader::EmitInstruction(SpirvRoutine *routine, InsnIterator insn) const
+	{
 		switch (insn.opcode())
 		{
 		case spv::OpTypeVoid:
@@ -1136,7 +1190,6 @@ namespace sw
 			break;
 		}
 	}
-	}
 
 	void SpirvShader::EmitVariable(InsnIterator insn, SpirvRoutine *routine) const
 	{

src/Pipeline/SpirvShader.hpp

@@ -234,6 +234,26 @@ namespace sw
 			} kind = Kind::Unknown;
 		};
 
+		// Block is an interval of SPIR-V instructions, starting with the
+		// opening OpLabel, and ending with a termination instruction.
+		class Block
+		{
+		public:
+			using ID = SpirvID<Block>;
+
+			Block() = default;
+			Block(const Block& other) = default;
+			explicit Block(InsnIterator begin, InsnIterator end) : begin_(begin), end_(end) {}
+
+			/* range-based-for interface */
+			inline InsnIterator begin() const { return begin_; }
+			inline InsnIterator end() const { return end_; }
+
+		private:
+			InsnIterator begin_;
+			InsnIterator end_;
+		};
+
 		struct TypeOrObject {}; // Dummy struct to represent a Type or Object.
 
 		// TypeOrObjectID is an identifier that represents a Type or an Object,
@@ -382,12 +402,24 @@ namespace sw
 			return it->second;
 		}
 
+		Block const &getBlock(Block::ID id) const
+		{
+			auto it = blocks.find(id);
+			ASSERT(it != blocks.end());
+			return it->second;
+		}
+
 	private:
 		const int serialID;
 		static volatile int serialCounter;
 		Modes modes;
 		HandleMap<Type> types;
 		HandleMap<Object> defs;
+		HandleMap<Block> blocks;
+		Block::ID mainBlockId; // Block of the entry point function.
+
+		void EmitBlock(SpirvRoutine *routine, Block const &block) const;
+		void EmitInstruction(SpirvRoutine *routine, InsnIterator insn) const;
 
 		// DeclareType creates a Type for the given OpTypeX instruction, storing
 		// it into the types map. It is called from the analysis pass (constructor).