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 ||
-			    (indirect.table_record_bytes != 0u &&
+			if (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);
-				if (!ReadScalarTable(table_base, table_size, 0u, runtime, keys))
-					return read_failed("image heap read", 0u, 0u, table_base, 0u, keys.size() * sizeof(uint32_t));
+				heap_words.resize(table_size / 4u);
+				if (!ReadScalarTable(table_base, table_size, 0u, runtime, heap_words))
+					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::min<size_t>(read_dwords, keys.size() - table_word),
+						std::copy_n(heap_words.begin() + 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} +
-			               (indirect.table_immediate & ~3u);
-			if (!ReadScalarTable(table_base, table_size, table_offset, runtime,
+			table_offset = byte_keys ? key
+			    : uint64_t {(key * indirect.table_stride + indirect.table_offset) & ~3u} +
+			          (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();
-		uint32_t lane = 0;
-		if (read != nullptr && read->GetOpcode() == ValueOpcode::CompositeExtractU32x2) {
-			if (!ImmediateU32(read->Arg(1), lane) || lane >= 2u) return false;
-			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) ||
+		DescriptorSource::IndirectDescriptor::SelectorRead selector;
+		const auto* read = IndexedMaterialRead(selected->Arg(1).Resolve().TryInstruction(),
+		                                       selected->Arg(0), selector);
+		if (read == nullptr || !MakeRuntimeTableSource(*read, material_source) ||
 		    !BoundedSelectedIndex(read->Arg(1), read->Arg(4), indirect)) return false;
-		indirect.selector.emplace(DescriptorSource::IndirectDescriptor::SelectorRead{.source = InternSource(material_source),
-		                              .offset = memory.offset + lane * 4u});
+		selector.source = InternSource(material_source);
+		indirect.selector = std::move(selector);
 		return finish();
 	}