shader: store SPIR-V definitions directly in IR nodes

src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.h

@@ -9,6 +9,7 @@
 
 namespace Libs::Graphics::ShaderRecompiler::Spirv {
 
+// Resets and writes IR backend IDs; do not emit the same program concurrently.
 std::vector<uint32_t> EmitProgram(const IR::Program& program,
                                   ShaderStageInputInfo input_info);
 

src/graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h

@@ -160,7 +160,6 @@ struct EmitterState {
 	std::vector<InputBinding>  inputs;
 	std::vector<OutputBinding> outputs;
 	std::vector<uint32_t>      interface_variables;
-	std::unordered_map<const IR::Block*, uint32_t> labels;
 };
 
 uint32_t TypeVoid(EmitterState& state);
@@ -255,12 +254,10 @@ struct ValueEmitContext {
 	const IR::Inst*       ImageAddress(IR::Value value);
 	const IR::MemoryInfo& Memory(const IR::Inst& inst) const;
 	const IR::ExportInfo& Export(const IR::Inst& inst) const;
-	uint32_t              Label(const IR::Block* block) const;
 	[[noreturn]] void     Fail(const char* reason) const;
 	[[noreturn]] void     Fail(const IR::Inst& inst, const char* reason) const;
 
 	EmitterState&                                                      state;
-	std::unordered_map<const IR::Inst*, uint32_t>                      definitions;
 	const std::unordered_map<const IR::Inst*, uint32_t>*               dispatcher_spills = nullptr;
 	std::unordered_map<const IR::Inst*, std::pair<uint32_t, uint32_t>> dispatcher_block_loads;
 	uint32_t                                                           scratch_u32_variable = 0;

src/graphics/shader/recompiler/backend/spirv/spirvEmitterProgram.cpp

@@ -106,13 +106,13 @@ void EmitStructuredTerminator(ValueEmitContext& ctx, const IR::Block* block,
 			const auto* merge = TargetBlock(program, term.merge_block);
 			const auto* cont  = TargetBlock(program, term.continue_block);
 			if (merge != nullptr && cont != nullptr) {
-				ctx.state.builder.AddFunction(spv::OpLoopMerge, ctx.Label(merge), ctx.Label(cont),
+				ctx.state.builder.AddFunction(spv::OpLoopMerge, merge->Definition(),
-				                              spv::LoopControlMaskNone);
+				                              cont->Definition(), spv::LoopControlMaskNone);
 			}
 		} else if (term.kind == CFG::TerminatorKind::ConditionalBranch &&
 		           term.merge_block != UINT32_MAX) {
 			if (const auto* merge = TargetBlock(program, term.merge_block); merge != nullptr) {
-				ctx.state.builder.AddFunction(spv::OpSelectionMerge, ctx.Label(merge),
+				ctx.state.builder.AddFunction(spv::OpSelectionMerge, merge->Definition(),
 				                              spv::SelectionControlMaskNone);
 			}
 		}
@@ -126,7 +126,7 @@ void EmitStructuredTerminator(ValueEmitContext& ctx, const IR::Block* block,
 				return;
 			}
 			emit_merge();
-			ctx.state.builder.AddFunction(spv::OpBranch, ctx.Label(target));
+			ctx.state.builder.AddFunction(spv::OpBranch, target->Definition());
 			return;
 		}
 		case CFG::TerminatorKind::ConditionalBranch: {
@@ -139,7 +139,7 @@ void EmitStructuredTerminator(ValueEmitContext& ctx, const IR::Block* block,
 			const auto condition = ctx.Def(info.condition);
 			emit_merge();
 			ctx.state.builder.AddFunction(spv::OpBranchConditional, condition,
-			                              ctx.Label(true_block), ctx.Label(false_block));
+			                              true_block->Definition(), false_block->Definition());
 			return;
 		}
 		default: EmitReturn(ctx); return;
@@ -267,7 +267,7 @@ void EmitStructuredInstruction(ValueEmitContext& ctx, StructuredFunctionState& s
 		}
 		for (size_t index = 0; index < inst.NumArgs(); index++) {
 			const auto* predecessor = inst.PhiBlock(index);
-			if (predecessor == nullptr || !ctx.state.labels.contains(predecessor)) {
+			if (predecessor == nullptr || predecessor->Definition() == 0) {
 				ctx.Fail(inst, "has a predecessor outside the structured function");
 			}
 		}
@@ -301,7 +301,7 @@ void EmitDispatcherInstruction(ValueEmitContext& ctx, const DispatcherFunctionSt
 template <typename EmitInstruction>
 void EmitBlock(ValueEmitContext& ctx, const IR::Block* block, EmitInstruction&& emit_instruction) {
 	ctx.state.current_block = block;
-	EmitLabel(ctx.state, ctx.Label(block));
+	EmitLabel(ctx.state, block->Definition());
 	bool emitted_non_phi = false;
 	for (const auto& inst: *block) {
 		if (inst.GetOpcode() == IR::ValueOpcode::Phi) {
@@ -342,7 +342,7 @@ void PatchStructuredPhis(ValueEmitContext& ctx, StructuredFunctionState& structu
 void EmitStructuredFunction(ValueEmitContext& ctx) {
 	const auto& program = ctx.state.program;
 	StructuredFunctionState structured;
-	ctx.state.builder.AddFunction(spv::OpBranch, ctx.Label(program.blocks.front()));
+	ctx.state.builder.AddFunction(spv::OpBranch, program.blocks.front()->Definition());
 	for (size_t index = 0; index < program.blocks.size(); index++) {
 		const auto* block = program.blocks[index];
 		EmitBlock(ctx, block, [&](ValueEmitContext& lane, const IR::Inst& inst) {
@@ -357,7 +357,7 @@ void EmitStructuredFunction(ValueEmitContext& ctx) {
 void EmitDispatcherFunction(ValueEmitContext& ctx, const DispatcherFunctionState& dispatcher) {
 	auto&       state = ctx.state;
 	const auto* entry = state.program.blocks.front();
-	state.builder.AddFunction(spv::OpBranch, ctx.Label(entry));
+	state.builder.AddFunction(spv::OpBranch, entry->Definition());
 	EmitBlock(ctx, entry, [&](ValueEmitContext& lane, const IR::Inst& inst) {
 		EmitDispatcherInstruction(lane, dispatcher, inst);
 	});
@@ -385,7 +385,7 @@ void EmitDispatcherFunction(ValueEmitContext& ctx, const DispatcherFunctionState
 	std::vector<uint32_t> words {spv::OpSwitch, pc, dispatcher.after_switch_label};
 	for (size_t index = 1; index < state.program.blocks.size(); index++) {
 		words.push_back(state.program.block_info[index].id);
-		words.push_back(ctx.Label(state.program.blocks[index]));
+		words.push_back(state.program.blocks[index]->Definition());
 	}
 	state.builder.AddFunction(words);
 	std::vector<uint32_t> next_pc_words {spv::OpPhi, TypeU32(state), next_pc,
@@ -551,23 +551,23 @@ uint32_t ValueEmitContext::Shuffle(const IR::Inst& inst, size_t index, uint32_t
 }
 
 uint32_t ValueEmitContext::Result(const IR::Inst& inst) {
-	if (const auto found = definitions.find(&inst); found != definitions.end()) {
+	if (const auto id = inst.Definition(half); id != 0) {
-		return found->second;
+		return id;
 	}
 	const auto id = state.builder.AllocateId();
-	definitions.emplace(&inst, id);
+	inst.SetDefinition(half, id);
 	return id;
 }
 
 uint32_t ValueEmitContext::Define(const IR::Inst& inst, uint32_t value) {
-	if (const auto found = definitions.find(&inst); found != definitions.end()) {
+	if (const auto id = inst.Definition(half); id != 0) {
-		if (found->second != value) {
+		if (id != value) {
 			state.builder.AddFunction(spv::OpCopyObject, TypeId(state, inst.GetType()),
-			                          found->second, value);
+			                          id, value);
 		}
-		return found->second;
+		return id;
 	}
-	definitions.emplace(&inst, value);
+	inst.SetDefinition(half, value);
 	return value;
 }
 
@@ -595,10 +595,6 @@ const IR::ExportInfo& ValueEmitContext::Export(const IR::Inst& inst) const {
 	return state.program.export_info.at(inst.Flags<IR::ExportFlags>().index);
 }
 
-uint32_t ValueEmitContext::Label(const IR::Block* block) const {
-	return state.labels.at(block);
-}
-
 [[noreturn]] void ValueEmitContext::Fail(const char* reason) const {
 	EXIT("SPIR-V emission failed: hash=0x%016" PRIx64 " stage=%u reason=%s\n",
 	     state.program.shader_hash, static_cast<unsigned>(state.program.stage), reason);
@@ -627,8 +623,11 @@ void EmitProgram(EmitterState& state) {
 		state.builder.AddName(state.pixel_valid_mask_variable, "pixel_valid_mask_active");
 	}
 	for (const auto* block: program.blocks) {
-		const auto label = state.builder.AllocateId();
+		// Re-emission may use a different host subgroup width and ID allocation order.
-		state.labels.emplace(block, label);
+		for (const auto& inst: *block) {
+			inst.ResetDefinitions();
+		}
+		block->SetDefinition(state.builder.AllocateId());
 	}
 	if (state.program.dispatcher_fallback) {
 		auto& dispatch = dispatcher.emplace();

src/graphics/shader/recompiler/ir/Block.h

@@ -37,6 +37,8 @@ public:
 	[[nodiscard]] const InstructionList&  Instructions() const;
 	[[nodiscard]] std::span<Block* const> ImmPredecessors() const;
 	[[nodiscard]] std::span<Block* const> ImmSuccessors() const;
+	[[nodiscard]] uint32_t Definition() const { return definition; }
+	void SetDefinition(uint32_t value) const { definition = value; }
 
 	void               SsaSeal();
 	[[nodiscard]] bool IsSsaSealed() const;
@@ -57,6 +59,7 @@ private:
 	std::vector<Block*> predecessors;
 	std::vector<Block*> successors;
 	bool                ssa_sealed = false;
+	mutable uint32_t    definition = 0;
 };
 
 using BlockList = std::vector<Block*>;

src/graphics/shader/recompiler/ir/Value.h

@@ -129,6 +129,12 @@ public:
 		}
 		return evaluation_index;
 	}
+	// Backend definitions belong to one emission. Wave64 emulation needs one per half.
+	[[nodiscard]] uint32_t Definition(uint32_t half) const { return definitions[half]; }
+	void SetDefinition(uint32_t half, uint32_t definition) const {
+		definitions[half] = definition;
+	}
+	void ResetDefinitions() const { definitions = {}; }
 
 	void SetParent(Block* block);
 	void SetArg(size_t index, Value value);
@@ -165,6 +171,7 @@ private:
 	uint8_t             num_args;
 	bool                live = false;
 	mutable uint32_t    evaluation_index = UINT32_MAX;
+	mutable std::array<uint32_t, 2> definitions {};
 	uint64_t            flags;
 	Block*              parent = nullptr;
 	union {
@@ -175,6 +182,6 @@ private:
 	std::vector<Use>    uses;
 };
 
-static_assert(sizeof(Inst) <= 112, "Inst operand storage unintentionally increased");
+static_assert(sizeof(Inst) <= 120, "Inst operand and backend storage unintentionally increased");
 
 } // namespace Libs::Graphics::ShaderRecompiler::IR

tests/shaderCfgTests.cpp

@@ -14692,14 +14692,39 @@ void TestCompilerStageInputOwnership() {
 void TestSpirvEmissionOwnsRequirements() {
   using namespace ShaderRecompiler;
   const uint32_t shader[] = {EncodeSopp(0x01)};
-  const auto options = MakeCompileOptions(ShaderType::Compute);
+  ShaderComputeInputInfo compute{};
+  compute.threads_num[0] = 64;
+  compute.threads_num[1] = compute.threads_num[2] = 1;
+  auto options = MakeCompileOptions(ShaderType::Compute);
+  options.wave_size = 64;
+  options.input_info.compute = &compute;
   auto compiled = RecompileForTest(shader, options);
   IR::IREmitter emitter(compiled.program.blocks.front());
-  emitter.Emit(IR::ValueOpcode::LaneId, {});
+  const auto lane = emitter.Emit(IR::ValueOpcode::LaneId, {});
+  const auto value = emitter.Emit(IR::ValueOpcode::IAdd32, {lane, IR::Value(1u)});
+  const auto first = emitter.Emit(IR::ValueOpcode::ReadFirstLane,
+                                  {value, IR::Value(true)});
+  emitter.Emit(IR::ValueOpcode::IAdd32, {first, lane});
   const auto binary = Spirv::EmitProgram(compiled.program, options.input_info);
   CheckSpirvBinaryValidates(binary);
   Check((DisassembleSpirvBinary(binary).find("BuiltIn SubgroupLocalInvocationId") != std::string::npos),
         "emission reused requirements from an earlier IR version");
+  Check(Spirv::EmitProgram(compiled.program, options.input_info) == binary,
+        "repeated emission reused IDs from the previous module");
+
+  compute.host_subgroup_size = 32;
+  const auto split = Spirv::EmitProgram(compiled.program, options.input_info);
+  CheckSpirvBinaryValidates(split);
+  Check(split != binary &&
+            Spirv::EmitProgram(compiled.program, options.input_info) == split,
+        "split-wave emission did not reset both halves' definitions");
+
+  compute.host_subgroup_size = 64;
+  Check(Spirv::EmitProgram(compiled.program, options.input_info) == binary,
+        "native-wave re-emission retained split-wave definitions");
+  compute.host_subgroup_size = 32;
+  Check(Spirv::EmitProgram(compiled.program, options.input_info) == split,
+        "split-wave re-emission retained native-wave definitions");
 }
 
 void TestRepeatedExportsHaveOneInterface() {