shader: bound indirect image selectors by counted scalar rows

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

@@ -504,11 +504,18 @@ struct ShaderInfo {
 struct DescriptorSource {
 	struct IndirectDescriptor {
 		struct SelectorRead {
+			struct CountedRows {
+				uint32_t source = UINT32_MAX;
+				uint32_t stride = 0; // Byte spacing between selector rows, not their length.
+
+				bool operator==(const CountedRows& other) const = default;
+			};
 			uint32_t source = UINT32_MAX;
 			uint32_t stride = 0;
 			uint32_t offset = 0;
 			// Four DWORD bitmap words per record; set-bit ranks subtract from the last key.
 			bool bitset = false;
+			std::optional<CountedRows> rows;
 			std::vector<std::pair<uint32_t, uint32_t>> byte_ranges;
 
 			bool operator==(const SelectorRead& other) const = default;

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

@@ -346,6 +346,50 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 			std::iota(keys.begin(), keys.end(), first_key);
 			if (!keys.empty() && keys.back() < keys.front())
 				std::rotate(keys.begin(), keys.begin() + (UINT32_MAX - first_key + 1ull), keys.end());
+		} else if (selector->rows) {
+			const auto& rows = *selector->rows;
+			DescriptorValue count_value;
+			ShaderBufferResource counts, material;
+			if (!clean.EvaluateDescriptor(rows.source, count_value) ||
+			    !DecodeBufferDescriptor(count_value, counts) ||
+			    !DecodeBufferDescriptor(material_value, material) || selector->stride != 4u ||
+			    selector->offset != 0u || rows.stride == 0u || (rows.stride & 3u) != 0u ||
+			    counts.GetSize() / 4u > MaxIndirectDescriptorProbes) return false;
+			std::vector<uint32_t> lengths(counts.GetSize() / 4u);
+			if (!ReadScalarTable(counts.Base48(), counts.GetSize(), 0u, runtime, lengths)) return false;
+			uint32_t total = 0;
+			for (size_t row = 0; row < lengths.size(); ++row) {
+				const auto length = lengths[row];
+				if (length > MaxIndirectDescriptorProbes - total ||
+				    (length != 0u && uint64_t {row} * rows.stride + uint64_t {length} * 4u >
+				                         UINT32_MAX + 1ull)) return false;
+				total += length;
+			}
+			keys.reserve(total);
+			std::vector<uint32_t> batch;
+			for (size_t row = 0; row < lengths.size();) {
+				if (lengths[row] == 0u) { ++row; continue; }
+				const auto first = uint64_t {row} * rows.stride / 4u;
+				auto end = first + lengths[row];
+				size_t next = row + 1u;
+				for (; next < lengths.size(); ++next) {
+					if (lengths[next] == 0u) continue;
+					const auto candidate_end =
+					    std::max(end, uint64_t {next} * rows.stride / 4u + lengths[next]);
+					if (candidate_end - first > MaxIndirectDescriptorProbes) break;
+					end = candidate_end;
+				}
+				// Coalesce small gaps, but interpret only selected prefixes. The strict
+				// reader captures the entire fetched span for writable-alias rejection.
+				batch.resize(end - first);
+				if (!ReadScalarTable(material.Base48(), material.GetSize(), first * 4u, runtime, batch))
+					return false;
+				for (; row < next; ++row) {
+					if (lengths[row] == 0u) continue;
+					const auto begin = batch.begin() + (uint64_t {row} * rows.stride / 4u - first);
+					keys.insert(keys.end(), begin, begin + lengths[row]);
+				}
+			}
 		} else if (selector->bitset) {
 			ShaderBufferResource material;
 			std::array<uint32_t, 4> masks {};
@@ -1074,8 +1118,11 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
 		for (const auto& source: plan.descriptor_sources) {
 			if (!source.indirect_descriptor.has_value()) continue;
 			const auto& indirect = *source.indirect_descriptor;
-			if (indirect.selector)
+			if (indirect.selector) {
 				MarkCleanReads(planned_reads, Source(plan, indirect.selector->source));
+				if (indirect.selector->rows)
+					MarkCleanReads(planned_reads, Source(plan, indirect.selector->rows->source));
+			}
 			MarkCleanReads(planned_reads, nullptr, indirect.selector_mask);
 			MarkCleanReads(planned_reads, nullptr, indirect.selector_first);
 			MarkCleanReads(planned_reads, nullptr, indirect.key_count);

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

@@ -1825,6 +1825,44 @@ private:
 		return bound;
 	}
 
+	bool MatchCountedRows(const Inst& read, DescriptorSource& count_source, uint32_t& row_stride) {
+		Value element;
+		uint32_t offset = 0, stride = 0;
+		if (!MatchTableOffset(read.Arg(1), element, offset, stride) || offset != 0u || stride != 4u)
+			return false;
+		const auto* sum = element.TryInstruction();
+		if (sum == nullptr || sum->GetOpcode() != ValueOpcode::IAdd32) return false;
+		for (uint32_t arm = 0; arm < 2u; ++arm) {
+			const auto induction = sum->Arg(arm).Resolve();
+			if (BoundedLoop(induction, read.Parent(), [&](const Inst& compare, uint32_t) {
+				if (compare.GetOpcode() != ValueOpcode::ULessThan32) return false;
+				const auto* count = compare.Arg(1).Resolve().TryInstruction();
+				if (count == nullptr || count->GetOpcode() != ValueOpcode::ReadConstBuffer) return false;
+				uint32_t memory_index = 0;
+				const auto* memory = ScalarReadMemory(*count, memory_index);
+				Value row;
+				uint32_t count_offset = 0, count_stride = 0;
+				if (memory == nullptr || memory->offset != 0u ||
+				    !MatchTableOffset(count->Arg(1), row, count_offset, count_stride) ||
+				    count_offset != 0u || count_stride != 4u) return false;
+				const auto affine = AffineIndex<3>(read.Arg(1), [](Value value) { return value.Resolve(); },
+				    [&](Value value) -> std::optional<size_t> {
+					    if (value == induction) return 1u;
+					    if (EquivalentValue(m_program, value, row)) return 2u;
+					    return {};
+				    });
+				if (!affine || (*affine)[0] != 0u || (*affine)[1] != 4u ||
+				    (*affine)[2] == 0u || ((*affine)[2] & 3u) != 0u ||
+				    !MakeRuntimeTableSource(*count, count_source)) return false;
+				// Count addresses identify rows modulo 2^30. A DWORD-aligned row
+				// spacing preserves those aliases in the U32 selector address too.
+				row_stride = static_cast<uint32_t>((*affine)[2]);
+				return true;
+			}) != nullptr) return true;
+		}
+		return false;
+	}
+
 	bool MatchScalarMaterialKey(Value key, DescriptorSource::IndirectDescriptor& indirect,
 	                            DescriptorSource& material_source) {
 		const auto* read = key.Resolve().TryInstruction();
@@ -1833,6 +1871,8 @@ private:
 		const auto* memory = ScalarReadMemory(*read, memory_index);
 		Value index;
 		Value count;
+		DescriptorSource count_source;
+		uint32_t row_stride = 0;
 		uint32_t offset = 0, stride = 0;
 		if (memory == nullptr || memory->kind != ResourceKind::ScalarBuffer ||
 		    memory->offset > INT32_MAX || !MemoryIndexBelongsTo(memory_index, *read)) return false;
@@ -1844,7 +1884,12 @@ private:
 				count = bound->Arg(1);
 			}
 		}
-		if (count.IsEmpty()) {
+		const bool rows = count.IsEmpty() && memory->offset == 0u &&
+		    MatchCountedRows(*read, count_source, row_stride);
+		if (rows) {
+			stride = 4u;
+			offset = 0u;
+		} else 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;
@@ -1870,6 +1915,9 @@ private:
 		    .source = InternSource(material_source), .stride = stride,
 		    .offset = offset + (memory->offset & ~3u)});
 		indirect.key_count = count;
+		if (rows) indirect.selector->rows.emplace(
+		    DescriptorSource::IndirectDescriptor::SelectorRead::CountedRows {
+		        .source = InternSource(count_source), .stride = row_stride});
 		if (!count.IsEmpty()) indirect.selector_first = Value(0u);
 		return true;
 	}