Resolve scalar image tables during resource planning

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

@@ -459,11 +459,13 @@ struct BlockInfo {
 
 struct DescriptorSource {
 	struct IndirectImage {
-		uint32_t material_source = 0;
-		uint32_t heap_source     = 0;
+		uint32_t material_source = UINT32_MAX;
+		uint32_t table_source    = 0;
 		uint32_t selector_stride = 0;
 		uint32_t selector_offset = 0;
 		uint32_t key_arg         = 0;
+		uint32_t table_offset    = 0;
+		uint32_t key_count       = 0;
 
 		bool operator==(const IndirectImage& other) const = default;
 	};

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

@@ -166,16 +166,15 @@ bool ReadSpecializationWord(const SrtRuntime& runtime, uint64_t address, uint32_
 	       runtime.read_specialization_memory(runtime.userdata, address, &word);
 }
 
-bool ReadScalarBufferWord(const ShaderBufferResource& descriptor, uint32_t dynamic_offset,
-                          uint32_t immediate_offset, const SrtRuntime& runtime, uint32_t& word) {
+bool ReadScalarTableWord(uint64_t base, uint64_t size, uint32_t dynamic_offset,
+                         uint32_t immediate_offset, const SrtRuntime& runtime, uint32_t& word) {
 	const auto byte_offset = static_cast<uint64_t>(dynamic_offset) + immediate_offset;
 	const auto aligned     = byte_offset & ~uint64_t {3};
-	const auto size        = descriptor.GetSize();
 	if (aligned > size || size - aligned < sizeof(uint32_t)) {
 		word = 0;
 		return true;
 	}
-	const auto base = descriptor.Base48() & ~uint64_t {3};
+	base &= AddressMask & ~uint64_t {3};
 	if (aligned > AddressMask - base) {
 		return false;
 	}
@@ -189,39 +188,59 @@ bool ReadScalarBufferWord(const ShaderBufferResource& descriptor, uint32_t dynam
 bool MaterializeIndirectImage(const ResourcePlan& program,
                               const DescriptorSource::IndirectImage& indirect,
                               const DescriptorValue& material_value,
-                              const DescriptorValue& heap_value, uint32_t image_index,
+                              const DescriptorValue& table_value, uint32_t image_index,
                               const SrtRuntime& runtime, ResourceSnapshot& snapshot,
                               ResourceSpecialization& specialization) {
-	ShaderBufferResource material;
-	ShaderBufferResource heap;
-	if (!DecodeBufferDescriptor(material_value, material) ||
-	    !DecodeBufferDescriptor(heap_value, heap) || material.Stride() != indirect.selector_stride) {
-		return false;
-	}
-	// Enumerate every wrapped scalar-buffer offset that can pass the descriptor bounds.
-	const auto step = std::gcd<uint64_t>(indirect.selector_stride, uint64_t {1} << 32u);
-	const auto residue = static_cast<uint64_t>(indirect.selector_offset) % step;
-	const auto limit = std::min<uint64_t>(UINT32_MAX, material.GetSize() + 3u);
-	const auto probe_count = residue <= limit ? (limit - residue) / step + 1u : 0u;
-	if (probe_count > MaxIndirectImageProbes) {
+	uint64_t table_base = 0;
+	uint64_t table_size = UINT64_MAX; // Scalar addresses have no buffer descriptor bounds.
+	ShaderBufferResource table;
+	if (table_value.dword_count == 2u) {
+		table_base = (static_cast<uint64_t>(table_value.dwords[1]) << 32u) | table_value.dwords[0];
+	} else if (DecodeBufferDescriptor(table_value, table)) {
+		table_base = table.Base48();
+		table_size = table.GetSize();
+	} else {
 		return false;
 	}
 	auto& keys = program.material_keys;
 	keys.clear();
-	keys.reserve(static_cast<size_t>(probe_count) + 1u);
-	keys.push_back(0u);
-	for (uint64_t offset = residue; offset <= limit && probe_count != 0u; offset += step) {
-		uint32_t key = 0;
-		if (!ReadScalarBufferWord(material, static_cast<uint32_t>(offset), 0u, runtime, key)) {
+	if (indirect.material_source == UINT32_MAX) {
+		if (table_value.dword_count != 2u || indirect.key_count == 0u ||
+		    indirect.key_count > MaxIndirectImageProbes) {
 			return false;
 		}
-		keys.push_back(key);
-		if (limit - offset < step) {
-			break;
+		keys.resize(indirect.key_count);
+		std::iota(keys.begin(), keys.end(), 0u);
+	} else {
+		ShaderBufferResource material;
+		if (!DecodeBufferDescriptor(material_value, material) || table_value.dword_count != 4u ||
+		    material.Stride() != indirect.selector_stride) {
+			return false;
 		}
+		// Enumerate every wrapped scalar-buffer offset that can pass the descriptor bounds.
+		const auto step = std::gcd<uint64_t>(indirect.selector_stride, uint64_t {1} << 32u);
+		const auto residue = static_cast<uint64_t>(indirect.selector_offset) % step;
+		const auto limit = std::min<uint64_t>(UINT32_MAX, material.GetSize() + 3u);
+		const auto probe_count = residue <= limit ? (limit - residue) / step + 1u : 0u;
+		if (probe_count > MaxIndirectImageProbes) {
+			return false;
+		}
+		keys.reserve(static_cast<size_t>(probe_count) + 1u);
+		keys.push_back(0u);
+		for (uint64_t offset = residue; offset <= limit && probe_count != 0u; offset += step) {
+			uint32_t key = 0;
+			if (!ReadScalarTableWord(material.Base48(), material.GetSize(),
+			                         static_cast<uint32_t>(offset), 0u, runtime, key)) {
+				return false;
+			}
+			keys.push_back(key);
+			if (limit - offset < step) {
+				break;
+			}
+		}
+		std::ranges::sort(keys);
+		keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
 	}
-	std::ranges::sort(keys);
-	keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
 
 	const auto children_begin = snapshot.images.size();
 	const auto mapping_offset = snapshot.flattened_srt.size();
@@ -232,10 +251,10 @@ bool MaterializeIndirectImage(const ResourcePlan& program,
 		const auto key = keys[entry];
 		DescriptorValue candidate;
 		candidate.dword_count = 8u;
-		const auto heap_offset = key << 5u;
+		const auto table_offset = (key << 5u) + indirect.table_offset;
 		for (uint32_t dword = 0; dword < candidate.dword_count; ++dword) {
-			if (!ReadScalarBufferWord(heap, heap_offset, dword * sizeof(uint32_t), runtime,
-			                          candidate.dwords[dword])) {
+			if (!ReadScalarTableWord(table_base, table_size, table_offset, dword * sizeof(uint32_t),
+			                         runtime, candidate.dwords[dword])) {
 				return false;
 			}
 		}
@@ -816,7 +835,7 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
 		plan.requires_specialization_memory = true;
 		MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->material_source),
 		                   plan.clean_flat_slots);
-		MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->heap_source),
+		MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->table_source),
 		                   plan.clean_flat_slots);
 	}
 	return plan;
@@ -892,10 +911,11 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 			}
 			const auto& indirect = *source->indirect_image;
 			DescriptorValue material;
-			DescriptorValue heap;
-			if (!clean.EvaluateDescriptor(indirect.material_source, material) ||
-			    !clean.EvaluateDescriptor(indirect.heap_source, heap) ||
-			    !MaterializeIndirectImage(program, indirect, material, heap, i, runtime, snapshot,
+			DescriptorValue table;
+			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, snapshot,
 			                              specialization)) {
 				return false;
 			}

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

@@ -344,16 +344,26 @@ private:
 	}
 
 	const MemoryInfo* ScalarReadMemory(const Inst& read, uint32_t& index) const {
-		if (read.GetOpcode() != ValueOpcode::ReadConstBuffer || read.NumArgs() != 2u) {
+		const bool address = read.GetOpcode() == ValueOpcode::LoadAddressU32;
+		if (!(address ? read.NumArgs() == 4u
+		              : read.GetOpcode() == ValueOpcode::ReadConstBuffer && read.NumArgs() == 2u)) {
 			return nullptr;
 		}
+		if (address) {
+			const auto high = read.Arg(2).Resolve();
+			const auto enabled = read.Arg(3).Resolve();
+			if (!high.IsImmediate() || high.GetType() != Type::U32 || high.U32() != 0u ||
+			    !enabled.IsImmediate() || enabled.GetType() != Type::U1 || !enabled.U1()) {
+				return nullptr;
+			}
+		}
 		index = read.Flags<MemoryFlags>().index;
 		if (index >= m_program.memory_info.size()) {
 			return nullptr;
 		}
 		const auto& memory = m_program.memory_info[index];
-		return memory.kind == ResourceKind::ScalarBuffer && memory.data_bits == 32u &&
-		               memory.data_dwords == 1u
+		return memory.kind == (address ? ResourceKind::ScalarAddress : ResourceKind::ScalarBuffer) &&
+		               memory.data_bits == 32u && memory.data_dwords == 1u
 		           ? &memory
 		           : nullptr;
 	}
@@ -376,18 +386,16 @@ private:
 		return true;
 	}
 
-	bool MakeRuntimeBufferSource(const Inst& handle, uint32_t pc, uint32_t& source,
-	                             DescriptorSource& descriptor) {
-		if (handle.GetOpcode() != ValueOpcode::GetBufferResource) {
+	bool MakeRuntimeTableSource(const Inst& handle, uint32_t pc,
+	                            DescriptorSource& descriptor) {
+		const auto width = handle.GetOpcode() == ValueOpcode::GetBufferResource ? 4u
+		                 : handle.GetOpcode() == ValueOpcode::GetAddressResource ? 2u : 0u;
+		if (width == 0u) {
 			return false;
 		}
-		MakeSource(handle, 4u, false, false, descriptor, pc);
+		MakeSource(handle, width, false, false, descriptor, pc);
 		uint32_t bad_dword = 0;
-		if (!ValidateSource(descriptor, bad_dword)) {
-			return false;
-		}
-		source = InternSource(descriptor);
-		return true;
+		return ValidateSource(descriptor, bad_dword);
 	}
 
 	bool MatchMaterialOffset(Value value, Value& selector, uint32_t& stride,
@@ -424,106 +432,183 @@ private:
 		       selector_inst->GetOpcode() == ValueOpcode::ReadFirstLane;
 	}
 
+	bool NonzeroOnEntry(Value value, const Block* block) const {
+		if (m_program.blocks.size() != m_program.block_info.size()) {
+			return false;
+		}
+		// Each unique predecessor must execute before this use. Stop at joins: an
+		// unrelated comparison is not a bound on FindILsb's zero-input sentinel.
+		for (size_t depth = 0; block != nullptr && depth < m_program.blocks.size(); ++depth) {
+			if (block->ImmPredecessors().size() != 1u) {
+				return false;
+			}
+			const auto* previous = block->ImmPredecessors()[0];
+			const auto current_it = std::ranges::find(m_program.blocks, block);
+			const auto previous_it = std::ranges::find(m_program.blocks, previous);
+			if (current_it == m_program.blocks.end() || previous_it == m_program.blocks.end()) {
+				return false;
+			}
+			const auto& info = m_program.block_info[previous_it - m_program.blocks.begin()];
+			const auto id = m_program.block_info[current_it - m_program.blocks.begin()].id;
+			const auto& term = info.terminator;
+			if (term.kind == CFG::TerminatorKind::ConditionalBranch &&
+			    (term.true_block == id) != (term.false_block == id)) {
+				bool positive = term.true_block == id;
+				const auto* test = info.condition.Resolve().TryInstruction();
+				while (test != nullptr && test->GetOpcode() == ValueOpcode::LogicalNot) {
+					positive = !positive;
+					test = test->Arg(0).Resolve().TryInstruction();
+				}
+				if (test != nullptr && test->NumArgs() == 2u &&
+				    ((test->GetOpcode() == ValueOpcode::INotEqual32 && positive) ||
+				     (test->GetOpcode() == ValueOpcode::IEqual32 && !positive))) {
+					for (uint32_t arg = 0; arg < 2u; ++arg) {
+						uint32_t immediate;
+						if (ImmediateU32(test->Arg(arg), immediate) && immediate == 0u &&
+						    EquivalentValue(m_program, test->Arg(arg ^ 1u), value)) {
+							return true;
+						}
+					}
+				}
+			}
+			block = previous;
+		}
+		return false;
+	}
+
+	bool MatchTableOffset(Value value, Value& key, uint32_t& offset) const {
+		offset = 0;
+		for (;;) {
+			const auto* inst = value.Resolve().TryInstruction();
+			if (inst == nullptr || inst->NumArgs() != 2u) {
+				return false;
+			}
+			uint32_t immediate;
+			if (inst->GetOpcode() == ValueOpcode::ShiftLeftLogical32 &&
+			    ImmediateU32(inst->Arg(1), immediate) && immediate == 5u) {
+				key = inst->Arg(0).Resolve();
+				return key.GetType() == Type::U32;
+			}
+			if (inst->GetOpcode() != ValueOpcode::IAdd32) {
+				return false;
+			}
+			if (ImmediateU32(inst->Arg(0), immediate)) {
+				value = inst->Arg(1);
+			} else if (ImmediateU32(inst->Arg(1), immediate)) {
+				value = inst->Arg(0);
+			} else {
+				return false;
+			}
+			// These additions are shader U32 arithmetic, before the scalar memory offset.
+			offset += immediate;
+		}
+	}
+
 	bool TryMakeIndirectImage(Inst& handle, uint32_t pc, IndirectImagePlan& plan) {
 		if (handle.GetOpcode() != ValueOpcode::GetImageResource || handle.NumArgs() != 8u) {
 			return false;
 		}
-
-		std::array<Inst*, 8> heap_reads {};
-		Inst*                heap_handle = nullptr;
-		Value                heap_offset;
-		for (uint32_t dword = 0; dword < heap_reads.size(); dword++) {
-			heap_reads[dword] = handle.Arg(dword).Resolve().TryInstruction();
-			if (heap_reads[dword] == nullptr) {
+		Inst* table_handle = nullptr;
+		Value key;
+		uint32_t table_offset = 0;
+		for (uint32_t dword = 0; dword < plan.reads.size(); ++dword) {
+			auto* read = handle.Arg(dword).Resolve().TryInstruction();
+			if (read == nullptr) {
 				return false;
 			}
-			uint32_t    memory_index = 0;
-			const auto* memory       = ScalarReadMemory(*heap_reads[dword], memory_index);
-			if (memory == nullptr || memory->offset != dword * sizeof(uint32_t) ||
-			    !MemoryIndexBelongsTo(memory_index, *heap_reads[dword])) {
+			uint32_t memory_index = 0;
+			const auto* memory = ScalarReadMemory(*read, memory_index);
+			if (memory == nullptr || memory->offset > INT32_MAX || (memory->offset & 3u) != 0u ||
+			    !MemoryIndexBelongsTo(memory_index, *read)) {
 				return false;
 			}
-			auto* current_handle = heap_reads[dword]->Arg(0).Resolve().TryInstruction();
+			auto* current_handle = read->Arg(0).Resolve().TryInstruction();
+			Value current_key;
+			uint32_t offset = 0;
 			if (current_handle == nullptr ||
-			    (heap_handle != nullptr && current_handle != heap_handle)) {
+			    current_handle->GetOpcode() != (memory->kind == ResourceKind::ScalarAddress
+			                                      ? ValueOpcode::GetAddressResource
+			                                      : ValueOpcode::GetBufferResource) ||
+			    (memory->kind == ResourceKind::ScalarAddress && read->Parent() != handle.Parent()) ||
+			    (table_handle != nullptr &&
+			     !EquivalentValue(m_program, Value(table_handle), Value(current_handle))) ||
+			    !MatchTableOffset(read->Arg(1), current_key, offset) ||
+			    memory->offset > UINT32_MAX - offset) {
 				return false;
 			}
-			heap_handle = current_handle;
+			offset += memory->offset;
 			if (dword == 0u) {
-				heap_offset = heap_reads[dword]->Arg(1).Resolve();
-			} else if (!EquivalentValue(m_program, heap_offset, heap_reads[dword]->Arg(1))) {
+				key = current_key;
+				table_offset = offset;
+			} else if (!EquivalentValue(m_program, key, current_key) ||
+			           static_cast<uint64_t>(table_offset) + dword * sizeof(uint32_t) != offset) {
 				return false;
 			}
-			plan.memory[dword] = memory_index;
-			plan.reads[dword]  = heap_reads[dword];
-		}
-
-		const auto* shift        = heap_offset.TryInstruction();
-		uint32_t    shift_amount = 0;
-		if (shift == nullptr || shift->GetOpcode() != ValueOpcode::ShiftLeftLogical32 ||
-		    shift->NumArgs() != 2u || !ImmediateU32(shift->Arg(1), shift_amount) ||
-		    shift_amount != 5u) {
-			return false;
-		}
-		auto* material_read = shift->Arg(0).Resolve().TryInstruction();
-		if (material_read == nullptr) {
-			return false;
-		}
-		uint32_t    material_memory_index = 0;
-		const auto* material_memory       = ScalarReadMemory(*material_read, material_memory_index);
-		if (material_memory == nullptr || material_memory->offset != 0u ||
-		    !MemoryIndexBelongsTo(material_memory_index, *material_read)) {
-			return false;
-		}
-		auto* material_handle = material_read->Arg(0).Resolve().TryInstruction();
-		if (material_handle == nullptr) {
-			return false;
-		}
-
-		Value    selector;
-		uint32_t selector_stride = 0;
-		uint32_t selector_offset = 0;
-		if (!MatchMaterialOffset(material_read->Arg(1), selector, selector_stride,
-		                         selector_offset)) {
-			return false;
-		}
-
-		const std::array<const Inst*, 1> material_users {shift};
-		std::array<const Inst*, 8>       heap_users {};
-		std::copy(heap_reads.begin(), heap_reads.end(), heap_users.begin());
-		const std::array<const Inst*, 1> image_users {&handle};
-		if (!UsesOnly(*material_read, material_users) || !UsesOnly(*shift, heap_users)) {
-			return false;
-		}
-		for (const auto* read: heap_reads) {
+			table_handle = current_handle;
+			const std::array<const Inst*, 1> image_users {&handle};
 			if (!UsesOnly(*read, image_users)) {
 				return false;
 			}
+			plan.memory[dword] = memory_index;
+			plan.reads[dword] = read;
 		}
 
-		DescriptorSource material_source;
-		DescriptorSource heap_source;
-		uint32_t         material_source_index = 0;
-		uint32_t         heap_source_index     = 0;
-		if (!MakeRuntimeBufferSource(*material_handle, pc, material_source_index,
-		                             material_source) ||
-		    !MakeRuntimeBufferSource(*heap_handle, pc, heap_source_index, heap_source)) {
+		DescriptorSource table_source;
+		if (!MakeRuntimeTableSource(*table_handle, pc, table_source)) {
 			return false;
 		}
-
+		DescriptorSource material_source;
+		DescriptorSource::IndirectImage indirect;
+		indirect.table_offset = table_offset;
+		if (table_source.dword_count == 2u) {
+			const auto* selector = key.Resolve().TryInstruction();
+			if (m_shader_writes || selector == nullptr ||
+			    selector->GetOpcode() != ValueOpcode::FindILsb32 || selector->NumArgs() != 1u ||
+			    !NonzeroOnEntry(selector->Arg(0), handle.Parent()) ||
+			    (table_offset & 3u) != 0u || table_offset > UINT32_MAX - (32u * 32u - 1u)) {
+				return false;
+			}
+			indirect.key_count = 32u;
+		} else {
+			auto* material_read = key.Resolve().TryInstruction();
+			uint32_t material_memory_index = 0;
+			const auto* memory = material_read != nullptr
+			                         ? ScalarReadMemory(*material_read, material_memory_index) : nullptr;
+			if (table_offset != 0u || memory == nullptr || memory->kind != ResourceKind::ScalarBuffer ||
+			    memory->offset != 0u || !MemoryIndexBelongsTo(material_memory_index, *material_read)) {
+				return false;
+			}
+			Value selector;
+			if (!MatchMaterialOffset(material_read->Arg(1), selector, indirect.selector_stride,
+			                         indirect.selector_offset)) {
+				return false;
+			}
+			const auto* shift = plan.reads[0]->Arg(1).Resolve().TryInstruction();
+			const std::array<const Inst*, 1> material_users {shift};
+			if (!UsesOnly(*material_read, material_users) || !UsesOnly(*shift, plan.reads)) {
+				return false;
+			}
+			const auto* material_handle = material_read->Arg(0).Resolve().TryInstruction();
+			if (material_handle == nullptr ||
+			    material_handle->GetOpcode() != ValueOpcode::GetBufferResource ||
+			    !MakeRuntimeTableSource(*material_handle, pc, material_source)) {
+				return false;
+			}
+			indirect.material_source = InternSource(material_source);
+		}
+		indirect.table_source = InternSource(table_source);
 		DescriptorSource image_source;
 		image_source.dword_count = 8u;
-		std::copy(material_source.dwords.begin(), material_source.dwords.begin() + 4u,
-		          image_source.dwords.begin());
-		std::copy(heap_source.dwords.begin(), heap_source.dwords.begin() + 4u,
-		          image_source.dwords.begin() + 4u);
-		image_source.indirect_image = DescriptorSource::IndirectImage {
-		    material_source_index, heap_source_index, selector_stride, selector_offset, 0u};
-
+		image_source.dwords.fill(Value(0u));
+		std::copy_n(material_source.dwords.begin(), material_source.dword_count,
+		            image_source.dwords.begin());
+		std::copy_n(table_source.dwords.begin(), table_source.dword_count,
+		            image_source.dwords.begin() + 4u);
+		image_source.indirect_image = indirect;
 		plan.handle = &handle;
 		plan.source = InternSource(image_source);
-		plan.key    = Value(material_read);
-		plan.roots  = image_source.dwords;
+		plan.key = key;
+		plan.roots = image_source.dwords;
 		return true;
 	}