shader: share nested pointer-table image and sampler selection

Co-authored-by: fxpw <fxpw@outlook.com>

src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp

@@ -826,8 +826,13 @@ std::array<ImageSampler, 2> LoadImageSamplers(ValueEmitContext& ctx, const IR::I
 			const auto* handle = inst.Arg(1).ResolveInstruction();
 			if (handle == nullptr || handle->GetOpcode() != IR::ValueOpcode::GetSamplerResource)
 				ctx.Fail(inst, "has invalid finite sampler key provenance");
+			auto selected = ctx.Def(handle->Arg(0));
+			if (samplers[index].indirect_root != IR::SamplerResource::NoIndirectSampler)
+				selected = EmitIndirectResourceIndex(state, selected,
+				    samplers[index].indirect_mapping_offset,
+				    samplers[index].indirect_search_iterations, 0u);
 			array_index = EmitIndexSwitch(
-			    state, ctx.Def(handle->Arg(0)), static_cast<uint32_t>(candidates.size()),
+			    state, selected, static_cast<uint32_t>(candidates.size()),
 			    TypeU32(state), [&](uint32_t ordinal) { return ConstantU32(state, candidates[ordinal]); });
 		}
 		sampler = {LoadSamplerDescriptor(state, index, array_index), array_index != 0u};

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

@@ -181,10 +181,15 @@ RawSignedStorageFormat(const ImageResource& image, Prospero::BufferFormat format
 }
 
 struct SamplerResource {
+	static constexpr uint32_t NoIndirectSampler = UINT32_MAX;
+
 	uint32_t source                = 0;
 	uint32_t first_use_pc          = 0;
 	// Native filtering/border variants share the original sampler's runtime descriptor.
 	uint32_t snapshot_index        = 0;
+	uint32_t indirect_root = NoIndirectSampler;
+	uint32_t indirect_mapping_offset = 0;
+	uint32_t indirect_search_iterations = 0;
 	std::vector<uint32_t> indirect_resources;
 	bool     force_point_filtering = false;
 	bool     depth_compare         = false;
@@ -520,6 +525,12 @@ struct ShaderInfo {
 
 struct DescriptorSource {
 	struct IndirectDescriptor {
+		struct PointerTable {
+			uint32_t maximum = 0;
+			uint32_t offset = 0;
+
+			bool operator==(const PointerTable&) const = default;
+		};
 		struct SelectorRead {
 			struct CountedRows {
 				uint32_t source = UINT32_MAX;
@@ -542,6 +553,7 @@ struct DescriptorSource {
 			bool operator==(const SelectorRead& other) const = default;
 		};
 		std::optional<SelectorRead> selector;
+		std::optional<PointerTable> pointer_table;
 		uint32_t table_source    = 0;
 		uint32_t table_offset    = 0;
 		uint32_t table_immediate = 0;

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

@@ -127,7 +127,7 @@ bool ValidatePackedUnormSampling(const ResourcePlan& program, const ResourceSnap
 	for (const auto& pair: program.info.sampled_pairs) {
 		if (pair.image != image_index) continue;
 		if (pair.sampler >= snapshot.samplers.size() ||
-		    !program.info.samplers[pair.sampler].indirect_resources.empty()) {
+		    program.descriptor_sources[program.info.samplers[pair.sampler].source].indirect_descriptor) {
 			return SpecializationFail("packed UNorm sampling requires a direct sampler");
 		}
 		ShaderSamplerResource sampler;
@@ -367,11 +367,16 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 			      (table.Stride() & 3u) != 0u || table.Stride() < indirect.table_record_bytes)))
 				return false;
 			const bool affine = indirect.table_stride != 0u;
-			const uint64_t step = affine ? uint64_t {1} << std::countr_zero(indirect.table_stride) : 4u;
-			const uint64_t first = ((indirect.table_offset % step) & ~uint64_t {3}) +
+			const auto* pointer = indirect.pointer_table ? &*indirect.pointer_table : nullptr;
+			const uint64_t step = pointer != nullptr ? indirect.table_stride
+			    : affine ? uint64_t {1} << std::countr_zero(indirect.table_stride) : 4u;
+			const uint64_t first = ((pointer != nullptr ? indirect.table_offset
+			    : indirect.table_offset % step) & ~uint64_t {3}) +
 			                       (indirect.table_immediate & ~3u);
+			const uint64_t domain = pointer != nullptr ? uint64_t {pointer->maximum} + 1u
+			                                           : (uint64_t {1} << 32u) / step;
 			const auto count = table_size >= first + 4u
-			    ? std::min((table_size - first - 4u) / step + 1u, (uint64_t {1} << 32u) / step)
+			    ? std::min((table_size - first - 4u) / step + 1u, domain)
 			    : 0u;
 			if (affine && count > MaxIndirectDescriptorProbes) return false;
 			// Affine SMEM offsets wrap in U32 before adding the aligned immediate.
@@ -385,7 +390,13 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 			snapshot.flattened_srt.resize(mapping_offset + 2u, 0u);
 			for (uint32_t word = 0; word < count; ++word) {
 				DescriptorValue candidate {.dword_count = dword_count};
-				if (affine) {
+				if (pointer != nullptr) {
+					std::array<uint32_t, 2> address {};
+					if (!ReadScalarTable(table_base, table_size, first + word * step, runtime, address) ||
+					    !ReadScalarTable((uint64_t {address[1]} << 32u) | address[0], UINT64_MAX,
+					                     pointer->offset, runtime,
+					                     std::span(candidate.dwords).first(dword_count))) return false;
+				} else if (affine) {
 					if (!ReadScalarTable(table_base, table_size, first + word * step, runtime,
 					                     std::span(candidate.dwords).first(dword_count))) return false;
 				} else {
@@ -1437,18 +1448,28 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 		}
 	}
 	snapshot.samplers.resize(program.info.samplers.size());
-	for (uint32_t i = 0; i < program.info.samplers.size(); ++i) {
-		auto source = program.info.samplers[i].source;
-		if (source >= program.descriptor_sources.size()) return false;
-		const auto& descriptor = program.descriptor_sources[source];
-		if (descriptor.indirect_descriptor) {
-			if (descriptor.indirect_descriptor->sources.empty()) return false;
-			source = descriptor.indirect_descriptor->sources[0];
+	specialization.samplers.assign(program.info.samplers.size(), {});
+	for (uint32_t i = 0; i < snapshot.samplers.size(); ++i) {
+		const auto base_index =
+		    i < program.info.samplers.size() ? i : specialization.samplers[i].indirect_root;
+		if (base_index >= program.info.samplers.size()) return false;
+		if (i < program.info.samplers.size()) {
+			auto source = program.info.samplers[i].source;
+			const auto* descriptor = Source(program, source);
+			if (descriptor == nullptr) return false;
+			const auto* indirect = descriptor->indirect_descriptor ? &*descriptor->indirect_descriptor : nullptr;
+			if (indirect != nullptr && indirect->sources.empty()) {
+				snapshot.samplers[i] = {.dword_count = 4u};
+				if ((active.empty() || active[source]) && !MaterializeIndirectDescriptor(
+				        program, *indirect, i, 4u, observed, clean, snapshot, snapshot.samplers,
+				        specialization.samplers, ShaderInfo::MaxSamplers,
+				        [](DescriptorValue&) { return true; })) return false;
+			} else {
+				if (indirect != nullptr) source = indirect->sources[0];
+				if (!evaluate(source, snapshot.samplers[i])) return false;
+			}
 		}
-		if (!evaluate(source, snapshot.samplers[i])) {
-			return false;
-		}
-		if (program.info.samplers[i].gather_lod) {
+		if (program.info.samplers[base_index].gather_lod) {
 			const auto control = snapshot.samplers[i].dwords[2];
 			const auto filter = (control >> 26u) & 3u;
 			// MipNone always selects the base level. Explicit point gathers currently require
@@ -1466,7 +1487,8 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 void ApplyResourceSpecialization(Program& program, const ResourceSpecialization& specialization) {
 	EXIT_IF(!program.resource_tracking_complete || program.shader_info_complete);
 	EXIT_IF(program.info.buffers.size() > specialization.buffers.size() ||
-	        program.info.images.size() > specialization.images.size());
+	        program.info.images.size() > specialization.images.size() ||
+	        program.info.samplers.size() > specialization.samplers.size());
 
 	auto& buffers = program.info.buffers;
 	const auto original_buffer_count = buffers.size();
@@ -1522,10 +1544,30 @@ void ApplyResourceSpecialization(Program& program, const ResourceSpecialization&
 		}
 	}
 
-	SamplerPlan sampler_plan;
-	EXIT_IF(!BuildSamplerPlan(program.info, sampler_plan));
 	auto& samplers      = program.info.samplers;
 	auto& sampled_pairs = program.info.sampled_pairs;
+	const auto base_sampler_count = samplers.size();
+	samplers.reserve(specialization.samplers.size());
+	for (uint32_t index = 0; index < specialization.samplers.size(); ++index) {
+		const auto& source = specialization.samplers[index];
+		if (index >= samplers.size()) {
+			EXIT_IF(source.indirect_root >= base_sampler_count);
+			samplers.push_back(samplers[source.indirect_root]);
+			samplers.back().indirect_resources.clear();
+		}
+		samplers[index].indirect_root = source.indirect_root;
+		samplers[index].indirect_mapping_offset = source.indirect_mapping_offset;
+		samplers[index].indirect_search_iterations = source.indirect_search_iterations;
+	}
+	for (uint32_t index = 0; index < samplers.size(); ++index) {
+		const auto root = samplers[index].indirect_root;
+		if (root != SamplerResource::NoIndirectSampler) {
+			EXIT_IF(root >= samplers.size());
+			samplers[root].indirect_resources.push_back(index);
+		}
+	}
+	SamplerPlan sampler_plan;
+	EXIT_IF(!BuildSamplerPlan(program.info, sampler_plan));
 	const auto original_sampler_count = samplers.size();
 	samplers.reserve(sampler_plan.sampler_count);
 	for (uint32_t index = 0; index < sampler_plan.sampler_count; index++) {

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

@@ -34,8 +34,16 @@ struct ResourceSpecialization {
 		bool                          operator==(const Image&) const = default;
 	};
 
+	struct Sampler {
+		uint32_t indirect_root = SamplerResource::NoIndirectSampler;
+		uint32_t indirect_mapping_offset = 0;
+		uint32_t indirect_search_iterations = 0;
+		bool operator==(const Sampler&) const = default;
+	};
+
 	std::vector<Buffer> buffers;
 	std::vector<Image>  images;
+	std::vector<Sampler> samplers;
 
 	bool operator==(const ResourceSpecialization&) const = default;
 };

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

@@ -45,10 +45,13 @@ uint32_t PossibleU32Bits(Value value, const auto& normalize) {
 				if (inst->NumArgs() != 0u) bits = merged;
 				break;
 			}
-			case ValueOpcode::ShiftLeftLogical32: {
+			case ValueOpcode::ShiftLeftLogical32:
+			case ValueOpcode::ShiftRightLogical32: {
 				const auto shift = inst->Arg(1).Resolve();
-				if (shift.IsImmediate() && shift.GetType() == Type::U32)
-					bits = arg(0) << (shift.U32() & 31u);
+				if (shift.IsImmediate() && shift.GetType() == Type::U32) {
+					bits = inst->GetOpcode() == ValueOpcode::ShiftLeftLogical32
+					    ? arg(0) << (shift.U32() & 31u) : arg(0) >> (shift.U32() & 31u);
+				}
 				break;
 			}
 			default: break;
@@ -397,6 +400,13 @@ public:
 		std::vector<std::pair<const IndirectDescriptorPlan*, Value>> buffer_selections;
 		for (const auto& plan: m_indirect_descriptors) {
 			auto key = plan.key;
+			if (m_sources[plan.source].indirect_descriptor->pointer_table) {
+				auto* block = plan.handle->Parent();
+				const auto where = std::ranges::find_if(block->Instructions(),
+				    [&](const Inst& inst) { return &inst == plan.handle; });
+				key = Value(&*block->PrependNewInst(
+				    where, ValueOpcode::ShiftLeftLogical32, {key, Value(2u)}));
+			}
 			if (plan.table_read != nullptr) {
 				const auto* anchor = plan.table_indexed ? plan.handle : plan.table_read;
 				auto* block = anchor->Parent();
@@ -964,7 +974,8 @@ private:
 			if (current.indirect_descriptor.has_value()) {
 				const auto& a = *current.indirect_descriptor;
 				const auto& b = *descriptor.indirect_descriptor;
-				if (a.selector != b.selector || a.table_source != b.table_source ||
+				if (a.selector != b.selector || a.pointer_table != b.pointer_table ||
+				    a.table_source != b.table_source ||
 				    a.table_offset != b.table_offset || a.table_immediate != b.table_immediate ||
 				    a.table_stride != b.table_stride || a.indexed != b.indexed ||
 				    a.workgroup_axis != b.workgroup_axis || a.sources != b.sources ||
@@ -2713,6 +2724,74 @@ private:
 		return true;
 	}
 
+	bool TryMakeNestedDescriptor(Inst& handle, const DescriptorSource& descriptor,
+	                             IndirectDescriptorPlan& plan) {
+		if (handle.GetOpcode() != ValueOpcode::GetImageResource &&
+		    handle.GetOpcode() != ValueOpcode::GetSamplerResource) return false;
+		const Inst* pointer = nullptr;
+		uint32_t descriptor_offset = 0;
+		for (uint32_t word = 0; word < handle.NumArgs(); ++word) {
+			const auto* read = descriptor.dwords[word].Resolve().TryInstruction();
+			uint32_t index, offset;
+			const auto* memory = read != nullptr ? ScalarReadMemory(*read, index) : nullptr;
+			if (memory == nullptr || memory->kind != ResourceKind::ScalarAddress ||
+			    memory->offset > INT32_MAX || !ImmediateU32(read->Arg(1), offset) ||
+			    !AvailableBefore(*read, handle)) return false;
+			const auto* address = read->Arg(0).Resolve().TryInstruction();
+			const uint64_t current_offset = uint64_t {offset & ~3u} + (memory->offset & ~3u);
+			if (address == nullptr || address->GetOpcode() != ValueOpcode::GetAddressResource ||
+			    address->NumArgs() != 2u || current_offset > UINT32_MAX ||
+			    (pointer != nullptr && (!EquivalentValue(m_program, Value(pointer), Value(address)) ||
+			                            current_offset != uint64_t {descriptor_offset} + word * 4u)))
+				return false;
+			if (word == 0u) descriptor_offset = static_cast<uint32_t>(current_offset);
+			pointer = address;
+		}
+		// An out-of-bounds pointer pair is zero. Guest mappings exclude the first
+		// 16 KiB page, so these pointed reads have the existing null-image semantics.
+		if (pointer == nullptr || uint64_t {descriptor_offset} + handle.NumArgs() * 4u > 0x4000u)
+			return false;
+		const auto* low = pointer->Arg(0).Resolve().TryInstruction();
+		const auto* high = pointer->Arg(1).Resolve().TryInstruction();
+		uint32_t low_index, high_index;
+		const auto* low_memory = low != nullptr ? ScalarReadMemory(*low, low_index) : nullptr;
+		const auto* high_memory = high != nullptr ? ScalarReadMemory(*high, high_index) : nullptr;
+		if (low_memory == nullptr || high_memory == nullptr ||
+		    low_memory->kind != ResourceKind::ScalarBuffer ||
+		    high_memory->kind != ResourceKind::ScalarBuffer ||
+		    low_memory->offset > INT32_MAX || high_memory->offset > INT32_MAX ||
+		    (low_memory->offset & ~3u) + 4u != (high_memory->offset & ~3u) ||
+		    !EquivalentValue(m_program, low->Arg(0), high->Arg(0)) ||
+		    !EquivalentValue(m_program, low->Arg(1), high->Arg(1)) ||
+		    !AvailableBefore(*low, *pointer) || !AvailableBefore(*high, *pointer)) return false;
+		Value key;
+		DescriptorSource::IndirectDescriptor indirect;
+		if (!MatchTableOffset(low->Arg(1), key, indirect.table_offset, indirect.table_stride) ||
+		    indirect.table_stride < 8u || (indirect.table_stride & 3u) != 0u) return false;
+		const auto maximum = SelectorMaximum(key);
+		if (maximum > UINT32_MAX / 4u ||
+		    uint64_t {maximum} * indirect.table_stride + indirect.table_offset > UINT32_MAX ||
+		    (key.TryInstruction() != nullptr && !AvailableBefore(*key.Instruction(), handle)))
+			return false;
+		DescriptorSource table_source;
+		if (!MakeRuntimeTableSource(*low, table_source) || table_source.dword_count != 4u)
+			return false;
+		indirect.table_source = InternSource(table_source);
+		indirect.table_immediate = low_memory->offset & ~3u;
+		indirect.table_scalar = true;
+		indirect.pointer_table = {{maximum, descriptor_offset}};
+		DescriptorSource source;
+		source.dword_count = handle.NumArgs();
+		source.dwords.fill(Value(0u));
+		std::copy_n(table_source.dwords.begin(), 4u, source.dwords.begin());
+		source.indirect_descriptor = indirect;
+		plan.handle = &handle;
+		plan.source = InternSource(source);
+		plan.key = key;
+		plan.roots = source.dwords;
+		return true;
+	}
+
 	bool TryMakeIndirectImage(Inst& handle, const DescriptorSource& descriptor,
 	                          IndirectDescriptorPlan& plan) {
 		if (handle.GetOpcode() != ValueOpcode::GetImageResource || handle.NumArgs() != 8u) return false;
@@ -3019,6 +3098,9 @@ private:
 			return SelectorMaximum(inst->Arg(arg), bounds, depth + 1u, guard);
 		};
 		switch (inst->GetOpcode()) {
+			case ValueOpcode::ReadLane:
+				// ReadLane ignores EXEC; only context-free bounds on every source lane apply.
+				return PossibleU32Bits(inst->Arg(0));
 			case ValueOpcode::Phi: {
 				const auto initial = InitialCandidateMask(value);
 				if (!initial.IsEmpty()) return SelectorMaximum(initial, bounds, depth + 1u, guard);
@@ -3240,10 +3322,11 @@ private:
 					uint32_t bad_dword = 0;
 					if (ValidateSource(descriptor, bad_dword)) return;
 					IndirectDescriptorPlan plan;
-					if (sampler ? TryMakeFiniteDescriptor(*handle, plan)
+					if (TryMakeNestedDescriptor(*handle, descriptor, plan) ||
+					    (sampler ? TryMakeFiniteDescriptor(*handle, plan)
 					             : (TryMakeIndirectImage(*handle, descriptor, plan) ||
 					                TryMakeFiniteDescriptor(*handle, plan) ||
-					                TryMakeIndirectBuffer(*handle, descriptor, plan))) {
+					                TryMakeIndirectBuffer(*handle, descriptor, plan)))) {
 						for (const auto& previous: m_indirect_descriptors) {
 							if (plan.reads[0] != nullptr && previous.reads[0] == plan.reads[0] &&
 							    !EquivalentValue(m_program, previous.key, plan.key))
@@ -3389,7 +3472,7 @@ private:
 		const auto* descriptor = Source(source);
 		if (descriptor != nullptr && descriptor->indirect_descriptor) {
 			const auto& sources = descriptor->indirect_descriptor->sources;
-			if (sources.empty()) Fail(pc, "finite sampler has no candidates");
+			if (sources.empty()) return root;
 			std::vector<uint32_t> resources {root};
 			for (size_t i = 1; i < sources.size(); ++i) {
 				const auto candidate = AddSampler(sources[i], pc);