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;
 	}

tests/ResourceTrackingTests.cpp

@@ -2052,6 +2052,268 @@ void TestScalarBufferSelectorDomain() {
         "large V# bounds extended the native U32 domain or invented an OOB zero key");
 }
 
+void TestCountedRowImageSelector() {
+  namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
+  enum class Variant { Valid, Signed, WrongStep, Bypass, WrongRow, CountStride,
+                       KeyStride, CountImmediate, KeyImmediate };
+  const auto make_plan = [](Variant variant = Variant::Valid, uint32_t capacity = 16u) {
+    Fixture fixture;
+    auto *entry = fixture.block;
+    auto *header = fixture.AddBlock();
+    auto *body = fixture.AddBlock();
+    auto *latch = fixture.AddBlock();
+    auto *exit = fixture.AddBlock();
+    const auto branch = [](Block *from, Block *to) {
+      from->AddBranch(to);
+      from->terminator = {.kind = CFG::TerminatorKind::Branch, .true_block = to};
+    };
+    branch(entry, header);
+    branch(body, latch);
+    branch(latch, header);
+    const auto keys = fixture.Buffer({fixture.UserData(0), fixture.UserData(1),
+                                      fixture.UserData(2), fixture.UserData(3)});
+    const auto counts = fixture.Buffer({fixture.UserData(4), fixture.UserData(5),
+                                        fixture.UserData(6), fixture.UserData(7)});
+    const auto material = fixture.Buffer({fixture.UserData(8), fixture.UserData(9),
+                                          fixture.UserData(10), fixture.UserData(11)});
+    const auto row = fixture.Emit(ValueOpcode::GetBuiltin,
+        {Value(static_cast<uint32_t>(StageInputKind::WorkgroupId)), Value(0u)});
+    auto count_row = row;
+    if (variant == Variant::WrongRow)
+      count_row = fixture.Emit(ValueOpcode::IAdd32, {row, Value(1u)});
+    const auto count_offset = fixture.Emit(ValueOpcode::ShiftLeftLogical32,
+        {count_row, Value(variant == Variant::CountStride ? 3u : 2u)});
+    MemoryInfo scalar;
+    scalar.kind = ResourceKind::ScalarBuffer;
+    scalar.offset = variant == Variant::CountImmediate ? 4u : 0u;
+    const auto count = fixture.Emit(ValueOpcode::ReadConstBuffer,
+        {counts, count_offset}, fixture.AddMemory(scalar, 0x10));
+    auto &induction = header->AppendNewInst(ValueOpcode::Phi, {}, uint64_t(Type::U32));
+    const auto guard = fixture.Emit(variant == Variant::Signed ? ValueOpcode::SLessThan32
+                                                              : ValueOpcode::ULessThan32,
+        {Value(&induction), count}, 0, header);
+    header->AddBranch(body);
+    auto *outside = variant == Variant::Bypass ? latch : exit;
+    header->AddBranch(outside);
+    header->condition = guard;
+    header->terminator = {.kind = CFG::TerminatorKind::ConditionalBranch,
+                          .true_block = body, .false_block = outside};
+    fixture.block = body;
+    const auto row_start = fixture.Emit(ValueOpcode::IMul32, {row, Value(capacity)});
+    auto index = Value(&induction);
+    if (variant == Variant::KeyStride)
+      index = fixture.Emit(ValueOpcode::ShiftLeftLogical32, {index, Value(1u)});
+    const auto key_offset = fixture.Emit(ValueOpcode::ShiftLeftLogical32,
+        {fixture.Emit(ValueOpcode::IAdd32, {row_start, index}), Value(2u)});
+    scalar.offset = variant == Variant::KeyImmediate ? 4u : 0u;
+    const auto key = fixture.Emit(ValueOpcode::ReadConstBuffer,
+        {keys, key_offset}, fixture.AddMemory(scalar, 0x20));
+    const auto record = fixture.Emit(ValueOpcode::IAdd32,
+        {fixture.Emit(ValueOpcode::IMul32, {key, Value(396u)}), Value(184u)});
+    std::array<Value, 8> image_words;
+    for (uint32_t word = 0; word < image_words.size(); ++word) {
+      scalar.offset = word * 4u;
+      image_words[word] = fixture.Emit(ValueOpcode::ReadConstBuffer,
+          {material, record}, fixture.AddMemory(scalar, 0x30));
+    }
+    MemoryInfo sample;
+    sample.kind = ResourceKind::Image;
+    sample.image_dimension = Decoder::ImageDimension::Dim2D;
+    fixture.Emit(ValueOpcode::ImageSampleRaw,
+        {fixture.Image(image_words), fixture.Sampler({Value(0u), Value(0u), Value(0u), Value(0u)}),
+         fixture.ImageAddress(), Value(true)}, fixture.AddMemory(sample, 0x40));
+    const auto step = fixture.Emit(ValueOpcode::IAdd32,
+        {Value(&induction), Value(variant == Variant::WrongStep ? 2u : 1u)}, 0, latch);
+    induction.AddPhiOperand(entry, Value(0u));
+    induction.AddPhiOperand(latch, step);
+    fixture.PlanAndTrack();
+    const auto &source = fixture.program.descriptor_sources[fixture.program.info.images[0].source];
+    Check(source.indirect_descriptor.has_value(), "counted-row image lost its indirect source");
+    const auto &selector = source.indirect_descriptor->selector;
+    Check((selector && selector->rows.has_value()) == (variant == Variant::Valid),
+          "counted-row selector lost its proof or accepted mismatched row/induction bounds");
+    if (variant == Variant::Valid)
+      Check(selector->stride == 4u && selector->offset == 0u &&
+                selector->rows->stride == uint64_t(capacity) * 4u,
+            "counted-row selector changed native row or element byte spacing");
+    return ExtractResourcePlan(fixture.program);
+  };
+  auto plan = make_plan();
+  for (const auto variant : {Variant::Signed, Variant::WrongStep, Variant::Bypass,
+                             Variant::WrongRow, Variant::CountStride, Variant::KeyStride,
+                             Variant::CountImmediate, Variant::KeyImmediate})
+    make_plan(variant);
+
+  constexpr uint32_t CountBase = 0x1000u, KeyBase = 0x2000u, MaterialBase = 0x4000u;
+  // This key wraps 396*key to72, aliasing an unrelated field at material byte256.
+  constexpr uint32_t PoisonKey = 0x345d1746u;
+  LinearTestMemory memory;
+  memory.words.resize((MaterialBase - memory.base + 4u * 396u) / 4u);
+  const auto at = [&](uint32_t address) -> uint32_t& {
+    return memory.words[(address - memory.base) / 4u];
+  };
+  for (uint32_t index = 0; index < 64u; ++index) at(KeyBase + index * 4u) = PoisonKey;
+  for (uint32_t key = 0; key < 4u; ++key) {
+    const auto address = MaterialBase + key * 396u + 184u;
+    at(address) = 0x100u + key;
+    at(address + 4u) = static_cast<uint32_t>(Libs::Graphics::Prospero::BufferFormat::k32_32_32_32Float) << 20u;
+    at(address + 12u) = Libs::Graphics::DstSel(4, 5, 6, 7) |
+        (static_cast<uint32_t>(Libs::Graphics::Prospero::ImageType::kColor2D) << 28u);
+  }
+  std::copy_n(memory.words.begin() + (MaterialBase + 184u - memory.base) / 4u, 8u,
+              memory.words.begin() + (MaterialBase + 32u - memory.base) / 4u);
+  at(MaterialBase + 32u) = 0x104u;
+  at(MaterialBase + 256u) = 0x200u;
+  at(MaterialBase + 260u) = static_cast<uint32_t>(Libs::Graphics::Prospero::BufferFormat::k8SScaled) << 20u;
+  at(MaterialBase + 268u) = at(MaterialBase + 196u);
+  std::array<uint32_t, 12> user_data{KeyBase, 4u << 16u, 64u, 0x14004u,
+      CountBase, 4u << 16u, 4u, 0x14004u, MaterialBase, 396u << 16u, 4u, 0x14004u};
+  const SrtRuntime runtime{.user_data = user_data, .userdata = &memory,
+      .read_specialization_memory = ReadLinearTestMemory};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  const auto captured = [&](uint64_t address, uint64_t size) {
+    return std::ranges::any_of(snapshot.specialization_reads, [&](const auto &span) {
+      return span.first <= address && address - span.first <= span.second &&
+          size <= span.second - (address - span.first);
+    });
+  };
+  const auto check = [&](std::initializer_list<uint32_t> selected) {
+    memory.watched_address = KeyBase;
+    memory.watched_reads = 0u;
+    Check(MaterializeResources(plan, runtime, snapshot, specialization),
+          "counted-row selector interpreted an unused slot or failed coherent discovery");
+    Check(captured(CountBase, user_data[6] * 4u),
+          "row counts escaped writable-overlap capture");
+    std::vector<uint32_t> expected(selected);
+    std::ranges::sort(expected);
+    expected.erase(std::unique(expected.begin(), expected.end()), expected.end());
+    if (expected.empty()) {
+      Check(memory.watched_reads == 0u && snapshot.images.size() == 1u &&
+                std::ranges::all_of(snapshot.images[0].dwords, [](uint32_t word) { return word == 0u; }),
+            "empty row counts read unused keys or retained stale image candidates");
+      return;
+    }
+    if (expected.size() > 1u) {
+      const auto map = specialization.images[0].indirect_mapping_offset;
+      Check(snapshot.flattened_srt[map] == expected.size(), "counted-row mapping lost selected keys");
+      for (uint32_t index = 0; index < expected.size(); ++index)
+        Check(snapshot.flattened_srt[map + 1u + index * 2u] == expected[index],
+              "counted-row mapping changed an unsigned key or retained an unused slot");
+    }
+    std::vector<uint32_t> image_addresses;
+    for (const auto key : expected) {
+      const auto offset = uint32_t(key * 396u) + 184u;
+      image_addresses.push_back(at(MaterialBase + offset));
+    }
+    std::ranges::sort(image_addresses);
+    image_addresses.erase(std::unique(image_addresses.begin(), image_addresses.end()), image_addresses.end());
+    Check(snapshot.images.size() == image_addresses.size() &&
+              std::ranges::all_of(snapshot.images, [&](const DescriptorValue &image) {
+                return std::ranges::find(image_addresses, image.dwords[0]) != image_addresses.end();
+              }), "counted rows lost a selected image or retained an unrelated descriptor");
+  };
+  check({});
+  at(CountBase) = 2u; at(CountBase + 8u) = 1u;
+  at(KeyBase) = 2u; at(KeyBase + 4u) = 1u; at(KeyBase + 128u) = 3u;
+  check({1u, 2u, 3u});
+  Check(captured(KeyBase, 8u) && captured(KeyBase + 128u, 4u),
+        "active row keys escaped writable-overlap capture");
+  at(KeyBase) = 0x80000001u;
+  check({1u, 3u, 0x80000001u});
+  at(KeyBase) = 0x1890cedeu; // 396*key +184 wraps to material byte32.
+  check({1u, 3u, 0x1890cedeu});
+  at(CountBase) = 0u;
+  user_data[2] = 1u; // The selected second row reads beyond the index V# and supplies key0.
+  check({0u});
+  user_data[2] = 64u;
+  at(CountBase + 8u) = 0u;
+  check({});
+  at(CountBase) = 1u;
+  at(KeyBase) = PoisonKey;
+  Check(!MaterializeResources(plan, runtime, snapshot, specialization),
+        "an actively selected malformed material descriptor was suppressed");
+  at(KeyBase) = 2u;
+  check({2u});
+  for (const auto fail : {CountBase, KeyBase}) {
+    memory.fail_address = fail;
+    Check(!MaterializeResources(plan, runtime, snapshot, specialization),
+          "failed count/key discovery reused stale selector data");
+  }
+  memory.fail_address = UINT64_MAX;
+  for (uint32_t index = 0; index < 17u; ++index) at(KeyBase + index * 4u) = 1u;
+  at(CountBase) = 17u;
+  check({1u}); // Native counts can spill past the row spacing into the next row's slots.
+  at(CountBase + 4u) = 1u;
+  check({1u}); // Overlapping nonempty prefixes preserve the union without duplicating images.
+  at(CountBase + 4u) = 0u;
+  for (uint32_t index = 0; index < 64u; ++index) at(KeyBase + index * 4u) = 1u;
+  at(CountBase) = MaxIndirectDescriptorProbes;
+  check({0u, 1u}); // The exact cap is valid; selected OOB scalar entries contribute key0.
+  for (const auto count : {65537u, 0x80000000u, UINT32_MAX}) {
+    at(CountBase) = count;
+    memory.watched_reads = 0u;
+    Check(!MaterializeResources(plan, runtime, snapshot, specialization) && memory.watched_reads == 0u,
+          "unsigned row count exceeded the probe cap or wrapped before index discovery");
+  }
+  at(CountBase) = 0u;
+  user_data[6] = MaxIndirectDescriptorProbes + 1u;
+  memory.reads = 0u;
+  Check(!MaterializeResources(plan, runtime, snapshot, specialization) && memory.reads == 0u,
+        "counted-row discovery exceeded the existing row probe cap");
+  user_data[6] = 0u;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) && memory.reads == 0u &&
+            snapshot.images.size() == 1u && snapshot.images[0].dwords[0] == 0u,
+        "empty count V# manufactured a selector key or retained stale images");
+
+  user_data[6] = 5u;
+  at(CountBase + 16u) = 1u;
+  auto wrap_plan = make_plan(Variant::Valid, 1u << 28u);
+  memory.watched_reads = 0u;
+  Check(!MaterializeResources(wrap_plan, runtime, snapshot, specialization) && memory.watched_reads == 0u,
+        "row address multiplication wrapped before the active prefix was validated");
+
+  struct OverflowMemory { uint32_t count_reads = 0, other_reads = 0; } overflow;
+  user_data[6] = MaxIndirectDescriptorProbes / 16u + 1u;
+  const auto read_overflow = +[](void *data, uint64_t address, std::span<uint32_t> words) {
+    auto &memory = *static_cast<OverflowMemory *>(data);
+    if (address == CountBase && words.size() == MaxIndirectDescriptorProbes / 16u + 1u) {
+      ++memory.count_reads;
+      std::ranges::fill(words, 16u);
+      return true;
+    }
+    ++memory.other_reads;
+    return false;
+  };
+  Check(!MaterializeResources(plan,
+            {.user_data = user_data, .userdata = &overflow, .read_specialization_memory = read_overflow},
+            snapshot, specialization) && overflow.count_reads == 1u && overflow.other_reads == 0u,
+        "sum of active row prefixes exceeded the probe cap or read keys before validation");
+
+  auto sparse_plan = make_plan(Variant::Valid, 1u << 16u);
+  user_data[2] = 2u * (1u << 16u) + 1u;
+  user_data[6] = 3u;
+  struct SparseRows { LinearTestMemory *material; uint32_t key_reads = 0u; } sparse{&memory};
+  const auto read_sparse = +[](void *data, uint64_t address, std::span<uint32_t> words) {
+    auto &memory = *static_cast<SparseRows *>(data);
+    if (address == CountBase && words.size() == 3u) {
+      words[0] = words[2] = 1u; words[1] = 0u;
+      return true;
+    }
+    if (address == KeyBase || address == KeyBase + 0x80000u) {
+      if (words.size() != 1u) return false;
+      ++memory.key_reads;
+      words[0] = address == KeyBase ? 2u : 3u;
+      return true;
+    }
+    return ReadLinearTestMemory(memory.material, address, words);
+  };
+  Check(MaterializeResources(sparse_plan,
+            {.user_data = user_data, .userdata = &sparse, .read_specialization_memory = read_sparse},
+            snapshot, specialization) && sparse.key_reads == 2u && snapshot.images.size() == 2u,
+        "sparse rows exceeded the batch cap or read unselected gaps");
+}
+
 void TestSharedAppendSelectorBounds() {
   namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
   enum class Variant {
@@ -4702,6 +4964,7 @@ int main() {
     Run("image descriptor fields", TestImageDescriptorFields);
     Run("draw-uniform scalar image", TestUniformScalarBufferImage);
     Run("scalar buffer selector domain", TestScalarBufferSelectorDomain);
+    Run("counted row image selector", TestCountedRowImageSelector);
     Run("shared append selector bounds", TestSharedAppendSelectorBounds);
     Run("descending image range", TestDescendingImageRange);
     Run("SRT runtime", TestSrtFlatteningAndRuntimeMemoization);

tests/ShaderRecompilerComputeTests.cpp

@@ -36176,6 +36176,120 @@ void CheckFiniteInlineSamplerPhi(VulkanHarness &vulkan) {
   std::printf("[compute] %-32s ok\n", test.name);
 }
 
+void CheckCountedRowImageSelector(VulkanHarness &vulkan) {
+  using O = ShaderOpcode;
+  using F = Prospero::BufferFormat;
+  constexpr u32 poison = 0x345d1746u;
+  const auto alignment = static_cast<u32>(vulkan.RuntimeContext().GetPhysicalDeviceProperties()
+                                             .limits.minStorageBufferOffsetAlignment);
+  const auto slot = std::max(256u, alignment);
+  const u32 srt = slot, counts = slot * 2u, keys = slot * 3u, material = slot * 4u;
+  for (const u32 wave : {32u, 64u}) {
+    TestCase test;
+    test.name = wave == 32u ? "CountedRowImageSelector32" : "CountedRowImageSelector64";
+    test.initial.resize((material + 396u * 4u) / 4u, 0u);
+    std::fill_n(test.initial.begin(), 64u, 0xdeadbeefu);
+    const std::array<u32, 12> descriptors{counts, 4u << 16u, 3u, 0x14004u,
+        keys, 4u << 16u, 48u, 0x14004u, material, 396u << 16u, 4u, 0x14004u};
+    std::copy(descriptors.begin(), descriptors.end(), test.initial.begin() + srt / 4u);
+    test.initial[counts / 4u] = 3u;
+    test.initial[counts / 4u + 2u] = 1u;
+    std::fill_n(test.initial.begin() + keys / 4u, 48u, poison);
+    test.initial[keys / 4u] = 2u;
+    test.initial[keys / 4u + 1u] = 0x80000001u;
+    test.initial[keys / 4u + 2u] = 0x1890cedeu; // Dynamic +184 wraps to material byte32.
+    test.initial[keys / 4u + 32u] = 3u;
+    const auto put_image = [&](u32 offset, u32 identity, F format) {
+      auto at = test.initial.begin() + (material + offset) / 4u;
+      at[0] = identity;
+      at[1] = static_cast<u32>(format) << 20u;
+      at[3] = DstSel(4, 5, 6, 7) |
+          (static_cast<u32>(Prospero::ImageType::kColor2D) << 28u);
+    };
+    for (u32 key = 0; key < 4u; ++key)
+      put_image(key * 396u + 184u, 0x100u + key, F::k32_32_32_32Float);
+    put_image(32u, 0x104u, F::k32_32_32_32Float);
+    put_image(256u, 0x200u, F::k8SScaled);
+    test.expected = test.initial;
+    test.expected[0] = 0x3f800002u;
+    test.expected[1] = 0x3f800001u;
+    test.expected[2] = 0x3f800004u;
+    test.expected[32] = 0x3f800003u;
+    test.user_data[0] = srt;
+    test.user_data[50] = 256u;
+    test.user_data[51] = 3u << 28u;
+    test.has_user_data = true;
+    test.has_compute_info = true;
+    test.compute_info.wave_size = wave;
+    test.compute_info.threads_num[0] = test.compute_info.threads_num[1] =
+        test.compute_info.threads_num[2] = 1u;
+    test.compute_info.group_id[0] = true;
+    test.compute_info.workgroup_register = 44u;
+    test.dispatch_x = 3u;
+    auto &code = test.code;
+    code.insert(code.end(), {EncodeSmem0(0x03, 8), EncodeSmem1(0, 125),
+                            EncodeSmem0(0x02, 16), EncodeSmem1(32, 125)});
+    code.push_back(EncodeSop2(0x1e, 22, 44, InlineU32(2)));
+    code.insert(code.end(), {EncodeSmem0(0x08, 21, 4), EncodeSmem1(0, 22)});
+    AppendSMovLiteral(&code, 20, 0u);
+    AppendVMovU32(&code, 0, 0x3f000000u);
+    AppendVMovU32(&code, 1, 0x3f000000u);
+    const auto header = code.size();
+    code.push_back(EncodeSopc(0x0a, 20, 21));
+    const auto exit_branch = code.size();
+    code.push_back(0u);
+    code.push_back(EncodeSop2(0x1e, 22, 44, InlineU32(4)));
+    code.push_back(EncodeSop2(0x02, 22, 22, 20));
+    code.push_back(EncodeSop2(0x1e, 22, 22, InlineU32(2)));
+    code.insert(code.end(), {EncodeSmem0(0x08, 23, 6), EncodeSmem1(0, 22)});
+    code.push_back(EncodeSop2(0x26, 24, 23, 255)); code.push_back(396u);
+    code.push_back(EncodeSop2(0x02, 24, 24, 255)); code.push_back(184u);
+    code.insert(code.end(), {EncodeSmem0(0x0b, 32, 8), EncodeSmem1(0, 24),
+                            EncodeMimg0(0x27, 1), EncodeMimg1(4, 0, 8, 10)});
+    code.push_back(EncodeVop1(0x01, 30, 22));
+    AppendBufferStoreDword(&code, 4, 30);
+    code.push_back(EncodeSop2(0x02, 20, 20, InlineU32(1)));
+    code.push_back(EncodeSopp(0x02, header - code.size() - 1u));
+    code[exit_branch] = EncodeSopp(0x04, code.size() - exit_branch - 1u);
+    AppendEnd(&code);
+    test.opcodes = {O::S_LOAD_DWORDX8, O::S_LOAD_DWORDX4, O::S_BUFFER_LOAD_DWORD,
+        O::S_BUFFER_LOAD_DWORDX8, O::S_MOV_B32, O::S_LSHL_B32, O::S_ADD_I32,
+        O::S_MUL_I32, O::S_CMP_LT_U32, O::S_CBRANCH_SCC0, O::S_BRANCH,
+        O::V_MOV_B32, O::IMAGE_SAMPLE, O::BUFFER_STORE_DWORD, O::S_ENDPGM};
+    const auto compiled = CompileCase(test, vulkan.SubgroupSize());
+    Require(test.name, "exact active images", compiled.resources.images.size() == 4u &&
+                std::ranges::all_of(compiled.resources.images, [](const auto &image) {
+                  return image.dwords[0] >= 0x101u && image.dwords[0] <= 0x104u;
+                }), "unused row entries selected malformed material fields");
+    for (u32 word = 0; word < 4u; ++word)
+      Require(test.name, "count descriptor capture",
+              std::ranges::any_of(compiled.resources.specialization_reads, [&](const auto &read) {
+                return read.first <= srt + word * 4u &&
+                    read.first + read.second >= srt + (word + 1u) * 4u;
+              }), "count V# loaded from SRT escaped writable-alias capture");
+    for (const auto &buffer : compiled.resources.buffers) {
+      ShaderBufferResource descriptor{};
+      std::copy_n(buffer.dwords.begin(), 4u, descriptor.fields);
+      test.storage_buffer_ranges.emplace_back(static_cast<u32>(descriptor.Base48()),
+                                              static_cast<u32>(descriptor.GetSize()));
+    }
+    std::vector<VulkanHarness::Image> images;
+    for (const auto &image : compiled.resources.images)
+      images.push_back(vulkan.CreateImageMips(test.name, 1, 1, vk::Format::eR32G32B32A32Sfloat,
+          vk::ImageUsageFlagBits::eSampled, {{0x3f800000u + image.dwords[0] - 0x100u, 0u, 0u, 0u}},
+          4u, vk::ImageLayout::eShaderReadOnlyOptimal, vk::ImageType::e2D, vk::ImageViewType::e2D, 1));
+    auto sampler = vulkan.CreateSampler(test.name, 1);
+    auto output = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size());
+    vulkan.Dispatch(test, compiled, output, nullptr, nullptr, nullptr, nullptr, sampler, images);
+    CompareWords(test, "active row images, wrapped keys and untouched unused slots", test.expected,
+                 vulkan.ReadBuffer(test.name, output, test.expected.size()));
+    vulkan.DestroyBuffer(&output);
+    vulkan.Device().destroySampler(sampler);
+    for (auto &image : images) vulkan.DestroyImage(&image);
+    std::printf("[compute] %-32s ok\n", test.name);
+  }
+}
+
 void CheckIndirectImageKeySwitch(VulkanHarness &vulkan) {
   constexpr const char *name = "IndirectImageKeySwitch";
   constexpr uint32_t mapping_capacity = 1793u;
@@ -44568,6 +44682,7 @@ int main(int argc, char **argv) {
   if (argc == 2 && std::strcmp(argv[1], "--indirect-image-only") == 0) {
     CheckImageSamplerSpecialization();
     VulkanHarness vulkan;
+    CheckCountedRowImageSelector(vulkan);
     CheckIndirectImageKeySwitch(vulkan);
     CheckFiniteInlineSamplerPhi(vulkan);
     RunCase(&vulkan, ImageCubeGradientsPreserveDerivatives());
@@ -44733,6 +44848,7 @@ int main(int argc, char **argv) {
   CheckIndirectBufferStore(vulkan);
   CheckRuntimeBufferRecords(vulkan);
   CheckComputeThreadDimensions(vulkan);
+  CheckCountedRowImageSelector(vulkan);
   CheckIndirectImageKeySwitch(vulkan);
   CheckDynamicSampleOffsets(vulkan, 32u);
   CheckDynamicSampleOffsets(vulkan, 64u);