shader: bound indexed material image selectors before heap expansion

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

@@ -292,11 +292,23 @@ bool NormalizeIndirectBuffer(DescriptorValue& value, bool indexed, const BufferR
 bool IsBoundedDescriptorTable(const ResourcePlan& program,
                               const DescriptorSource::IndirectDescriptor& indirect) {
 	const auto* table = Source(program, indirect.table_source);
-	return indirect.sources.empty() && !indirect.selector &&
+	return indirect.sources.empty() && (!indirect.selector || indirect.table_scalar) &&
 	       (indirect.key_count.IsEmpty() || indirect.table_maximum) && indirect.workgroup_axis == UINT32_MAX &&
 	       table != nullptr && table->dword_count == 4u;
 }
 
+bool CanReadIndexedSelector(const ShaderBufferResource& material, uint32_t offset) {
+	if (material.Type() != 0u || material.SwizzleEnabled() || material.AddTid() ||
+	    material.OutOfBounds() != 0u) return false;
+	if (material.Format() == Prospero::BufferFormat::kInvalid || material.NumRecords() == 0u ||
+	    offset >= material.Stride()) return true;
+	return material.Base48() != 0u &&
+	       ((material.Base48() | material.Stride() | offset) & 3u) == 0u &&
+	       material.NumRecords() <= MaxIndirectDescriptorProbes &&
+	       uint64_t {material.NumRecords() - 1u} * material.Stride() + offset + 4u <=
+	           uint64_t {UINT32_MAX} + 1u;
+}
+
 bool EvaluateIndirectCount(const DescriptorSource::IndirectDescriptor& indirect,
                            const SrtRuntime& runtime, SrtWalker& clean, uint32_t& count) {
 	if (indirect.workgroup_axis != UINT32_MAX) {
@@ -490,6 +502,15 @@ bool ResolveSelectorKeys(const DescriptorSource::IndirectDescriptor& indirect,
 			keys.push_back(key);
 			mask &= mask - 1u;
 		}
+	} else if (selector->stride == 0u) {
+		ShaderBufferResource material;
+		if (!DecodeBufferDescriptor(material_value, material) ||
+		    !CanReadIndexedSelector(material, selector->offset)) return false;
+		if (material.Format() != Prospero::BufferFormat::kInvalid && material.NumRecords() != 0u &&
+		    selector->offset < material.Stride() &&
+		    !read_keys(material.Base48(), material.GetSize(), selector->offset,
+		               material.Stride(), material.NumRecords())) return false;
+		keys.push_back(0u); // Every out-of-bounds indexed load returns zero.
 	} else {
 		ShaderBufferResource material;
 		if (!DecodeBufferDescriptor(material_value, material) || selector->stride < 4u ||
@@ -529,9 +550,11 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 	const auto  descriptor_bytes = read_dwords * sizeof(uint32_t);
 	const auto& sources = indirect.sources;
 	const auto* selector = indirect.selector ? &*indirect.selector : nullptr;
+	const bool byte_keys = IsBoundedDescriptorTable(program, indirect);
 	auto* selector_keys = selector != nullptr
 	    ? &program.selector_keys[program.selector_groups[source_index]] : nullptr;
 	auto& keys = selector_keys != nullptr ? selector_keys->values : program.material_keys;
+	auto& byte_offsets = program.material_keys;
 	const auto children_begin = descriptors.size();
 	const auto root_resource = specializations[resource_index];
 	const auto intern_candidate = [&](const DescriptorValue& candidate, uint32_t& ordinal) {
@@ -583,7 +606,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 	if (sources.empty()) {
 		if (selector == nullptr) keys.clear();
 		DescriptorValue material_value;
-		if ((selector != nullptr && !selector_keys->resolved &&
+		if ((selector != nullptr && (byte_keys || !selector_keys->resolved) &&
 		     !clean.EvaluateDescriptor(selector->source, material_value)) ||
 		    !clean.EvaluateDescriptor(indirect.table_source, table_value)) {
 			return false;
@@ -596,16 +619,40 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 		} else {
 			return false;
 		}
-		if (IsBoundedDescriptorTable(program, indirect)) {
+		if (byte_keys && table_size > UINT32_MAX) return failed("table geometry");
+		bool use_selector = selector != nullptr;
+		if (byte_keys && selector != nullptr) {
+			ShaderBufferResource material;
+			if (!DecodeBufferDescriptor(material_value, material)) return false;
+			use_selector = CanReadIndexedSelector(material, selector->offset);
+			if (use_selector && material.Format() != Prospero::BufferFormat::kInvalid &&
+			    material.NumRecords() != 0u && selector->offset < material.Stride()) {
+				const auto begin = material.Base48() + selector->offset;
+				const auto end = begin + uint64_t {material.NumRecords() - 1u} * material.Stride() + 4u;
+				for (uint32_t i = 0; i < program.info.buffers.size(); ++i) {
+					const auto& buffer = program.info.buffers[i];
+					if (!buffer.written) continue;
+					ShaderBufferResource written;
+					if (program.descriptor_sources[buffer.source].indirect_descriptor ||
+					    !DecodeBufferDescriptor(snapshot.buffers[i], written) ||
+					    (written.Base48() != 0u && written.GetSize() != 0u &&
+					     begin < written.Base48() + written.GetSize() &&
+					     written.Base48() < end)) {
+						use_selector = false;
+						break;
+					}
+				}
+			}
+		}
+		if (byte_keys && !use_selector) {
 			if (!indirect.table_scalar && table.Format() == Prospero::BufferFormat::kInvalid) {
 				descriptors[resource_index] = {.dword_count = dword_count};
 				return true;
 			}
-			if (table_size > UINT32_MAX ||
+			if (indirect.table_record_bytes != 0u &&
-			    (indirect.table_record_bytes != 0u &&
 			     (table.Type() != 0u || table.SwizzleEnabled() || table.AddTid() ||
 			      table.OutOfBounds() != 0u || (table_base & 3u) != 0u ||
-			      (table.Stride() & 3u) != 0u || table.Stride() < indirect.table_record_bytes)))
+			      (table.Stride() & 3u) != 0u || table.Stride() < indirect.table_record_bytes))
 				return failed("table geometry");
 			const auto* pointer = indirect.pointer_table ? &*indirect.pointer_table : nullptr;
 			const bool indexed_pointer = pointer != nullptr && pointer->indexed;
@@ -640,10 +687,11 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 				return failed("candidate count", fmt::format("rows={} limit={}", count, MaxIndirectDescriptorProbes));
 			// Affine SMEM offsets wrap in U32 before adding the aligned immediate.
 			// Image heaps retain one bulk snapshot and their absolute byte-offset mapping.
+			auto& heap_words = program.material_keys;
 			if (image_heap) {
-				keys.resize(table_size / 4u);
+				heap_words.resize(table_size / 4u);
-				if (!ReadScalarTable(table_base, table_size, 0u, runtime, keys))
+				if (!ReadScalarTable(table_base, table_size, 0u, runtime, heap_words))
-					return read_failed("image heap read", 0u, 0u, table_base, 0u, keys.size() * sizeof(uint32_t));
+					return read_failed("image heap read", 0u, 0u, table_base, 0u, heap_words.size() * sizeof(uint32_t));
 			}
 			const auto mapping_offset = snapshot.flattened_srt.size();
 			descriptors[resource_index] = {.dword_count = dword_count};
@@ -670,8 +718,8 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 						                     std::span(candidate.dwords).first(read_dwords))) return false;
 					} else {
 						const auto table_word = table_offset / 4u;
-						std::copy_n(keys.begin() + table_word,
+						std::copy_n(heap_words.begin() + table_word,
-						            std::min<size_t>(read_dwords, keys.size() - table_word),
+						            std::min<size_t>(read_dwords, heap_words.size() - table_word),
 						            candidate.dwords.begin());
 					}
 					if (!normalize(candidate))
@@ -722,21 +770,38 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 				selector_keys->resolved = true;
 			}
 		}
+		if (byte_keys) {
+			// Keep cached selector values independent of each destination table's
+			// U32 arithmetic, and use the same byte keys as the full-heap GPU path.
+			byte_offsets.clear();
+			byte_offsets.reserve(keys.size());
+			for (const auto key: keys) {
+				const auto offset = uint64_t {(key * indirect.table_stride + indirect.table_offset) & ~3u} +
+				                    (indirect.table_immediate & ~3u);
+				byte_offsets.push_back(offset <= UINT32_MAX ? static_cast<uint32_t>(offset) : UINT32_MAX);
+			}
+			std::ranges::sort(byte_offsets);
+			byte_offsets.erase(std::unique(byte_offsets.begin(), byte_offsets.end()), byte_offsets.end());
+		}
 	}
 
 	const auto mapping_offset = snapshot.flattened_srt.size();
-	const auto key_count = sources.empty() ? keys.size() : sources.size();
+	const auto selected_keys = byte_keys ? std::span<const uint32_t>(byte_offsets)
+	                                    : std::span<const uint32_t>(keys);
+	const auto key_count = sources.empty() ? selected_keys.size() : sources.size();
 	snapshot.flattened_srt.resize(mapping_offset + 1u + key_count * 2u);
 	snapshot.flattened_srt[mapping_offset] = static_cast<uint32_t>(key_count);
 	for (uint32_t entry = 0; entry < key_count; ++entry) {
-		const auto key = sources.empty() ? keys[entry] : entry;
+		const auto key = sources.empty() ? selected_keys[entry] : entry;
 		uint64_t table_offset = 0;
 		DescriptorValue candidate;
 		candidate.dword_count = dword_count;
 		if (sources.empty()) {
-			table_offset = uint64_t {(key * indirect.table_stride + indirect.table_offset) & ~3u} +
+			table_offset = byte_keys ? key
-			               (indirect.table_immediate & ~3u);
+			    : uint64_t {(key * indirect.table_stride + indirect.table_offset) & ~3u} +
-			if (!ReadScalarTable(table_base, table_size, table_offset, runtime,
+			          (indirect.table_immediate & ~3u);
+			if (!(byte_keys && key == UINT32_MAX) &&
+			    !ReadScalarTable(table_base, table_size, table_offset, runtime,
 			                     std::span(candidate.dwords).first(read_dwords))) {
 				return read_failed("table descriptor read", key, 0u, table_base, table_offset, descriptor_bytes);
 			}
@@ -746,7 +811,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 		if (!normalize(candidate))
 			return normalization_failed(table_offset, key, indirect.table_stride, "key");
 		uint32_t ordinal = 0;
-		if (entry == 0) {
+		if (entry == 0 && !byte_keys) {
 			descriptors[resource_index] = candidate;
 		} else if (!intern_candidate(candidate, ordinal)) {
 			return false;
@@ -913,6 +978,7 @@ static bool BuildResourceSpecialization(const ResourcePlan& program, ResourceSna
 			if (i >= program.info.images.size() && source != nullptr &&
 			    source->indirect_descriptor &&
 			    IsBoundedDescriptorTable(program, *source->indirect_descriptor) &&
+			    specialization.images[base_index].indirect_search_iterations == 0u &&
 			    !source->indirect_descriptor->pointer_table) {
 				const auto& indirect = *source->indirect_descriptor;
 				const size_t mapping = specialization.images[base_index].indirect_mapping_offset;
@@ -997,7 +1063,6 @@ static bool BuildResourceSpecialization(const ResourcePlan& program, ResourceSna
 			    is_1d_or_2d(image_class.dimension);
 			if ((!same_coordinates && !(is_2d(image.dimension) && is_2d(image_class.dimension)) &&
 			     !unfiltered_coordinates) ||
-			    image.mip_count != image_class.mip_count ||
 			    image.conversion_format != image_class.conversion_format ||
 			    (!storage_swizzle && image.shader_swizzle != image_class.shader_swizzle)) {
 				return SpecializationFail(

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

@@ -1558,14 +1558,49 @@ private:
 		return read;
 	}
 
+	const Inst* IndexedMaterialRead(const Inst* read, Value active,
+	                               DescriptorSource::IndirectDescriptor::SelectorRead& selector) const {
+		uint32_t lane = 0;
+		if (read != nullptr && read->GetOpcode() == ValueOpcode::CompositeExtractU32x2) {
+			if (!ImmediateU32(read->Arg(1), lane) || lane >= 2u) return nullptr;
+			read = read->Arg(0).Resolve().TryInstruction();
+			if (read == nullptr || read->GetOpcode() != ValueOpcode::LoadBufferU32x2) return nullptr;
+		} else if (read == nullptr || read->GetOpcode() != ValueOpcode::LoadBufferU32) {
+			return nullptr;
+		}
+		const auto& memory = m_program.memory_info[read->Flags<MemoryFlags>().index];
+		if (memory.kind != ResourceKind::Buffer || memory.typed || memory.formatted ||
+		    !memory.idxen || memory.offen || memory.offset > UINT32_MAX - lane * 4u ||
+		    read->Arg(2).Resolve() != Value(0u) || read->Arg(3).Resolve() != Value(0u) ||
+		    !EquivalentValue(m_program, active, read->Arg(4))) return nullptr;
+		selector.offset = memory.offset + lane * 4u;
+		return read;
+	}
+
 	bool MatchUniformizedMaterialKey(Value key, const Inst& image,
 	                                DescriptorSource::IndirectDescriptor& indirect,
-	                                DescriptorSource& material_source) {
+	                                DescriptorSource& material_source, bool indexed) {
 		const auto* selected = UniformizedMaterialValue(key, image);
 		if (selected == nullptr) return false;
 		const auto active = selected->Arg(0).Resolve();
 		DescriptorSource::IndirectDescriptor::SelectorRead selector;
 		const auto* read = SelectorReadValue(selected->Arg(1), active, selector);
+		if (indexed) {
+			read = IndexedMaterialRead(read, active, selector);
+			// This optional domain is checked against the dispatch's writable buffers.
+			// Other stages and unbounded/image writes need the original heap domain.
+			if (read == nullptr || m_program.stage != ShaderType::Compute ||
+			    std::ranges::any_of(m_program.blocks, [](const Block* block) {
+				    return std::ranges::any_of(*block, [](const Inst& inst) {
+					    const auto op = inst.GetOpcode();
+					    return ImageOpcodeInfoOf(op).resource_class == ImageResourceClass::Storage ||
+					           AddressOpcodeInfoOf(op).access == AddressAccess::Write;
+				    });
+			    }) || !MakeRuntimeTableSource(*read, material_source)) return false;
+			selector.source = InternSource(material_source);
+			indirect.selector = std::move(selector);
+			return true;
+		}
 		if (read == nullptr || read->GetOpcode() != ValueOpcode::LoadAddressU32 ||
 		    read->NumArgs() != 4u || !Implies(active, read->Arg(3)))
 			return false;
@@ -2971,6 +3006,13 @@ private:
 			// 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 (!sampler && !plan.table_indexed && table_source.dword_count == 4u &&
+		           MatchUniformizedMaterialKey(key, handle, indirect, material_source, true)) {
+			// Preserve the heap's byte-offset key so runtime geometry can choose the
+			// tighter selector domain without changing the GPU's descriptor reads.
+			indirect.table_scalar = true;
+			plan.table_read = plan.reads[0];
+			plan.retain_reads = true;
 		} else if (sampler && table_source.dword_count == 4u && !plan.table_indexed &&
 		           SetExactTableMaximum(key, indirect, &handle)) {
 			indirect.table_scalar = true;
@@ -3020,7 +3062,7 @@ private:
 				}
 			}
 			if (indirect.key_count.IsEmpty() &&
-			    !MatchUniformizedMaterialKey(key, handle, indirect, material_source) &&
+			    !MatchUniformizedMaterialKey(key, handle, indirect, material_source, false) &&
 			    !MatchReversedImageRange(key, indirect)) {
 				const auto maximum = SelectorMaximum(key);
 				if (maximum >= MaxIndirectDescriptorProbes) return false;
@@ -3064,6 +3106,10 @@ private:
 					indirect.selector_mask = selected.selector_mask;
 					indirect.signed_count = selected.signed_count;
 					indirect.workgroup_axis = selected.workgroup_axis;
+					if (selected.table_scalar) {
+						indirect.table_scalar = true;
+						plan.table_read = plan.reads[0];
+					}
 					material_source = m_sources[selected.selector->source];
 					shared_selector = true;
 					break;
@@ -3264,26 +3310,13 @@ private:
 			if (value == nullptr || (selected != nullptr && selected != value)) return false;
 			selected = value;
 		}
-		const auto* read = selected->Arg(1).Resolve().TryInstruction();
+		DescriptorSource::IndirectDescriptor::SelectorRead selector;
-		uint32_t lane = 0;
+		const auto* read = IndexedMaterialRead(selected->Arg(1).Resolve().TryInstruction(),
-		if (read != nullptr && read->GetOpcode() == ValueOpcode::CompositeExtractU32x2) {
+		                                       selected->Arg(0), selector);
-			if (!ImmediateU32(read->Arg(1), lane) || lane >= 2u) return false;
+		if (read == nullptr || !MakeRuntimeTableSource(*read, material_source) ||
-			read = read->Arg(0).Resolve().TryInstruction();
-			if (read == nullptr || read->GetOpcode() != ValueOpcode::LoadBufferU32x2) return false;
-		} else if (read == nullptr || read->GetOpcode() != ValueOpcode::LoadBufferU32) {
-			return false;
-		}
-		const auto& memory = m_program.memory_info[read->Flags<MemoryFlags>().index];
-		uint32_t offset = 1u, scalar = 1u;
-		if (memory.kind != ResourceKind::Buffer || memory.typed || memory.formatted ||
-		    !memory.idxen || memory.offen || memory.offset > UINT32_MAX - lane * 4u ||
-		    !ImmediateU32(read->Arg(2), offset) || offset != 0u ||
-		    !ImmediateU32(read->Arg(3), scalar) || scalar != 0u ||
-		    !EquivalentValue(m_program, selected->Arg(0), read->Arg(4))) return false;
-		if (!MakeRuntimeTableSource(*read, material_source) ||
 		    !BoundedSelectedIndex(read->Arg(1), read->Arg(4), indirect)) return false;
-		indirect.selector.emplace(DescriptorSource::IndirectDescriptor::SelectorRead{.source = InternSource(material_source),
+		selector.source = InternSource(material_source);
-		                              .offset = memory.offset + lane * 4u});
+		indirect.selector = std::move(selector);
 		return finish();
 	}
 

tests/ResourceTrackingTests.cpp

@@ -593,8 +593,9 @@ void TestExactUnfilteredImageCandidates() {
     second[3] |= 1u << 16u;
     fill();
     plan.info.images[0].mip_mode = ImageMipMode::Dynamic;
-    Check(!MaterializeResources(plan, runtime, snapshot, specialization),
+    Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
-          "mixed-dimension acceptance discarded mip-count compatibility");
+              specialization.images[0].mip_count == 1u && specialization.images[1].mip_count == 2u,
+          "mixed-dimension candidates lost their individual mip counts");
   }
   // The same image may be merged from several uses; one query keeps the guard strict.
   auto merged = MakeIndirectImageFixture(false, 32u, 0u, 0u, false, ValueOpcode::ImageRead, true);
@@ -3164,6 +3165,237 @@ void TestUniformizedMaterialImageKeys() {
                "packed selector accepted an unproved OR guard, offset domain, or field extraction");
 }
 
+void TestIndexedMaterialImageKeys() {
+  namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
+  const auto make_plan = [](ShaderType stage = ShaderType::Compute, bool shared = false) {
+    Fixture f(stage);
+    auto *entry = f.block, *body = f.AddBlock();
+    auto *later = shared ? f.AddBlock() : nullptr;
+    auto *done = f.AddBlock();
+    entry->AddBranch(body);
+    entry->AddBranch(done);
+    body->AddBranch(shared ? later : done);
+    if (shared) later->AddBranch(done);
+    const auto lane = f.Emit(ValueOpcode::GetBuiltin,
+        {Value(static_cast<uint32_t>(StageInputKind::LocalInvocationId)), Value(0u)});
+    const auto active = f.Emit(ValueOpcode::INotEqual32, {lane, Value(UINT32_MAX)});
+    entry->condition = f.Emit(ValueOpcode::ConditionRef, {active}, CFG::BranchCondition::ExecNonZero);
+    entry->terminator = {.kind = CFG::TerminatorKind::ConditionalBranch,
+                         .true_block = body, .false_block = done};
+    body->terminator = {.kind = CFG::TerminatorKind::Branch, .true_block = shared ? later : done};
+    if (shared) later->terminator = {.kind = CFG::TerminatorKind::Branch, .true_block = done};
+    done->terminator.kind = CFG::TerminatorKind::Return;
+    const auto material = f.Buffer({f.UserData(0), f.UserData(1), f.UserData(2), f.UserData(3)});
+    const auto heap = f.Buffer({f.UserData(4), f.UserData(5), f.UserData(6), f.UserData(7)});
+    const auto output = f.Buffer({f.UserData(8), f.UserData(9), f.UserData(10), f.UserData(11)});
+    const auto sampler = f.Sampler({Value(0u), Value(0u), Value(1u << 26u), Value(0u)});
+    const auto address = f.ImageAddress();
+    Value owner_key, shared_image;
+    for (uint32_t index = 0; index < (shared ? 3u : 2u); ++index) {
+      const uint32_t field = index == 1u ? 68u : 64u;
+      auto *use = index == 2u ? later : body;
+      auto key = owner_key;
+      if (index < 2u) {
+        MemoryInfo memory;
+        memory.kind = ResourceKind::Buffer;
+        memory.idxen = true;
+        memory.offset = field;
+        const auto loaded = f.Emit(ValueOpcode::LoadBufferU32,
+            {material, lane, Value(0u), Value(0u), active}, f.AddMemory(memory, field), entry);
+        const auto selected = f.Emit(ValueOpcode::SelectU32, {active, loaded, Value(0xdeadbeefu)}, 0, entry);
+        key = shared && index == 0u
+            ? f.Emit(ValueOpcode::ReadLane, {selected, Value(0u)}, 0, entry)
+            : f.Emit(ValueOpcode::ReadFirstLane, {selected, active}, 0, body);
+        if (shared && index == 0u) {
+          owner_key = key;
+          const auto equal = f.Emit(ValueOpcode::IEqual32, {key, selected}, 0, entry);
+          const auto guard = f.Emit(ValueOpcode::LogicalAnd, {active, equal}, 0, entry);
+          entry->condition = f.Emit(ValueOpcode::ConditionRef, {guard}, CFG::BranchCondition::ExecNonZero, entry);
+        }
+      }
+      const auto offset = f.Emit(ValueOpcode::ShiftLeftLogical32, {key, Value(5u)}, 0, use);
+      std::array<Value, 8> words;
+      for (uint32_t word = 0; word < words.size(); ++word) {
+        MemoryInfo scalar;
+        scalar.kind = ResourceKind::ScalarBuffer;
+        scalar.offset = word * 4u;
+        words[word] = f.Emit(ValueOpcode::ReadConstBuffer, {heap, offset},
+                            f.AddMemory(scalar, 0x100u + field), use);
+      }
+      const auto image_handle = [&] {
+        return f.Emit(ValueOpcode::GetImageResource,
+            {words[0], words[1], words[2], words[3], words[4], words[5], words[6], words[7]}, 0, use);
+      };
+      MemoryInfo gather;
+      gather.kind = ResourceKind::Image;
+      gather.image_dimension = Decoder::ImageDimension::Dim2D;
+      if (index == 0u)
+        f.Emit(ValueOpcode::ImageQueryDimensions, {image_handle(), address}, f.AddMemory(gather, 0x180u), body);
+      const auto image = image_handle();
+      if (index == 2u) shared_image = image;
+      gather.image_sample_flags = Decoder::ImageSampleFlagLod;
+      f.Emit(ValueOpcode::ImageGatherRaw,
+          {image, sampler, address, words[1], words[3], Value(0u), Value(1u << 26u), active},
+          f.AddMemory(gather, 0x200u + field), use);
+    }
+    MemoryInfo store;
+    store.kind = ResourceKind::Buffer;
+    f.Emit(ValueOpcode::StoreBufferU32,
+        {output, Value(0u), Value(0u), Value(0u), Value(1u), Value(true)},
+        f.AddMemory(store, 0x300u), done);
+    f.PlanAndTrack();
+    if (shared) {
+      owner_key.ResolveInstruction()->ReplaceUsesWith(Value(12994u));
+      uint32_t offset = 0u;
+      Check(SrtWalker(f.program, {}).Evaluate(shared_image.ResolveInstruction()->Arg(0), offset) &&
+                offset == 12994u * 32u,
+            "inherited indexed selector changed GPU byte offsets into raw material keys");
+    }
+    return ExtractResourcePlan(f.program);
+  };
+  auto plan = make_plan();
+  Check(plan.info.images.size() == 3u && plan.selector_keys.size() == 2u,
+        "indexed material fields did not share their query/gather selector domain");
+  for (const auto &image : plan.info.images) {
+    const auto &source = *plan.descriptor_sources[image.source].indirect_descriptor;
+    Check(source.table_scalar && source.selector && source.selector->stride == 0u &&
+              source.key_count.IsEmpty() && source.selector_first.IsEmpty(),
+          "indexed material field fell back to an unconditional heap scan");
+  }
+  constexpr uint32_t base = 0x1000u, heap = 0x2000u, records = 14517u;
+  struct Memory {
+    std::vector<uint32_t> words = std::vector<uint32_t>((heap - base + records * 32u) / 4u);
+    uint32_t selectors = 0, descriptors = 0;
+    bool narrow = true, unreadable = false;
+  } memory;
+  const auto put = [&](uint32_t key, uint32_t mips) {
+    auto *words = memory.words.data() + (heap - base + key * 32u) / 4u;
+    words[0] = 0x100u + key;
+    words[1] = static_cast<uint32_t>(Libs::Graphics::Prospero::BufferFormat::k16UNorm) << 20u;
+    words[3] = (9u << 28u) | ((mips - 1u) << 16u) | Libs::Graphics::DstSel(4, 5, 6, 7);
+  };
+  for (uint32_t key = 0; key < 64u; ++key) put(key, 1u);
+  put(12994u, 8u);
+  put(12996u, 8u);
+  for (uint32_t row = 0; row < 10u; ++row) {
+    memory.words[(row * 96u + 64u) / 4u] = 12994u;
+    memory.words[(row * 96u + 68u) / 4u] = 12996u;
+  }
+  std::array<uint32_t, 12> user_data{base, 96u << 16u, 10u, 0x5204u,
+      heap, 32u << 16u, records, 0x5204u, 0x90000u, 4u << 16u, 16u, 0x5204u};
+  const auto read = +[](void *data, uint64_t address, std::span<uint32_t> words) {
+    auto &m = *static_cast<Memory *>(data);
+    if (address < base || address - base + words.size_bytes() > m.words.size() * 4u) return false;
+    if (address < heap) {
+      ++m.selectors;
+      if (m.unreadable) return false;
+    } else {
+      ++m.descriptors;
+      const auto key = (address - heap) / 32u;
+      if (m.narrow && (words.size() != 8u || (key != 0u && key != 12994u && key != 12996u))) return false;
+    }
+    std::copy_n(m.words.begin() + (address - base) / 4u, words.size(), words.begin());
+    return true;
+  };
+  const SrtRuntime runtime{.user_data = user_data, .userdata = &memory,
+                           .read_specialization_memory = read};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+            memory.selectors == 2u && memory.descriptors == 6u && snapshot.images.size() == 9u,
+        "indexed selector scanned the broad heap or did not share its captured keys");
+  for (uint32_t root = 0; root < plan.info.images.size(); ++root) {
+    const auto &image = specialization.images[root];
+    const auto offset = image.indirect_mapping_offset;
+    Check(snapshot.images[root].dwords[0] == 0u && image.indirect_search_iterations == 2u &&
+              snapshot.flattened_srt[offset] == 2u && snapshot.flattened_srt[offset + 1u] == 0u &&
+              snapshot.flattened_srt[offset + 3u] == (root == 2u ? 12996u : 12994u) * 32u,
+          "byte-key selector lost native OOB zero, sorted offsets, or the null default");
+  }
+  Check(std::ranges::count_if(specialization.images, [](const auto &image) {
+          return image.mip_count == 8u;
+        }) == 2u, "indirect gathers lost candidate-specific mip counts");
+  for (const uint32_t field : {64u, 68u})
+    Check(std::ranges::find(snapshot.specialization_reads,
+              std::pair<uint64_t, uint64_t>{base + field, 868u}) != snapshot.specialization_reads.end(),
+          "indexed selector batch escaped writable-alias tracking");
+  auto shared = make_plan(ShaderType::Compute, true);
+  Check(MaterializeResources(shared, runtime, snapshot, specialization) && snapshot.images.size() == 9u,
+        "inherited indexed selector lost its bounded material domain");
+  memory.words[64u / 4u] = 12996u;
+  memory.selectors = 0;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+            memory.selectors == 2u && snapshot.images.size() == 11u,
+        "changed material rows reused stale selector keys");
+  memory.words[64u / 4u] = 12994u;
+  memory.unreadable = true;
+  Check(!MaterializeResources(plan, runtime, snapshot, specialization),
+        "unreadable eligible selector silently fell back to the heap");
+  memory.narrow = false;
+  const auto original = user_data;
+  for (const uint32_t variant : {0u, 1u, 2u, 3u, 4u, 5u, 6u, 7u}) {
+    user_data = original;
+    if (variant == 0u) user_data[1] |= 1u << 31u;
+    if (variant == 1u) user_data[3] |= 1u << 23u;
+    if (variant == 2u) user_data[3] |= 1u << 28u;
+    if (variant == 3u) user_data[0] += 1u;
+    if (variant == 4u) user_data[2] = MaxIndirectDescriptorProbes + 1u;
+    if (variant == 5u) user_data[8] = base + 68u;
+    if (variant == 6u) user_data[1] = 66u << 16u;
+    if (variant == 7u) user_data[0] = 0u;
+    memory.selectors = 0;
+    Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+              memory.selectors == 0u && snapshot.images.size() > 100u,
+          "unsupported geometry or writable selector lost the original heap domain");
+  }
+  user_data = original;
+  auto graphics = make_plan(ShaderType::Pixel);
+  Check(MaterializeResources(graphics, runtime, snapshot, specialization) && memory.selectors == 0u,
+        "graphics material lookup acquired unproved selector read dependencies");
+  memory.unreadable = false;
+  for (uint32_t row = 0; row < 10u; ++row)
+    memory.words[(row * 96u + 64u) / 4u] = 0x08000000u; // U32 key*32 wraps to zero.
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) && snapshot.images.size() == 7u,
+        "indexed selector lost wrapped scalar descriptor offsets");
+  for (auto &source : plan.descriptor_sources) {
+    if (!source.indirect_descriptor) continue;
+    source.indirect_descriptor->table_offset = 0xffffffe0u;
+    source.indirect_descriptor->table_immediate = 32u;
+  }
+  for (uint32_t row = 0; row < 10u; ++row) {
+    memory.words[(row * 96u + 64u) / 4u] = row == 0u ? 1u : 0u;
+    memory.words[(row * 96u + 68u) / 4u] = 0u;
+  }
+  memory.selectors = memory.descriptors = 0u;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+            memory.selectors == 2u && memory.descriptors == 2u && snapshot.images.size() == 5u,
+        "overflowed scalar immediate read a descriptor or lost valid wrapped keys");
+  for (uint32_t root = 0; root < 2u; ++root) {
+    const auto offset = specialization.images[root].indirect_mapping_offset;
+    Check(snapshot.images[root].dwords[0] == 0u && snapshot.flattened_srt[offset] == 2u &&
+              snapshot.flattened_srt[offset + 1u] == 32u && snapshot.flattened_srt[offset + 3u] == UINT32_MAX &&
+              snapshot.flattened_srt[offset + 4u] == 0u,
+          "invalid byte key selected a real texture instead of the null root");
+  }
+  for (auto &source : plan.descriptor_sources) {
+    if (!source.indirect_descriptor) continue;
+    source.indirect_descriptor->table_offset = source.indirect_descriptor->table_immediate = 0u;
+  }
+  memory.narrow = true;
+  memory.unreadable = true;
+  for (const uint32_t variant : {0u, 1u, 2u, 3u}) {
+    user_data = original;
+    if (variant == 0u) user_data[2] = 0u;
+    if (variant == 1u) user_data[1] = 0u;
+    if (variant == 2u) user_data[3] = 0u;
+    if (variant == 3u) user_data[1] = 64u << 16u;
+    memory.selectors = 0;
+    Check(MaterializeResources(plan, runtime, snapshot, specialization) && memory.selectors == 0u &&
+              snapshot.images.size() == 6u,
+          "native zero/OOB indexed load read its unavailable payload");
+  }
+}
+
 void TestImageDescriptorFields() {
   constexpr std::array<std::pair<uint32_t, uint32_t>, 5> reserved{
       {{1u, 0x20000000u}, {2u, 0x70003000u}, {4u, 0xe000e000u},
@@ -8366,6 +8598,7 @@ int main() {
     Run("guarded direct image table", TestGuardedDirectImageTable);
     Run("bounded compute image loop", TestBoundedComputeImageLoop);
     Run("uniformized material image keys", TestUniformizedMaterialImageKeys);
+    Run("indexed material image keys", TestIndexedMaterialImageKeys);
     Run("image descriptor fields", TestImageDescriptorFields);
     Run("draw-uniform scalar image", TestUniformScalarBufferImage);
     Run("compact scalar image table", TestCompactScalarImageTable);

tests/ShaderRecompilerComputeTests.cpp

@@ -40392,12 +40392,13 @@ void CheckIndirectImageKeySwitch(VulkanHarness &vulkan) {
         resource == 4u ? vk::ImageViewType::e2DArray : vk::ImageViewType::e2D, layers));
   }
   const auto native_sampler = vulkan.CreateSampler(test.name);
-  enum class Access { Sample, Gather, GatherHomogeneous, Fetch, Fetch1D };
+  enum class Access { Sample, Gather, GatherHomogeneous, GatherMips, Fetch, Fetch1D };
   auto *access = &sample;
   for (const auto mode : {Access::Sample, Access::Gather,
-                          Access::GatherHomogeneous, Access::Fetch, Access::Fetch1D}) {
+                          Access::GatherHomogeneous, Access::GatherMips, Access::Fetch, Access::Fetch1D}) {
     const bool fetch = mode == Access::Fetch || mode == Access::Fetch1D;
-    const bool gather = mode == Access::Gather || mode == Access::GatherHomogeneous;
+    const bool gather = mode == Access::Gather || mode == Access::GatherHomogeneous || mode == Access::GatherMips;
+    const bool mixed_mips = mode == Access::GatherMips;
     const bool homogeneous = mode == Access::GatherHomogeneous;
     if (mode == Access::Gather) {
       test.name = "IndirectImageGatherContiguousRuns";
@@ -40426,7 +40427,21 @@ void CheckIndirectImageKeySwitch(VulkanHarness &vulkan) {
                                      : root.indirect_resources;
     if (homogeneous) test.name = "IndirectImageGatherHomogeneous";
     compiled.resources.flattened_srt[6] = homogeneous ? 4u : ordinals[5];
+    if (mixed_mips) {
+      test.name = "IndirectImageGatherMixedMips";
+      compiled.program.info.images[3].mip_mode = ImageMipMode::Dynamic;
+      compiled.program.info.images[3].mip_count = 2u;
+      compiled.program.memory_info[0].image_sample_flags = ShaderRecompiler::Decoder::ImageSampleFlagLod;
+      address.SetArg(0, Value(std::bit_cast<u32>(0.75f)));
+      address.SetArg(1, Value(std::bit_cast<u32>(0.25f)));
+      address.SetArg(2, Value(0u));
+      address.SetArg(3, Value(std::bit_cast<u32>(1.0f)));
+      access->SetArg(5, Value(0xfffu << 12u));
+      access->SetArg(6, Value(1u << 26u));
+    }
     if (fetch) {
+      compiled.program.info.images[3].mip_mode = ImageMipMode::None;
+      compiled.program.info.images[3].mip_count = 1u;
       test.name = "IndirectImageFetchContiguousRuns";
       const auto position = std::ranges::find_if(*block, [&](const Inst &inst) {
         return &inst == access;
@@ -40484,12 +40499,15 @@ void CheckIndirectImageKeySwitch(VulkanHarness &vulkan) {
       const auto image_opcode = fetch ? "OpImageFetch" : gather ? "OpImageGather"
                                                           : "OpImageSampleExplicitLod";
       Require(test.name, "native array image runs", tools.Disassemble(compiled.spirv, &text) &&
-                  CountText(text, image_opcode) == (homogeneous ? 1u : 2u) * halves,
+                  CountText(text, image_opcode) == (homogeneous ? 1u : mixed_mips ? 4u : 2u) * halves,
               "mixed runs expanded into per-image accesses");
       if (homogeneous) {
         Require(test.name, "homogeneous gather array", text.find("OpSwitch") == std::string::npos,
                 "homogeneous gather emitted a descriptor switch");
       }
+      if (mixed_mips)
+        Require(test.name, "mixed mip emitted size", compiled.spirv.size() <= 8192u * halves,
+                "mixed mip selection exceeded its 32 KiB per-wave regression budget");
       if (fetch || gather) {
         std::vector<u32> decorated;
         std::vector<u32> image_operands;
@@ -40513,7 +40531,8 @@ void CheckIndirectImageKeySwitch(VulkanHarness &vulkan) {
       test.expected.assign(wave_size, 0u);
       for (u32 lane = 0; lane < ordinals.size(); ++lane) {
         const auto resource = selected_resources[compiled.resources.flattened_srt[1u + lane]];
-        test.expected[lane] = std::bit_cast<u32>(float(resource + (fetch && resource != 0u ? 10u : 0u)));
+        test.expected[lane] = std::bit_cast<u32>(float(resource +
+            ((fetch && resource != 0u) || (mixed_mips && resource == 3u) ? 10u : 0u)));
       }
       auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size());
       vulkan.Dispatch(test, compiled, buffer, nullptr, nullptr, nullptr, nullptr,
@@ -40521,7 +40540,7 @@ void CheckIndirectImageKeySwitch(VulkanHarness &vulkan) {
       const auto actual = vulkan.ReadBuffer(test.name, buffer, test.expected.size());
       vulkan.DestroyBuffer(&buffer);
       CompareWords(test, "root, native slots, array runs and mapping bounds", test.expected, actual);
-      if (wave_size == 32u && !homogeneous) {
+      if (wave_size == 32u && !homogeneous && !mixed_mips) {
         // A single integer candidate must keep raw register bits across the
         // same descriptor runs used by samples, gathers, and image fetches.
         const auto coordinate_conversions = CountText(text, "OpConvertUToF");