shader: preserve descriptor key SSA across native-resolved image uses

Co-authored-by: TheCruZ <thecruzwtf@gmail.com>

src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp

@@ -371,10 +371,11 @@ public:
 			}
 		}
 		for (const auto& plan: m_indirect_descriptors) {
-			plan.handle->SetArg(0, plan.key);
 			if (plan.handle->GetOpcode() == ValueOpcode::GetBufferResource) {
+				plan.handle->SetArg(0, plan.key);
 				for (uint32_t word = 1; word < 4u; ++word) plan.handle->SetArg(word, Value(0u));
 			} else {
+				if (plan.reads[0] == nullptr) plan.handle->SetArg(0, plan.key);
 				for (uint32_t dword = 0; dword < 4u; dword++) {
 					plan.handle->SetArg(dword + 1u, plan.roots[dword + 4u]);
 				}
@@ -387,6 +388,20 @@ public:
 					m_program.memory_info[plan.memory[word]].planning_only = true;
 			}
 		}
+		// Project descriptor-only SSA onto its key. Existing Phis preserve dominance;
+		// native descriptor resolution excludes their stale incoming values at each use.
+		for (const auto& plan: m_indirect_descriptors) {
+			if (plan.handle->GetOpcode() != ValueOpcode::GetImageResource || plan.reads[0] == nullptr)
+				continue;
+			for (const auto* read: plan.reads) {
+				const auto first = std::ranges::find_if(m_indirect_descriptors, [&](const auto& candidate) {
+					return candidate.reads[0] == read;
+				});
+				const auto key = first != m_indirect_descriptors.end() ? first->key : Value(0u);
+				const auto uses = read->Uses();
+				for (const auto& use: uses) use.user->SetArg(use.operand, key);
+			}
+		}
 		m_program.descriptor_sources         = std::move(m_sources);
 		m_program.info                       = std::move(m_info);
 		m_program.resource_tracking_complete = true;
@@ -857,10 +872,17 @@ private:
 		return true;
 	}
 
-	static bool UsesOnly(const Inst& value, std::span<const Inst* const> users) {
+	static bool UsesOnlyImageDescriptors(const Inst& value) {
-		return !value.Uses().empty() && std::ranges::all_of(value.Uses(), [&](const Use& use) {
+		if (value.Uses().empty()) return false;
-			return std::ranges::find(users, use.user) != users.end();
+		std::vector<const Inst*> values {&value};
-		});
+		for (size_t index = 0; index < values.size(); ++index) {
+			for (const auto& use: values[index]->Uses()) {
+				if (use.user->GetOpcode() == ValueOpcode::GetImageResource) continue;
+				if (use.user->GetOpcode() != ValueOpcode::Phi) return false;
+				if (std::ranges::find(values, use.user) == values.end()) values.push_back(use.user);
+			}
+		}
+		return true;
 	}
 
 	const MemoryInfo* ScalarReadMemory(const Inst& read, uint32_t& index) const {
@@ -1497,11 +1519,12 @@ private:
 		return false;
 	}
 
-	bool MatchDescriptorTable(Inst& handle, IndirectDescriptorPlan& plan,
+	bool MatchDescriptorTable(Inst& handle, const DescriptorSource& descriptor,
+	                          IndirectDescriptorPlan& plan,
 	                          DescriptorSource& table_source, Value& key, uint32_t& table_offset) {
 		Inst* table_handle = nullptr;
 		for (uint32_t dword = 0; dword < handle.NumArgs(); ++dword) {
-			auto* read = handle.Arg(dword).Resolve().TryInstruction();
+			auto* read = descriptor.dwords[dword].Resolve().TryInstruction();
 			if (read == nullptr) {
 				return false;
 			}
@@ -1534,8 +1557,7 @@ private:
 				return false;
 			}
 			table_handle = current_handle;
-			const std::array<const Inst*, 1> image_users {&handle};
+			if (handle.NumArgs() == 8u ? !UsesOnlyImageDescriptors(*read) :
-			if (handle.NumArgs() == 8u ? !UsesOnly(*read, image_users) :
 			    std::ranges::any_of(read->Uses(), [](const Use& use) {
 				return use.user->GetOpcode() != ValueOpcode::GetBufferResource ||
 				       std::ranges::any_of(use.user->Uses(), [](const Use& consumer) {
@@ -1584,12 +1606,13 @@ private:
 		return true;
 	}
 
-	bool TryMakeIndirectImage(Inst& handle, IndirectDescriptorPlan& plan) {
+	bool TryMakeIndirectImage(Inst& handle, const DescriptorSource& descriptor,
+	                          IndirectDescriptorPlan& plan) {
 		if (handle.GetOpcode() != ValueOpcode::GetImageResource || handle.NumArgs() != 8u) return false;
 		DescriptorSource table_source;
 		Value key;
 		uint32_t table_offset = 0;
-		if (!MatchDescriptorTable(handle, plan, table_source, key, table_offset)) return false;
+		if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset)) return false;
 		DescriptorSource material_source;
 		DescriptorSource::IndirectDescriptor indirect;
 		indirect.table_offset = table_offset;
@@ -1628,14 +1651,6 @@ private:
 			const auto step = std::gcd<uint64_t>(indirect.selector_stride, uint64_t {1} << 32u);
 			indirect.selector_offset =
 			    (static_cast<uint32_t>(indirect.selector_offset % step) & ~3u) + (memory->offset & ~3u);
-			const auto* shift = plan.reads[0]->Arg(1).Resolve().TryInstruction();
-			const std::array<const Inst*, 1> material_users {key_shift != nullptr ? key_shift : shift};
-			const std::array<const Inst*, 1> key_users {shift};
-			if (!UsesOnly(*material_read, material_users) ||
-			    (key_shift != nullptr && !UsesOnly(*key_shift, key_users)) ||
-			    !UsesOnly(*shift, plan.reads)) {
-				return false;
-			}
 			const auto* material_handle = material_read->Arg(0).Resolve().TryInstruction();
 			if (material_handle == nullptr ||
 			    material_handle->GetOpcode() != ValueOpcode::GetBufferResource ||
@@ -1660,7 +1675,8 @@ private:
 		return true;
 	}
 
-	bool TryMakeIndirectBuffer(Inst& handle, IndirectDescriptorPlan& plan) {
+	bool TryMakeIndirectBuffer(Inst& handle, const DescriptorSource& descriptor,
+	                           IndirectDescriptorPlan& plan) {
 		if (handle.GetOpcode() != ValueOpcode::GetBufferResource || handle.NumArgs() != 4u ||
 		    handle.Uses().empty()) return false;
 		for (const auto& use: handle.Uses()) {
@@ -1675,7 +1691,7 @@ private:
 		DescriptorSource table_source;
 		Value key;
 		uint32_t table_offset = 0;
-		if (!MatchDescriptorTable(handle, plan, table_source, key, table_offset) ||
+		if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset) ||
 		    table_source.dword_count != 4u || table_offset != 0u) return false;
 		const auto* lane_read = key.Resolve().TryInstruction();
 		if (lane_read == nullptr || lane_read->GetOpcode() != ValueOpcode::ReadFirstLane)
@@ -1896,9 +1912,19 @@ private:
 				if (handle == nullptr || FindIndirectDescriptor(*handle) != nullptr) {
 					continue;
 				}
+				const auto flags = inst.Flags<MemoryFlags>();
+				const auto& memory = m_program.memory_info[flags.index];
+				DescriptorSource descriptor;
+				MakeSource(*handle, inst.GetOpcode() == ValueOpcode::StoreBufferU32 ? 4u : 8u,
+				           false, false, memory.resource * 4u, descriptor, flags.pc);
 				IndirectDescriptorPlan plan;
-				if (TryMakeIndirectImage(*handle, plan) || TryMakeFiniteImage(*handle, plan) ||
+				if (TryMakeIndirectImage(*handle, descriptor, plan) || TryMakeFiniteImage(*handle, plan) ||
-				    TryMakeIndirectBuffer(*handle, plan)) {
+				    TryMakeIndirectBuffer(*handle, descriptor, plan)) {
+					for (const auto& previous: m_indirect_descriptors) {
+						if (plan.reads[0] != nullptr && previous.reads[0] == plan.reads[0] &&
+						    !EquivalentValue(m_program, previous.key, plan.key))
+							Fail(flags.pc, "shared descriptor read has incompatible keys");
+					}
 					m_indirect_descriptors.push_back(std::move(plan));
 				}
 			}

tests/ShaderRecompilerComputeTests.cpp

@@ -35627,18 +35627,21 @@ void CheckImageSamplerSpecialization() {
         0u, 64u, false, static_cast<u32>(Prospero::BufferFormat::k32UInt));
     std::copy_n(output.begin(), 4, user_data.begin() + 48);
     user_data[48] = 0x3000u;
+    if (scalar) user_data[30] = 1u;
     std::vector<u32> code{EncodeVopc(0xc5, InlineU32(0), 0),
                           EncodeSop1(0x24, 24, 106), EncodeSopp(0x08, 0)};
-    size_t scalar_shift = 0;
+    const auto append_scalar_key = [&](bool numeric_use) {
-    if (scalar) {
       code.insert(code.end(), {EncodeVop1(0x02, 20, Vgpr(0)),
                                EncodeSop2(0x00, 20, 20, 21),
                                EncodeSop2(0x26, 20, 20, InlineU32(3)),
                                EncodeSop2(0x1e, 20, 20, InlineU32(2)),
                                EncodeSop2(0x00, 20, 20, InlineU32(8)),
                                EncodeSmem0(0x08, 20, 2), EncodeSmem1(0, 20)});
-      scalar_shift = code.size();
+      if (numeric_use) code.push_back(EncodeVop1(0x01, 17, 20));
       code.push_back(EncodeSop2(0x20, 20, 20, InlineU32(4)));
+    };
+    if (scalar) {
+      append_scalar_key(true);
     } else {
       code.insert(code.end(), {EncodeSop1(0x04, 26, 126), EncodeVop1(0x01, 16, InlineU32(1))});
       // An earlier descriptor waterfall separates the EXECZ witness from the material load.
@@ -35658,13 +35661,41 @@ void CheckImageSamplerSpecialization() {
       code.insert(code.end(), {EncodeSop1(0x04, 28, 126),
                                EncodeVop1(0x02, 20, Vgpr(15)), EncodeVopc(0xd2, 20, 15)});
     code.insert(code.end(), {EncodeSop2(0x1e, 20, 20, InlineU32(5)),
-                             EncodeSmem0(0x0b, 12), EncodeSmem1(0, 20),
+                             EncodeSmem0(0x0b, 12), EncodeSmem1(0, 20)});
-                             EncodeMimg0(0x24, 1), EncodeMimg1(4, 0, 3, 2)});
+    const auto table_read_end = code.size();
+    code.insert(code.end(), {EncodeMimg0(0x24, 1), EncodeMimg1(4, 0, 3, 2)});
     AppendBufferStoreDword(&code, 4, 0);
     if (!scalar) {
       code.insert(code.end(), {EncodeVop1(0x01, 16, InlineU32(0)),
                                EncodeSop1(0x04, 126, 28), EncodeVopc(0xd5, InlineU32(0), 16)});
       code.push_back(EncodeSopp(0x09, static_cast<int16_t>(image_loop - code.size() - 1u)));
+    } else {
+      // Crossed exits introduce stale s16:19 Phi arms before the later s16:23 samples.
+      code.push_back(EncodeVop1(0x01, 5, InlineU32(0)));
+      code.push_back(EncodeSopc(0x06, 30, InlineU32(0)));
+      const auto skip_load = code.size();
+      code.push_back(EncodeSopp(0x05, 0));
+      append_scalar_key(false);
+      code.insert(code.end(), {EncodeSop2(0x1e, 20, 20, InlineU32(5)),
+                               EncodeSmem0(0x0b, 16), EncodeSmem1(0, 20),
+                               EncodeSopc(0x06, 31, InlineU32(0))});
+      const auto later_sample = code.size();
+      code.push_back(EncodeSopp(0x05, 0));
+      code.insert(code.end(), {EncodeMimg0(0x24, 1), EncodeMimg1(4, 0, 4, 2)});
+      const auto early_exit = code.size();
+      code.push_back(EncodeSopp(0x02, 0));
+      code[later_sample] = EncodeSopp(0x05, static_cast<int16_t>(code.size() - later_sample - 1u));
+      code.insert(code.end(), {EncodeMimg0(0x24, 1), EncodeMimg1(4, 0, 4, 2),
+                               EncodeMimg0(0x24, 1), EncodeMimg1(5, 0, 4, 2)});
+      const auto leave_samples = code.size();
+      code.push_back(EncodeSopp(0x02, 0));
+      code[skip_load] = EncodeSopp(0x05, static_cast<int16_t>(code.size() - skip_load - 1u));
+      code.push_back(EncodeVop1(0x01, 4, InlineU32(0)));
+      code[early_exit] = EncodeSopp(0x02, static_cast<int16_t>(code.size() - early_exit - 1u));
+      code[leave_samples] = EncodeSopp(0x02, static_cast<int16_t>(code.size() - leave_samples - 1u));
+      AppendBufferStoreDword(&code, 4, 0);
+      AppendBufferStoreDword(&code, 5, 0);
+      AppendBufferStoreDword(&code, 17, 0);
     }
     code[2] = EncodeSopp(0x08, static_cast<int16_t>(code.size() - 3u));
     AppendEnd(&code);
@@ -35678,6 +35709,15 @@ void CheckImageSamplerSpecialization() {
     options.user_data = user_data;
     options.input_info.compute = &compute;
     auto translated = ShaderRecompiler::TranslateProgram(code, options);
+    if (scalar) {
+      bool key_phi = false;
+      for (const auto *block : translated.program.blocks)
+        for (const auto &inst : *block)
+          key_phi |= inst.GetOpcode() == ValueOpcode::GetImageResource &&
+                     inst.Arg(0).Resolve().IsPhi();
+      Require(name, "descriptor key reconvergence", key_phi,
+              "native crossed exits did not retain a descriptor key Phi");
+    }
     auto plan = ExtractResourcePlan(translated.program);
     ResourceSnapshot snapshot;
     ResourceSpecialization specialization;
@@ -35724,10 +35764,10 @@ void CheckImageSamplerSpecialization() {
             "GPU-selected texture materialization bypassed strict read provenance");
 #if KYTY_PLATFORM != KYTY_PLATFORM_WINDOWS
     if (scalar) {
-      ExpectFatal("ScalarMaterialNumericConsumer", [&] {
+      ExpectFatal("ScalarImageDescriptorNumericConsumer", [&] {
-        code.insert(code.begin() + scalar_shift, EncodeVop1(0x01, 17, 20));
+        code.insert(code.begin() + table_read_end, EncodeVop1(0x01, 18, 12));
         code.pop_back();
-        AppendBufferStoreDword(&code, 17, 0);
+        AppendBufferStoreDword(&code, 18, 0);
         code[2] = EncodeSopp(0x08, static_cast<int16_t>(code.size() - 3u));
         AppendEnd(&code);
         (void)ShaderRecompiler::TranslateProgram(code, options);