Reuse scalar evaluation storage

Co-authored-by: MehmetCambaz <mcambaz18@gmail.com>
Co-authored-by: Leclowndu93150 <xdr.a123454321@gmail.com>

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

@@ -14,6 +14,7 @@
 
 #include <array>
 #include <bit>
+#include <deque>
 #include <list>
 #include <memory>
 #include <optional>
@@ -509,6 +510,16 @@ struct UniformFillPlan {
 // Immutable runtime resource analysis retained by the shader cache. It owns descriptor/SRT,
 // uniform condition and fill values without retaining translated blocks.
 struct ResourcePlan {
+	struct EvaluationContext {
+		struct Entry {
+			uint64_t value      = 0;
+			uint64_t generation = 0;
+		};
+
+		std::vector<Entry> values;
+		uint64_t           generation = 0;
+	};
+
 	ResourcePlan() = default;
 	~ResourcePlan();
 
@@ -533,6 +544,10 @@ struct ResourcePlan {
 	bool                                resource_tracking_complete = false;
 	ShaderInfo                          info;
 	UniformFillPlan                     uniform_fill;
+	// GPU-thread scratch. Each nested clean/EXEC evaluation has its own reusable memo.
+	mutable std::deque<EvaluationContext> evaluation_contexts;
+	mutable uint32_t                       evaluation_value_count = 0;
+	mutable uint32_t                       evaluation_depth       = 0;
 };
 
 struct Program: ResourcePlan {

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

@@ -121,6 +121,13 @@ public:
 	[[nodiscard]] Block*                  PhiBlock(size_t index) const;
 	[[nodiscard]] Block*                  Parent() const;
 	[[nodiscard]] const std::vector<Use>& Uses() const;
+	// Runtime indices belong to the resource plan that owns this instruction.
+	[[nodiscard]] uint32_t EvaluationIndex(uint32_t& count) const {
+		if (evaluation_index == UINT32_MAX) {
+			evaluation_index = count++;
+		}
+		return evaluation_index;
+	}
 
 	void SetParent(Block* block);
 	void SetArg(size_t index, Value value);
@@ -155,6 +162,7 @@ private:
 	std::vector<Value>  args;
 	std::vector<Block*> phi_blocks;
 	std::vector<Use>    uses;
+	mutable uint32_t    evaluation_index = UINT32_MAX;
 };
 
 } // namespace Libs::Graphics::ShaderRecompiler::IR

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

@@ -8,7 +8,6 @@
 #include <cmath>
 #include <cstring>
 #include <fmt/format.h>
-#include <unordered_map>
 #include <unordered_set>
 
 namespace Libs::Graphics::ShaderRecompiler::IR {
@@ -464,7 +463,12 @@ public:
 	          std::span<const uint8_t> clean_flat_slots = {}, Evaluator* clean_evaluator = nullptr,
 	          Value active_mask = {})
 	    : m_program(program), m_runtime(runtime), m_clean_flat_slots(clean_flat_slots),
-	      m_clean_evaluator(clean_evaluator), m_active_mask(active_mask.Resolve()) {}
+	      m_clean_evaluator(clean_evaluator), m_active_mask(active_mask.Resolve()),
+	      m_context(AcquireContext(program)) {}
+
+	~Evaluator() { --m_program.evaluation_depth; }
+	Evaluator(const Evaluator&)            = delete;
+	Evaluator& operator=(const Evaluator&) = delete;
 
 	bool Evaluate(Value value, uint32_t& result) {
 		uint64_t wide = 0;
@@ -476,6 +480,15 @@ public:
 	}
 
 private:
+	static ResourcePlan::EvaluationContext& AcquireContext(const ResourcePlan& program) {
+		if (program.evaluation_depth == program.evaluation_contexts.size()) {
+			program.evaluation_contexts.emplace_back();
+		}
+		auto& context = program.evaluation_contexts[program.evaluation_depth++];
+		context.generation += 2;
+		return context;
+	}
+
 	static float Float32(uint64_t bits) {
 		return std::bit_cast<float>(static_cast<uint32_t>(bits));
 	}
@@ -497,30 +510,33 @@ private:
 		if (inst == nullptr) {
 			return false;
 		}
-		if (!m_reserved) {
-			m_cache.reserve(m_program.value_storage.size());
-			m_visiting.reserve(m_program.value_storage.size());
-			m_reserved = true;
-		}
 		if (!m_active_mask.IsEmpty() && IsRuntimeSelect(inst->GetOpcode()) &&
 		    inst->NumArgs() == 3 && inst->Arg(0).Resolve() == m_active_mask) {
 			return EvaluateWide(inst->Arg(1), result);
 		}
-		if (const auto found = m_cache.find(inst); found != m_cache.end()) {
+		const auto index = inst->EvaluationIndex(m_program.evaluation_value_count);
-			result = found->second;
+		if (index >= m_context.values.size()) {
+			m_context.values.resize(m_program.evaluation_value_count);
+		}
+		if (m_context.values[index].generation == m_context.generation) {
+			result = m_context.values[index].value;
 			return true;
 		}
-		if (std::ranges::find(m_visiting, inst) != m_visiting.end()) {
+		// The low generation bit marks an instruction that is still being evaluated.
+		if (m_context.values[index].generation == (m_context.generation | 1u)) {
 			return false;
 		}
-		m_visiting.push_back(inst);
+		m_context.values[index].generation = m_context.generation | 1u;
 		uint64_t out = 0;
 		const bool evaluated = EvaluateInst(*inst, out);
-		m_visiting.pop_back();
+		// Recursive evaluation may grow the dense memo vector.
+		auto& memo = m_context.values[index];
 		if (!evaluated) {
+			memo.generation = 0;
 			return false;
 		}
-		m_cache.emplace(inst, out);
+		memo.value      = out;
+		memo.generation = m_context.generation;
 		result = out;
 		return true;
 	}
@@ -977,9 +993,7 @@ private:
 	std::span<const uint8_t>                  m_clean_flat_slots;
 	Evaluator*                                m_clean_evaluator = nullptr;
 	Value                                     m_active_mask;
-	std::unordered_map<const Inst*, uint64_t> m_cache;
+	ResourcePlan::EvaluationContext&          m_context;
-	std::vector<const Inst*>                  m_visiting;
-	bool                                      m_reserved = false;
 };
 
 const DescriptorSource* Source(const ResourcePlan& program, uint32_t source) {

tests/ResourceTrackingTests.cpp

@@ -972,11 +972,13 @@ void TestFmaskLoadSpecialization() {
         "FMASK load did not lower to a value vector");
   result = vector->Arg(0);
   uint32_t value = 0;
-  Check(EvaluateUniformValues(plan, {&result, 1}, {.user_data = user_data}, {&value, 1}) &&
+  Check(EvaluateUniformValues(fixture.program, {&result, 1},
+                              {.user_data = user_data}, {&value, 1}) &&
             value == 0x76543210u,
         "FMASK load did not return the native sample-to-fragment mapping");
   user_data[8] = 1;
-  Check(EvaluateUniformValues(plan, {&result, 1}, {.user_data = user_data}, {&value, 1}) &&
+  Check(EvaluateUniformValues(fixture.program, {&result, 1},
+                              {.user_data = user_data}, {&value, 1}) &&
             value == 0u,
         "inactive FMASK load did not preserve the execution mask");
   ShaderComputeInputInfo compute{};
@@ -1184,6 +1186,14 @@ void TestSrtFlatteningAndRuntimeMemoization() {
             flat == std::vector<uint32_t>{0xdeadbeefu} && memory.reads == 1,
         "descriptor and flat SRT evaluation did not share one memoized read");
 
+  memory.reads = 0;
+  memory.words[1] = 0x12345678u;
+  Check(EvaluateRuntimeSources(fixture.program, std::span{&request, 1}, runtime,
+                               descriptors, flat, {}, active_sources) &&
+            descriptors[0].dwords[0] == 0x12345678u &&
+            flat == std::vector<uint32_t>{0x12345678u} && memory.reads == 1,
+        "repeated runtime evaluation reused stale scalar memory");
+
   memory.reads = 0;
   memory.fail_after = 0;
   descriptors = {{{1u}, 1u}};
@@ -1196,6 +1206,24 @@ void TestSrtFlatteningAndRuntimeMemoization() {
             active_sources == std::vector<uint8_t>{3u},
         "runtime evaluation failure was not transactional");
 
+  memory.fail_after = UINT32_MAX;
+  Check(EvaluateRuntimeSources(fixture.program, std::span{&request, 1}, runtime,
+                               descriptors, flat, {}, active_sources) &&
+            descriptors[0].dwords[0] == 0x12345678u && memory.reads == 1,
+        "failed runtime evaluation left a value marked as visiting");
+
+  auto detached = ExtractResourcePlan(fixture.program);
+  Check(EvaluateRuntimeSources(detached, std::span{&request, 1}, runtime,
+                               descriptors, flat, {}, active_sources),
+        "detached resource plan did not evaluate");
+  auto moved = std::move(detached);
+  memory.reads = 0;
+  memory.words[1] = 0x87654321u;
+  Check(EvaluateRuntimeSources(moved, std::span{&request, 1}, runtime,
+                               descriptors, flat, {}, active_sources) &&
+            descriptors[0].dwords[0] == 0x87654321u && memory.reads == 1,
+        "moving a cached resource plan lost its evaluation state");
+
   ShaderComputeInputInfo compute{};
   CollectShaderInfo(fixture.program, {.compute = &compute});
   AllocateBindings(fixture.program);