Bound uniformized image keys by guarded material mask

Co-authored-by: TarkusR <90155059+TarkusR@users.noreply.github.com>

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

@@ -465,6 +465,7 @@ struct DescriptorSource {
 		uint32_t selector_offset = 0;
 		uint32_t table_offset    = 0;
 		Value    key_count;
+		Value    selector_mask;
 
 		bool operator==(const IndirectImage& other) const = default;
 	};
@@ -540,6 +541,7 @@ struct ResourcePlan {
 	std::vector<SrtRead>                srt_reads;
 	std::vector<uint8_t>                clean_flat_slots;
 	bool                                requires_specialization_memory = false;
+	bool                                has_address_writes = false;
 	bool                                srt_plan_complete          = false;
 	bool                                resource_tracking_complete = false;
 	ShaderInfo                          info;
@@ -552,6 +554,7 @@ struct ResourcePlan {
 	mutable std::vector<uint8_t>            visited_blocks;
 	mutable std::vector<uint32_t>           pending_blocks;
 	mutable std::vector<uint32_t>           material_keys;
+	mutable std::vector<std::pair<uint64_t, uint64_t>> specialization_reads;
 };
 
 struct Program: ResourcePlan {

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

@@ -138,6 +138,50 @@ bool DecodeBufferDescriptor(const DescriptorValue& descriptor, ShaderBufferResou
 	return true;
 }
 
+struct ReadCapture {
+	SrtRuntime                                  source;
+	std::vector<std::pair<uint64_t, uint64_t>>& ranges;
+};
+
+bool CaptureStrictRead(void* userdata, uint64_t address, std::span<uint32_t> values) {
+	auto& capture = *static_cast<ReadCapture*>(userdata);
+	if (!capture.source.read_specialization_memory(capture.source.userdata, address, values)) {
+		return false;
+	}
+	capture.ranges.emplace_back(address, values.size_bytes());
+	return true;
+}
+
+bool CaptureOrdinaryRead(void* userdata, uint64_t address, std::span<uint32_t> values) {
+	auto& capture = *static_cast<ReadCapture*>(userdata);
+	if (!capture.source.read_memory(capture.source.userdata, address, values)) {
+		return false;
+	}
+	capture.ranges.emplace_back(address, values.size_bytes());
+	return true;
+}
+
+bool WrittenBuffersDisjoint(const ResourcePlan& program, const ResourceSnapshot& snapshot,
+                            std::span<const std::pair<uint64_t, uint64_t>> reads) {
+	for (uint32_t i = 0; i < program.info.buffers.size(); ++i) {
+		if (!program.info.buffers[i].written) continue;
+		ShaderBufferResource buffer;
+		if (!DecodeBufferDescriptor(snapshot.buffers[i], buffer)) return false;
+		const auto base = buffer.Base48();
+		const auto size = buffer.GetSize();
+		if (size == 0u) continue;
+		if (size - 1u > AddressMask - base) return false;
+		for (const auto [address, bytes]: reads) {
+			if (bytes == 0u) continue;
+			if (address > AddressMask || bytes - 1u > AddressMask - address ||
+			    (address <= base ? base - address < bytes : address - base < size)) {
+				return false;
+			}
+		}
+	}
+	return true;
+}
+
 const DescriptorSource* Source(const ResourcePlan& program, uint32_t source) {
 	if (source >= program.descriptor_sources.size()) {
 		return nullptr;
@@ -146,12 +190,15 @@ const DescriptorSource* Source(const ResourcePlan& program, uint32_t source) {
 }
 
 void MarkCleanFlatSlots(const ResourcePlan& program, const DescriptorSource* source,
-                        std::span<uint8_t> slots) {
-	if (source == nullptr) {
+                        std::span<uint8_t> slots, Value extra = {}) {
+	if (source == nullptr && extra.IsEmpty()) {
 		return;
 	}
-	std::vector<Value>       pending(source->dwords.begin(),
-	                                 source->dwords.begin() + source->dword_count);
+	std::vector<Value>       pending;
+	if (source != nullptr) {
+		pending.assign(source->dwords.begin(), source->dwords.begin() + source->dword_count);
+	}
+	if (!extra.IsEmpty()) pending.push_back(extra);
 	std::vector<const Inst*> visited;
 	while (!pending.empty()) {
 		auto value = pending.back().Resolve();
@@ -225,6 +272,31 @@ bool MaterializeIndirectImage(const ResourcePlan& program,
 		}
 		keys.resize(key_count);
 		std::iota(keys.begin(), keys.end(), 0u);
+	} else if (!indirect.selector_mask.IsEmpty()) {
+		uint32_t mask = 0;
+		uint32_t count = 0;
+		if (material_value.dword_count != 2u || table_value.dword_count != 2u ||
+		    !clean.Evaluate(indirect.selector_mask, mask) ||
+		    !clean.Evaluate(indirect.key_count, count) || count == 0u || count > 32u) {
+			return false;
+		}
+		if (count < 32u) mask &= (1u << count) - 1u;
+		const auto material_base =
+		    (static_cast<uint64_t>(material_value.dwords[1]) << 32u) | material_value.dwords[0];
+		keys.reserve(std::popcount(mask));
+		while (mask != 0u) {
+			const auto index = std::countr_zero(mask);
+			const auto offset = static_cast<uint64_t>(indirect.selector_offset) +
+			                    static_cast<uint64_t>(index) * indirect.selector_stride;
+			if (offset > UINT32_MAX) return false;
+			uint32_t key = 0;
+			if (!ReadScalarTable(material_base, UINT64_MAX, static_cast<uint32_t>(offset),
+			                     runtime, {&key, 1})) return false;
+			keys.push_back(key);
+			mask &= mask - 1u;
+		}
+		std::ranges::sort(keys);
+		keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
 	} else {
 		ShaderBufferResource material;
 		if (!DecodeBufferDescriptor(material_value, material) || table_value.dword_count != 4u ||
@@ -784,6 +856,7 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
 	plan.memory_info                = program.memory_info;
 	plan.srt_plan_complete          = program.srt_plan_complete;
 	plan.resource_tracking_complete = program.resource_tracking_complete;
+	plan.has_address_writes         = program.has_address_writes;
 
 	std::unordered_map<const Inst*, Inst*> cloned;
 	std::function<Value(Value)>            Clone = [&](Value value) -> Value {
@@ -823,6 +896,7 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
 		target.indirect_image = source.indirect_image;
 		if (target.indirect_image.has_value()) {
 			target.indirect_image->key_count = Clone(target.indirect_image->key_count);
+			target.indirect_image->selector_mask = Clone(target.indirect_image->selector_mask);
 		}
 		for (uint32_t dword = 0; dword < source.dword_count; dword++) {
 			target.dwords[dword] = Clone(source.dwords[dword]);
@@ -848,7 +922,7 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
 		}
 		plan.requires_specialization_memory = true;
 		MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->material_source),
-		                   plan.clean_flat_slots);
+		                   plan.clean_flat_slots, source->indirect_image->selector_mask);
 		MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->table_source),
 		                   plan.clean_flat_slots);
 	}
@@ -861,8 +935,29 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 	    (program.requires_specialization_memory && runtime.read_specialization_memory == nullptr)) {
 		return false;
 	}
-	SrtWalker clean(program, CleanRuntime(runtime));
-	SrtWalker walker(program, runtime, program.clean_flat_slots, &clean);
+	const bool masked_image = std::ranges::any_of(program.info.images, [&](const auto& image) {
+		const auto* source = Source(program, image.source);
+		return source != nullptr && source->indirect_image.has_value() &&
+		       !source->indirect_image->selector_mask.IsEmpty();
+	});
+	if (masked_image &&
+	    (program.has_address_writes ||
+	     std::ranges::any_of(program.info.images, &ImageResource::written))) {
+		return false;
+	}
+	const bool capture_reads = masked_image &&
+	    std::ranges::any_of(program.info.buffers, &BufferResource::written);
+	auto& reads = program.specialization_reads;
+	ReadCapture capture {runtime, reads};
+	SrtRuntime observed = runtime;
+	if (capture_reads) {
+		reads.clear();
+		observed.userdata = &capture;
+		observed.read_specialization_memory = CaptureStrictRead;
+		if (observed.read_memory != nullptr) observed.read_memory = CaptureOrdinaryRead;
+	}
+	SrtWalker clean(program, CleanRuntime(observed));
+	SrtWalker walker(program, observed, program.clean_flat_slots, &clean);
 	const auto active = clean.FindActiveSources();
 	if (!walker.RefreshFlatBuffer(snapshot.flattened_srt)) {
 		return false;
@@ -896,6 +991,15 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 			return false;
 		}
 	}
+	if (capture_reads) {
+		for (uint32_t i = 0; i < program.info.buffers.size(); ++i) {
+			const auto& buffer = program.info.buffers[i];
+			if (!buffer.written || (!active.empty() && !active[buffer.source])) continue;
+			DescriptorValue strict;
+			if (!clean.EvaluateDescriptor(buffer.source, strict) ||
+			    strict != snapshot.buffers[i]) return false;
+		}
+	}
 	snapshot.images.resize(program.info.images.size());
 	specialization.images.clear();
 	specialization.images.reserve(program.info.images.size());
@@ -929,7 +1033,7 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 			if ((indirect.material_source != UINT32_MAX &&
 			     !clean.EvaluateDescriptor(indirect.material_source, material)) ||
 			    !clean.EvaluateDescriptor(indirect.table_source, table) ||
-			    !MaterializeIndirectImage(program, indirect, material, table, i, runtime, clean, snapshot,
+			    !MaterializeIndirectImage(program, indirect, material, table, i, observed, clean, snapshot,
 			                              specialization)) {
 				return false;
 			}
@@ -948,6 +1052,7 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 			return false;
 		}
 	}
+	if (capture_reads && !WrittenBuffersDisjoint(program, snapshot, reads)) return false;
 	snapshot.user_data.assign(runtime.user_data.begin(), runtime.user_data.end());
 	return BuildResourceSpecialization(program, snapshot, specialization);
 }

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

@@ -322,7 +322,10 @@ private:
 				if (a.material_source != b.material_source || a.table_source != b.table_source ||
 				    a.selector_stride != b.selector_stride || a.selector_offset != b.selector_offset ||
 				    a.table_offset != b.table_offset ||
-				    !EquivalentValue(m_program, a.key_count, b.key_count)) continue;
+				    !EquivalentValue(m_program, a.key_count, b.key_count) ||
+				    a.selector_mask.IsEmpty() != b.selector_mask.IsEmpty() ||
+				    (!a.selector_mask.IsEmpty() &&
+				     !EquivalentValue(m_program, a.selector_mask, b.selector_mask))) continue;
 			}
 			bool same = true;
 			for (uint32_t i = 0; i < descriptor.dword_count; i++) {
@@ -534,6 +537,277 @@ private:
 		return false;
 	}
 
+	bool ConditionalEdge(const Block* from, const Block* to, Value& condition,
+	                     bool& positive) const {
+		const auto position = std::ranges::find(m_program.blocks, from);
+		const auto target = std::ranges::find(m_program.blocks, to);
+		if (position == m_program.blocks.end() || target == m_program.blocks.end()) return false;
+		const auto& info = m_program.block_info[position - m_program.blocks.begin()];
+		const auto& term = info.terminator;
+		const auto id = m_program.block_info[target - m_program.blocks.begin()].id;
+		if (term.kind != CFG::TerminatorKind::ConditionalBranch ||
+		    (term.true_block == id) == (term.false_block == id)) return false;
+		condition = info.condition;
+		positive = term.true_block == id;
+		while (const auto* inst = condition.Resolve().TryInstruction()) {
+			if (inst->GetOpcode() != ValueOpcode::LogicalNot) break;
+			condition = inst->Arg(0);
+			positive = !positive;
+		}
+		condition = condition.Resolve();
+		return true;
+	}
+
+	Value PositiveUseGuard(const Block* use) const {
+		if (use == nullptr || use->ImmPredecessors().size() != 1u) return {};
+		Value condition;
+		bool positive = false;
+		return ConditionalEdge(use->ImmPredecessors()[0], use, condition, positive) && positive
+		           ? condition : Value {};
+	}
+
+	Value SimplifyGuard(Value guard) const {
+		const auto invariant = ResolveInvariantPhi(m_program, guard);
+		return (invariant.IsEmpty() ? guard : invariant).Resolve();
+	}
+
+	bool Implies(Value guard, Value required) const {
+		guard = SimplifyGuard(guard);
+		required = required.Resolve();
+		if (EquivalentValue(m_program, guard, required)) return true;
+		const auto* inst = guard.TryInstruction();
+		return inst != nullptr && inst->GetOpcode() == ValueOpcode::LogicalAnd &&
+		       (Implies(inst->Arg(0), required) || Implies(inst->Arg(1), required));
+	}
+
+	Value EqualLocalKey(Value guard, Value key) const {
+		guard = SimplifyGuard(guard);
+		const auto* inst = guard.TryInstruction();
+		if (inst == nullptr) return {};
+		if (inst->GetOpcode() == ValueOpcode::LogicalAnd) {
+			const auto left = EqualLocalKey(inst->Arg(0), key);
+			return left.IsEmpty() ? EqualLocalKey(inst->Arg(1), key) : left;
+		}
+		if (inst->GetOpcode() != ValueOpcode::IEqual32 || inst->NumArgs() != 2u) return {};
+		if (EquivalentValue(m_program, inst->Arg(0), key)) return inst->Arg(1).Resolve();
+		if (EquivalentValue(m_program, inst->Arg(1), key)) return inst->Arg(0).Resolve();
+		return {};
+	}
+
+	struct AffineOffset {
+		Value    index;
+		uint64_t stride = 0;
+		uint64_t offset = 0;
+	};
+
+	bool MatchAffineOffset(Value value, Value guard, AffineOffset& out,
+	                       uint32_t depth = 0) const {
+		if (depth > 16u) return false;
+		value = value.Resolve();
+		uint32_t immediate = 0;
+		if (ImmediateU32(value, immediate)) {
+			out.offset = immediate;
+			return true;
+		}
+		const auto* inst = value.TryInstruction();
+		if (inst == nullptr) return false;
+		if (inst->GetOpcode() == ValueOpcode::SelectU32 && inst->NumArgs() == 3u &&
+		    Implies(guard, inst->Arg(0))) {
+			return MatchAffineOffset(inst->Arg(1), guard, out, depth + 1u);
+		}
+		if (inst->GetOpcode() == ValueOpcode::IAdd32 && inst->NumArgs() == 2u) {
+			AffineOffset left, right;
+			if (!MatchAffineOffset(inst->Arg(0), guard, left, depth + 1u) ||
+			    !MatchAffineOffset(inst->Arg(1), guard, right, depth + 1u) ||
+			    (!left.index.IsEmpty() && !right.index.IsEmpty() &&
+			     !EquivalentValue(m_program, left.index, right.index))) return false;
+			out.index = left.index.IsEmpty() ? right.index : left.index;
+			out.stride = left.stride + right.stride;
+			out.offset = left.offset + right.offset;
+			return out.stride <= UINT32_MAX && out.offset <= UINT32_MAX;
+		}
+		if (inst->GetOpcode() == ValueOpcode::ShiftLeftLogical32 && inst->NumArgs() == 2u &&
+		    ImmediateU32(inst->Arg(1), immediate) && immediate < 32u) {
+			if (!MatchAffineOffset(inst->Arg(0), guard, out, depth + 1u)) return false;
+			out.stride <<= immediate;
+			out.offset <<= immediate;
+			return out.stride <= UINT32_MAX && out.offset <= UINT32_MAX;
+		}
+		out.index = value;
+		out.stride = 1u;
+		return value.GetType() == Type::U32;
+	}
+
+	bool ImpliesNonzero(Value guard, Value value) const {
+		guard = SimplifyGuard(guard);
+		const auto* inst = guard.TryInstruction();
+		if (inst == nullptr) return false;
+		if (inst->GetOpcode() == ValueOpcode::LogicalAnd) {
+			return ImpliesNonzero(inst->Arg(0), value) ||
+			       ImpliesNonzero(inst->Arg(1), value);
+		}
+		if (inst->GetOpcode() != ValueOpcode::INotEqual32 || inst->NumArgs() != 2u)
+			return false;
+		uint32_t zero = 1u;
+		return (ImmediateU32(inst->Arg(0), zero) && zero == 0u &&
+		        EquivalentValue(m_program, inst->Arg(1), value)) ||
+		       (ImmediateU32(inst->Arg(1), zero) && zero == 0u &&
+		        EquivalentValue(m_program, inst->Arg(0), value));
+	}
+
+	bool ImpliesIndexBelow32(Value guard, Value index) const {
+		guard = SimplifyGuard(guard);
+		const auto* inst = guard.TryInstruction();
+		if (inst == nullptr) return false;
+		if (inst->GetOpcode() == ValueOpcode::LogicalAnd) {
+			return ImpliesIndexBelow32(inst->Arg(0), index) ||
+			       ImpliesIndexBelow32(inst->Arg(1), index);
+		}
+		if (inst->NumArgs() != 2u) return false;
+		uint32_t limit = 0;
+		return (inst->GetOpcode() == ValueOpcode::SLessThan32 &&
+		        EquivalentValue(m_program, inst->Arg(0), index) &&
+		        ImmediateU32(inst->Arg(1), limit) && limit == 32u) ||
+		       (inst->GetOpcode() == ValueOpcode::SGreaterThan32 &&
+		        ImmediateU32(inst->Arg(0), limit) && limit == 32u &&
+		        EquivalentValue(m_program, inst->Arg(1), index));
+	}
+
+	bool ClearsFirstSetBit(Value value, Value mask) const {
+		const auto* update = value.Resolve().TryInstruction();
+		if (update == nullptr || update->GetOpcode() != ValueOpcode::BitwiseXor32 ||
+		    update->NumArgs() != 2u) return false;
+		Value bit;
+		if (EquivalentValue(m_program, update->Arg(0), mask)) bit = update->Arg(1);
+		else if (EquivalentValue(m_program, update->Arg(1), mask)) bit = update->Arg(0);
+		else return false;
+		const auto* shift = bit.Resolve().TryInstruction();
+		uint32_t one = 0;
+		if (shift == nullptr || shift->GetOpcode() != ValueOpcode::ShiftLeftLogical32 ||
+		    shift->NumArgs() != 2u || !ImmediateU32(shift->Arg(0), one) || one != 1u)
+			return false;
+		Value position = shift->Arg(1).Resolve();
+		const auto* masked = position.TryInstruction();
+		uint32_t lane_mask = 0;
+		if (masked != nullptr && masked->GetOpcode() == ValueOpcode::BitwiseAnd32 &&
+		    masked->NumArgs() == 2u) {
+			if (ImmediateU32(masked->Arg(0), lane_mask) && lane_mask == 31u)
+				position = masked->Arg(1);
+			else if (ImmediateU32(masked->Arg(1), lane_mask) && lane_mask == 31u)
+				position = masked->Arg(0);
+		}
+		const auto* first = position.Resolve().TryInstruction();
+		return first != nullptr && first->GetOpcode() == ValueOpcode::FindILsb32 &&
+		       first->NumArgs() == 1u && EquivalentValue(m_program, first->Arg(0), mask);
+	}
+
+	Value InitialClearingMask(Value value, const Block* update_block) const {
+		const auto* phi = value.Resolve().TryInstruction();
+		if (phi == nullptr || phi->GetOpcode() != ValueOpcode::Phi ||
+		    phi->NumArgs() != 2u || phi->GetType() != Type::U32) return {};
+		for (uint32_t back = 0; back < 2u; ++back) {
+			if (phi->PhiBlock(back) != update_block ||
+			    !ClearsFirstSetBit(phi->Arg(back), value)) continue;
+			const auto initial = phi->Arg(back ^ 1u).Resolve();
+			return ValidateRuntimeValue(m_program, initial, RuntimeValueType::Integer)
+			           ? initial : Value {};
+		}
+		return {};
+	}
+
+	bool EdgeOn(const Block* from, const Block* to, Value predicate, bool positive) const {
+		Value condition;
+		bool taken = false;
+		return ConditionalEdge(from, to, condition, taken) && taken == positive &&
+		       EquivalentValue(m_program, condition, predicate);
+	}
+
+	Value BoundedSetBitMask(Value index, Value guard) const {
+		const auto* phi = index.Resolve().TryInstruction();
+		if (phi == nullptr || phi->GetOpcode() != ValueOpcode::Phi ||
+		    phi->NumArgs() != 3u || phi->GetType() != Type::U32 ||
+		    !ImpliesIndexBelow32(guard, index)) return {};
+		for (uint32_t bit_arm = 0; bit_arm < 3u; ++bit_arm) {
+			const auto bit_guard = PositiveUseGuard(phi->PhiBlock(bit_arm));
+			const auto* selected = phi->Arg(bit_arm).Resolve().TryInstruction();
+			if (selected == nullptr || selected->GetOpcode() != ValueOpcode::SelectU32 ||
+			    selected->NumArgs() != 3u ||
+			    !Implies(bit_guard, selected->Arg(0))) continue;
+			const auto* first = selected->Arg(1).Resolve().TryInstruction();
+			if (first == nullptr || first->GetOpcode() != ValueOpcode::FindILsb32 ||
+			    first->NumArgs() != 1u || !ImpliesNonzero(bit_guard, first->Arg(0)))
+				continue;
+			const auto initial_mask = InitialClearingMask(first->Arg(0), phi->PhiBlock(bit_arm));
+			if (initial_mask.IsEmpty()) continue;
+			for (uint32_t sentinel_arm = 0; sentinel_arm < 3u; ++sentinel_arm) {
+				if (sentinel_arm == bit_arm) continue;
+				const auto sentinel_guard = PositiveUseGuard(phi->PhiBlock(sentinel_arm));
+				const auto* sentinel = phi->Arg(sentinel_arm).Resolve().TryInstruction();
+				uint32_t bound = 0;
+				if (sentinel == nullptr || sentinel->GetOpcode() != ValueOpcode::SelectU32 ||
+				    sentinel->NumArgs() != 3u ||
+				    !Implies(sentinel_guard, sentinel->Arg(0)) ||
+				    !ImmediateU32(sentinel->Arg(1), bound) || bound != 32u) continue;
+				const auto loop_active = sentinel->Arg(0).Resolve();
+				const auto* active_phi = loop_active.TryInstruction();
+				if (active_phi == nullptr || active_phi->GetOpcode() != ValueOpcode::Phi ||
+				    active_phi->NumArgs() != 2u) continue;
+				bool invariant = false;
+				for (uint32_t initial = 0; initial < 2u; ++initial) {
+					invariant = Implies(guard, active_phi->Arg(initial)) &&
+					            EdgeOn(active_phi->PhiBlock(initial ^ 1u), active_phi->Parent(),
+					                   active_phi->Arg(initial ^ 1u), true);
+					if (invariant) break;
+				}
+				if (!invariant) continue;
+				const auto other_arm = 3u - bit_arm - sentinel_arm;
+				if (EdgeOn(phi->PhiBlock(other_arm), phi->Parent(), loop_active, false))
+					return initial_mask;
+			}
+		}
+		return {};
+	}
+
+	bool MatchUniformizedMaterialKey(Value key, const Inst& image,
+	                                DescriptorSource::IndirectImage& indirect,
+	                                DescriptorSource& material_source, uint32_t pc) {
+		const auto guard = PositiveUseGuard(image.Parent());
+		if (guard.IsEmpty()) return false;
+		const auto local = EqualLocalKey(guard, key);
+		const auto* selected = local.Resolve().TryInstruction();
+		if (selected == nullptr || selected->GetOpcode() != ValueOpcode::SelectU32 ||
+		    selected->NumArgs() != 3u || !Implies(guard, selected->Arg(0))) return false;
+		const auto active = selected->Arg(0).Resolve();
+		const auto* read = selected->Arg(1).Resolve().TryInstruction();
+		if (read == nullptr || read->GetOpcode() != ValueOpcode::LoadAddressU32 ||
+		    read->NumArgs() != 4u || !EquivalentValue(m_program, read->Arg(3), active))
+			return false;
+		uint32_t high = 1u;
+		if (!ImmediateU32(read->Arg(2), high) || high != 0u) return false;
+		const auto memory_index = read->Flags<MemoryFlags>().index;
+		if (memory_index >= m_program.memory_info.size()) return false;
+		const auto& memory = m_program.memory_info[memory_index];
+		if (memory.kind != ResourceKind::Global || memory.data_bits != 32u ||
+		    memory.data_dwords != 1u || !MemoryIndexBelongsTo(memory_index, *read)) return false;
+		const auto* material_handle = read->Arg(0).Resolve().TryInstruction();
+		if (material_handle == nullptr ||
+		    material_handle->GetOpcode() != ValueOpcode::GetAddressResource ||
+		    !MakeRuntimeTableSource(*material_handle, pc, material_source)) return false;
+		AffineOffset offset;
+		if (!MatchAffineOffset(read->Arg(1), active, offset) || offset.index.IsEmpty() ||
+		    offset.stride == 0u || offset.offset + memory.offset > UINT32_MAX ||
+		    offset.offset + memory.offset + 31u * offset.stride + 4u >
+		        static_cast<uint64_t>(UINT32_MAX) + 1u) return false;
+		const auto mask = BoundedSetBitMask(offset.index, active);
+		if (mask.IsEmpty()) return false;
+		indirect.material_source = InternSource(material_source);
+		indirect.selector_stride = static_cast<uint32_t>(offset.stride);
+		indirect.selector_offset = static_cast<uint32_t>(offset.offset + memory.offset);
+		indirect.key_count = Value(32u);
+		indirect.selector_mask = mask;
+		return true;
+	}
+
 	Value BoundedLoopCount(Value key, const Block* use) const {
 		const auto* phi = key.Resolve().TryInstruction();
 		if (m_shader_writes || phi == nullptr || phi->GetOpcode() != ValueOpcode::Phi ||
@@ -656,7 +930,11 @@ private:
 			} else {
 				indirect.key_count = BoundedLoopCount(key, handle.Parent());
 			}
-			if (indirect.key_count.IsEmpty() || (table_offset & 3u) != 0u ||
+			if (indirect.key_count.IsEmpty() &&
+			    !MatchUniformizedMaterialKey(key, handle, indirect, material_source, pc)) {
+				return false;
+			}
+			if ((table_offset & 3u) != 0u ||
 			    (bitscan && table_offset > UINT32_MAX - (32u * 32u - 1u))) return false;
 		} else {
 			auto* material_read = key.Resolve().TryInstruction();
@@ -962,6 +1240,9 @@ private:
 				return;
 			}
 			ValidateAddressHandle(inst.Arg(0), flags.pc);
+			if (address_info.access == AddressAccess::Write) {
+				m_program.has_address_writes = true;
+			}
 			m_info.uses_dma = true;
 			return;
 		}