shader: drop quadratic scans from resource tracking

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

@@ -16,6 +16,7 @@
 #include <span>
 #include <tuple>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #define XXH_INLINE_ALL
 #include <xxhash.h>
@@ -563,6 +564,11 @@ private:
 		bool     has_sampler = false;
 	};
 
+	struct MemoryIndexUse {
+		const Inst* owner = nullptr;
+		uint32_t    count = 0;
+	};
+
 	struct ResolvedHandle {
 		const Inst* handle;
 		uint32_t pc;
@@ -829,6 +835,7 @@ private:
 		if (value.IsImmediate()) return;
 		auto* inst = value.TryInstruction();
 		if (inst == nullptr) Fail(use_pc, "invalid typed planning value");
+		if (m_srt_visited.contains(inst)) return;
 		const auto cycle = std::ranges::find(m_srt_visiting, inst);
 		if (cycle != m_srt_visiting.end()) {
 			if (std::any_of(cycle, m_srt_visiting.end(), [](const Inst* value) {
@@ -836,7 +843,6 @@ private:
 			})) return;
 			Fail(use_pc, "cyclic typed planning value without a phi");
 		}
-		if (std::ranges::find(m_srt_visited, inst) != m_srt_visited.end()) return;
 		m_srt_visiting.push_back(inst);
 		uint32_t memory_index = 0;
 		const auto* memory = ScalarReadMemory(*inst, memory_index);
@@ -858,7 +864,7 @@ private:
 				CollectScalarRead(inst->Arg(arg), use_pc);
 		}
 		m_srt_visiting.pop_back();
-		m_srt_visited.push_back(inst);
+		m_srt_visited.insert(inst);
 		if (memory == nullptr) return;
 		const auto offset = inst->Arg(1).Resolve();
 		if (!offset.IsImmediate() || offset.GetType() != Type::U32) return;
@@ -1098,22 +1104,25 @@ private:
 		           : nullptr;
 	}
 
-	bool MemoryIndexBelongsTo(uint32_t index, const Inst& owner) const {
+	bool MemoryIndexBelongsTo(uint32_t index, const Inst& owner) {
-		for (const auto* block: m_program.blocks) {
+		if (m_memory_index_uses.empty()) {
-			for (const auto& inst: *block) {
+			// Only descriptor planning asks, and it adds no memory instructions.
-				const auto op = inst.GetOpcode();
+			m_memory_index_uses.resize(m_program.memory_info.size());
-				if ((BufferAccessOf(op) == BufferAccess::None &&
+			for (const auto* block: m_program.blocks) {
-				     AddressOpcodeInfoOf(op).access == AddressAccess::None &&
+				for (const auto& inst: *block) {
-				     ImageOpcodeInfoOf(op).access == ImageAccess::None) ||
+					const auto op = inst.GetOpcode();
-				    &inst == &owner) {
+					if (BufferAccessOf(op) != BufferAccess::None ||
-					continue;
+					    AddressOpcodeInfoOf(op).access != AddressAccess::None ||
-				}
+					    ImageOpcodeInfoOf(op).access != ImageAccess::None) {
-				if (inst.Flags<MemoryFlags>().index == index) {
+						auto& use = m_memory_index_uses[inst.Flags<MemoryFlags>().index];
-					return false;
+						use.owner = &inst;
+						use.count++;
+					}
 				}
 			}
 		}
-		return true;
+		const auto& use = m_memory_index_uses[index];
+		return use.count == 0 || (use.count == 1 && use.owner == &owner);
 	}
 
 	bool MakeRuntimeTableSource(const Inst& read, DescriptorSource& descriptor) {
@@ -4342,7 +4351,8 @@ private:
 	std::vector<Program::ScalarWrite>          m_scalar_writes;
 	std::vector<ResolvedHandle>                m_resolved_handles;
 	std::vector<const Inst*>                   m_srt_visiting;
-	std::vector<const Inst*>                   m_srt_visited;
+	std::unordered_set<const Inst*>            m_srt_visited;
+	std::vector<MemoryIndexUse>                m_memory_index_uses;
 	std::vector<Inst*>                         m_scalar_reads;
 	ShaderInfo                                 m_info;
 	std::vector<DescriptorSource>              m_sources;

tests/ResourceTrackingTests.cpp

@@ -439,6 +439,16 @@ void TestInvariantIndirectImageMaterialization(bool nested) {
              "noncontiguous descriptor DWORDs entered the table plan");
   Check(!malformed->program.resource_tracking_complete && malformed->program.descriptor_sources.empty(),
         "malformed image table tracking was not transactional");
+  auto aliased = MakeIndirectImageFixture(false, 48u, 0u, 16u, nested);
+  const auto image = std::ranges::find_if(*aliased->block, [](const Inst &inst) {
+    return inst.GetOpcode() == ValueOpcode::GetImageResource;
+  });
+  const auto *read = image->Arg(0).Resolve().TryInstruction();
+  aliased->Emit(ValueOpcode::ReferenceU32,
+                {aliased->Emit(ValueOpcode::ReadConstBuffer, {read->Arg(0), read->Arg(1)},
+                               read->Flags<uint64_t>())});
+  CheckFatal([&] { aliased->PlanAndTrack(); }, "not a valid runtime value",
+             "a descriptor read sharing its memory index entered the table plan");
 }
 
 void TestResidentImageHeap() {