shader: resolve GPU scalar image selectors within buffer bounds

Co-authored-by: Ali Almohaya <AliAlmohaya@gmail.com>

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

@@ -394,7 +394,22 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 				keys.push_back(key);
 				mask &= mask - 1u;
 			}
-		} else return false;
+		} else {
+			ShaderBufferResource material;
+			if (!DecodeBufferDescriptor(material_value, material) || selector->stride < 4u ||
+			    !std::has_single_bit(selector->stride) || (selector->offset & 3u) != 0u)
+				return false;
+			// Empty selector bounds denote the complete U32 congruence class. SMEM
+			// adds its aligned immediate after the dynamic offset, without U32 wrap.
+			const auto domain_count = (uint64_t {1} << 32u) / selector->stride;
+			const auto available = material.GetSize() > selector->offset
+			    ? material.GetSize() - selector->offset : 0u;
+			const auto count = std::min(domain_count,
+			    available < 4u ? 0u : (available - 4u) / selector->stride + 1u);
+			if (count > MaxIndirectDescriptorProbes ||
+			    !read_keys(material, selector->offset, selector->stride, count)) return false;
+			if (count < domain_count) keys.push_back(0u); // An out-of-bounds scalar read returns zero.
+		}
 		if (selector != nullptr && selector->byte_ranges.empty()) {
 			std::ranges::sort(keys);
 			keys.erase(std::unique(keys.begin(), keys.end()), keys.end());

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

@@ -8,6 +8,7 @@
 #include <algorithm>
 #include <fmt/format.h>
 #include <map>
+#include <numeric>
 #include <optional>
 #include <span>
 #include <tuple>
@@ -1152,6 +1153,7 @@ private:
 
 	bool HasActiveLane(Value mask, const Block* use) const {
 		mask = mask.Resolve();
+		if (mask.IsImmediate()) return mask.GetType() == Type::U1 && mask.U1();
 		const auto incoming_is_nonempty = [&](const Block* from, const Block* to,
 		                                     Value incoming, const Block* header) {
 			const auto accepts = [&](const EdgePredicate& edge) {
@@ -1534,20 +1536,37 @@ private:
 		uint32_t memory_index = 0;
 		const auto* memory = ScalarReadMemory(*read, memory_index);
 		Value index;
+		Value count;
 		uint32_t offset = 0, stride = 0;
 		if (memory == nullptr || memory->kind != ResourceKind::ScalarBuffer ||
-		    memory->offset > INT32_MAX || !MemoryIndexBelongsTo(memory_index, *read) ||
-		    !MatchTableOffset(read->Arg(1), index, offset, stride) ||
-		    ((offset | stride | memory->offset) & 3u) != 0u ||
-		    uint64_t {offset} + memory->offset > UINT32_MAX) return false;
-		const auto* bound = BoundedLoop(index, read->Parent());
-		if (bound == nullptr || bound->GetOpcode() != ValueOpcode::ULessThan32 ||
-		    !MakeRuntimeTableSource(*read, material_source)) return false;
+		    memory->offset > INT32_MAX || !MemoryIndexBelongsTo(memory_index, *read)) return false;
+		if (MatchTableOffset(read->Arg(1), index, offset, stride) &&
+		    ((offset | stride | memory->offset) & 3u) == 0u &&
+		    uint64_t {offset} + memory->offset <= UINT32_MAX) {
+			const auto* bound = BoundedLoop(index, read->Parent());
+			if (bound != nullptr && bound->GetOpcode() == ValueOpcode::ULessThan32) {
+				count = bound->Arg(1);
+			}
+		}
+		if (count.IsEmpty()) {
+			const auto* lane = read->Arg(1).Resolve().TryInstruction();
+			if (lane == nullptr || lane->GetOpcode() != ValueOpcode::ReadFirstLane ||
+			    !HasActiveLane(lane->Arg(1), lane->Parent())) return false;
+			AffineOffset affine;
+			if (!MatchAffineOffset(lane->Arg(0), lane->Arg(1), affine) ||
+			    affine.index.IsEmpty() || affine.stride == 0u) return false;
+			// U32 multiplication can wrap. Its complete byte-offset domain is one
+			// congruence class, not the source V#'s structured record stride.
+			const auto divisor = std::gcd(affine.stride, uint64_t {1} << 32u);
+			stride = static_cast<uint32_t>(std::max<uint64_t>(divisor, 4u));
+			offset = static_cast<uint32_t>(affine.offset % divisor) & ~3u;
+		}
+		if (!MakeRuntimeTableSource(*read, material_source)) return false;
 		indirect.selector.emplace(DescriptorSource::IndirectDescriptor::SelectorRead {
 		    .source = InternSource(material_source), .stride = stride,
-		    .offset = offset + memory->offset});
-		indirect.selector_first = Value(0u);
-		indirect.key_count = bound->Arg(1);
+		    .offset = offset + (memory->offset & ~3u)});
+		indirect.key_count = count;
+		if (!count.IsEmpty()) indirect.selector_first = Value(0u);
 		return true;
 	}
 
@@ -1768,6 +1787,11 @@ private:
 			// The descriptor also supplies dimensions to shader arithmetic. Keep its reads;
 			// only the image handle is projected onto the bounded workgroup key.
 			plan.retain_reads = true;
+		} else if (!plan.table_indexed &&
+		           MatchScalarMaterialKey(key, indirect, material_source)) {
+			// Keep live descriptor payloads, but project descriptor-only SMEM onto
+			// the GPU key using the same plan as other indirect image reads.
+			plan.retain_reads |= table_source.dword_count == 4u;
 		} else if (table_source.dword_count == 2u) {
 			if (table_stride != 32u) return false;
 			const auto* selector = key.Resolve().TryInstruction();
@@ -1786,10 +1810,6 @@ private:
 			}
 			if ((table_offset & 3u) != 0u ||
 			    (bitscan && table_offset > UINT32_MAX - (32u * 32u - 1u))) return false;
-		} else if (!plan.table_indexed &&
-		           MatchScalarMaterialKey(key, indirect, material_source)) {
-			// The host enumerates the guarded scalar loop; its key and T# payload stay on the GPU.
-			plan.retain_reads = true;
 		} else {
 			// A bounded V# supplies the complete image table. Leave every GPU selector
 			// and descriptor read in the shader; the host only translates table bytes.