shader: share bounded record descriptors for packed loads and byte stores

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

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

@@ -537,7 +537,7 @@ struct DescriptorSource {
 		uint32_t table_stride    = 0;
 		uint32_t table_record_bytes = 0;
 		bool     table_scalar = false;
-		bool     indexed_store = false;
+		bool     indexed = false;
 		bool     signed_count = false;
 		uint32_t workgroup_axis  = UINT32_MAX;
 		Value    key_count;

src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.inc

@@ -235,6 +235,7 @@ VALUE_OPCODE(BpermuteU32, U32, U32, U32, U1)
 // dense bindings are assigned by resource tracking.
 VALUE_OPCODE(GetSrtResource, SrtResource)
 VALUE_OPCODE(GetBufferResource, BufferResource, U32, U32, U32, U32)
+VALUE_OPCODE(GetBufferResourceIndex, U32, U32)
 VALUE_OPCODE(GetAddressResource, AddressResource, U32, U32)
 VALUE_OPCODE(GetScratchResource, AddressResource)
 VALUE_OPCODE(GetImageResource, ImageResource, U32, U32, U32, U32, U32, U32, U32, U32)

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

@@ -227,16 +227,38 @@ bool ReadScalarTable(uint64_t base, uint64_t size, uint64_t offset,
 	return ReadTable(base & AddressMask & ~uint64_t {3}, size, offset & ~uint64_t {3}, runtime, words);
 }
 
-bool NormalizeIndirectStoreBuffer(DescriptorValue& value, bool indexed, bool formatted) {
+bool NormalizeIndirectBuffer(DescriptorValue& value, bool indexed, const BufferResource& resource) {
 	ShaderBufferResource descriptor;
 	if (!DecodeBufferDescriptor(value, descriptor)) return false;
+	const bool formatted_read = resource.formatted && resource.read && !resource.written;
 	if (descriptor.Type() != 0u || descriptor.Format() == Prospero::BufferFormat::kInvalid ||
-	    descriptor.Base48() == 0u || descriptor.GetSize() == 0u) {
+	    (!formatted_read && (descriptor.Base48() == 0u || descriptor.GetSize() == 0u))) {
 		value.dwords.fill(0u);
-	} else if (descriptor.Stride() != (indexed ? 4u : 1u) || descriptor.SwizzleEnabled() ||
-	           descriptor.AddTid() || descriptor.OutOfBounds() != (indexed ? 0u : 3u)) {
+		return true;
+	}
+	if (descriptor.SwizzleEnabled() || descriptor.AddTid() ||
+	    descriptor.OutOfBounds() != (indexed ? 0u : 3u)) {
+		return SpecializationFail("indirect buffers require exact linear index or byte bounds");
+	}
+	if (formatted_read) {
+		const auto format = Format::GetFormatInfo(descriptor.Format());
+		if (!indexed || format.type == Format::ComponentType::Unknown)
+			return SpecializationFail("indirect formatted loads require a supported indexed format");
+		if (descriptor.Base48() == 0u || descriptor.Stride() == 0u || descriptor.NumRecords() == 0u) {
+			// Every index is out of bounds. Preserve the format and swizzle constants.
+			value.dwords[1] &= ~(0x3fffu << 16u);
+			value.dwords[2] = 0u;
+			return true;
+		}
+		const auto alignment = std::min(4u, format.byte_size);
+		if (descriptor.Stride() < format.byte_size ||
+		    (descriptor.Stride() & (alignment - 1u)) != 0u ||
+		    (descriptor.Base48() & (alignment - 1u)) != 0u)
+			return SpecializationFail("indirect formatted loads require aligned full-format records");
+	} else if ((!indexed && descriptor.Stride() != 1u) ||
+	           (indexed && resource.max_byte_extent != 1u && descriptor.Stride() != 4u)) {
 		return SpecializationFail("indirect buffer stores require exact linear DWORD or byte bounds");
-	} else if (formatted && (descriptor.Format() != Prospero::BufferFormat::k32UInt ||
+	} else if (resource.formatted && (descriptor.Format() != Prospero::BufferFormat::k32UInt ||
 	                         descriptor.DstSelXYZW() != DstSel(4, 0, 0, 0) ||
 	                         (descriptor.Base48() & 3u) != 0u)) {
 		return SpecializationFail("indirect formatted X stores require an aligned 32-bit UInt identity mapping");
@@ -248,7 +270,8 @@ bool IsBoundedDescriptorTable(const ResourcePlan& program,
                               const DescriptorSource::IndirectDescriptor& indirect) {
 	const auto* table = Source(program, indirect.table_source);
 	return indirect.sources.empty() && !indirect.selector &&
-	       indirect.workgroup_axis == UINT32_MAX && table != nullptr && table->dword_count == 4u;
+	       indirect.key_count.IsEmpty() && indirect.workgroup_axis == UINT32_MAX &&
+	       table != nullptr && table->dword_count == 4u;
 }
 
 template <typename Specialization, typename Normalize>
@@ -372,7 +395,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 				key_count = runtime.workgroup_counts[indirect.workgroup_axis];
 				if (key_count == 0u) return false;
 			} else {
-				if (table_value.dword_count != 2u || !evaluate_count(key_count))
+				if (!evaluate_count(key_count))
 					return false;
 			}
 			if (!indirect.selector_first.IsEmpty() && !clean.Evaluate(indirect.selector_first, first_key))
@@ -1308,8 +1331,8 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 					        program, *source->indirect_descriptor, i, 4u, observed, clean, snapshot,
 					        snapshot.buffers, specialization.buffers, ShaderInfo::MaxBuffers,
 					        [&](DescriptorValue& value) {
-						        return NormalizeIndirectStoreBuffer(
-						            value, source->indirect_descriptor->indexed_store, base.formatted);
+						        return NormalizeIndirectBuffer(
+						            value, source->indirect_descriptor->indexed, base);
 					        }))
 						return false;
 				}
@@ -1526,7 +1549,8 @@ void ApplyResourceSpecialization(Program& program, const ResourceSpecialization&
 			if (BufferAccessOf(inst.GetOpcode()) == BufferAccess::Read) {
 				const auto& memory = memory_info[inst.Flags<MemoryFlags>().index];
 				if (memory.kind == ResourceKind::Buffer &&
-				    specialization.buffers[memory.resource].zero_stride_oob) {
+				    specialization.buffers[memory.resource].zero_stride_oob &&
+				    buffers[memory.resource].indirect_root == BufferResource::NoIndirectBuffer) {
 					// Bounds mode 0 checks offset >= stride, so zero stride
 					// makes every vector read out of bounds regardless of its address.
 					const auto count = BufferComponentCount(inst.GetOpcode());

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

@@ -394,6 +394,7 @@ public:
 				memory.sampler = patch.sampler;
 			}
 		}
+		std::vector<std::pair<const IndirectDescriptorPlan*, Value>> buffer_selections;
 		for (const auto& plan: m_indirect_descriptors) {
 			auto key = plan.key;
 			if (plan.table_read != nullptr) {
@@ -452,8 +453,32 @@ public:
 					plan.handle->SetArg(word, Value(0u));
 				continue;
 			}
-			if (plan.handle->NumArgs() == 4u) {
-				plan.handle->SetArg(0, plan.key);
+			if (plan.handle->GetOpcode() == ValueOpcode::GetBufferResource) {
+				Value selected;
+				for (const auto& [previous, value]: buffer_selections) {
+					if (plan.selection_anchor != nullptr && previous->reads == plan.reads &&
+					    previous->source == plan.source &&
+					    previous->handle->Flags<uint32_t>() == plan.handle->Flags<uint32_t>() &&
+					    EquivalentValue(m_program, previous->key, plan.key) &&
+					    AvailableBefore(*value.Instruction(), *plan.handle)) {
+						selected = value;
+						break;
+					}
+				}
+				if (selected.IsEmpty()) {
+					const auto* anchor = plan.selection_anchor != nullptr ? plan.selection_anchor : plan.handle;
+					auto* block = anchor->Parent();
+					const auto where = std::ranges::find_if(block->Instructions(),
+					    [&](const Inst& inst) { return &inst == anchor; });
+					auto& selection = *block->PrependNewInst(where, ValueOpcode::GetBufferResourceIndex, {plan.key});
+					selection.SetFlags(plan.handle->Flags<uint32_t>());
+					selected = Value(&selection);
+					buffer_selections.emplace_back(&plan, selected);
+				}
+				plan.handle->SetArg(0, selected);
+				for (uint32_t word = 1; word < 4u; ++word) plan.handle->SetArg(word, Value(0u));
+			} else if (plan.handle->NumArgs() == 4u) {
+				plan.handle->SetArg(0, key);
 				for (uint32_t word = 1; word < 4u; ++word) plan.handle->SetArg(word, Value(0u));
 			} else {
 				if (plan.reads[0] == nullptr) plan.handle->SetArg(0, key);
@@ -519,6 +544,7 @@ private:
 		bool retain_reads = false;
 		bool table_indexed = false;
 		const Inst* table_read = nullptr;
+		const Inst* selection_anchor = nullptr;
 	};
 
 	[[noreturn]] void Fail(uint32_t pc, const std::string& reason) const {
@@ -931,7 +957,7 @@ private:
 				const auto& b = *descriptor.indirect_descriptor;
 				if (a.selector != b.selector || 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_store != b.indexed_store ||
+				    a.table_stride != b.table_stride || a.indexed != b.indexed ||
 				    a.workgroup_axis != b.workgroup_axis || a.sources != b.sources ||
 				    a.signed_count != b.signed_count ||
 				    !EquivalentValue(m_program, a.key_count, b.key_count) ||
@@ -2522,6 +2548,25 @@ private:
 		return true;
 	}
 
+	bool IndirectBufferAccess(const Inst& inst) const {
+		const auto opcode = inst.GetOpcode();
+		if (opcode != ValueOpcode::StoreBufferU32 && opcode != ValueOpcode::StoreBufferU8 &&
+		    opcode != ValueOpcode::LoadBufferU32x2) return false;
+		const auto index = inst.Flags<MemoryFlags>().index;
+		if (index >= m_program.memory_info.size()) return false;
+		const auto& memory = m_program.memory_info[index];
+		if (memory.kind != ResourceKind::Buffer || memory.typed || memory.idxen == memory.offen ||
+		    (memory.idxen && (memory.offset != 0u || inst.Arg(2).Resolve() != Value(0u))) ||
+		    inst.Arg(3).Resolve() != Value(0u)) return false;
+		if (opcode == ValueOpcode::StoreBufferU32)
+			return !memory.formatted || (memory.idxen && memory.data_bits == 32u && memory.data_dwords == 1u);
+		if (!memory.idxen) return false;
+		if (opcode == ValueOpcode::StoreBufferU8)
+			return !memory.formatted && memory.data_bits == 8u;
+		return memory.formatted && memory.data_bits == 16u && memory.data_dwords == 2u &&
+		       memory.component_count == 4u;
+	}
+
 	bool MatchDescriptorTable(Inst& handle, const DescriptorSource& descriptor,
 	                          IndirectDescriptorPlan& plan,
 	                          DescriptorSource& table_source, Value& key, uint32_t& table_offset,
@@ -2599,10 +2644,10 @@ private:
 			if (handle.NumArgs() == 8u) {
 				if (memory->kind == ResourceKind::ScalarAddress)
 					plan.retain_reads |= !UsesOnlyImageDescriptors(*read);
-			} else if (std::ranges::any_of(read->Uses(), [](const Use& use) {
+			} else if (std::ranges::any_of(read->Uses(), [&](const Use& use) {
 				return use.user->GetOpcode() != ValueOpcode::GetBufferResource ||
-				       std::ranges::any_of(use.user->Uses(), [](const Use& consumer) {
-					       return consumer.user->GetOpcode() != ValueOpcode::StoreBufferU32;
+				       std::ranges::any_of(use.user->Uses(), [&](const Use& consumer) {
+					       return !IndirectBufferAccess(*consumer.user);
 				       });
 			    })) {
 				return false;
@@ -2780,19 +2825,14 @@ private:
 	                           IndirectDescriptorPlan& plan) {
 		if (handle.GetOpcode() != ValueOpcode::GetBufferResource || handle.NumArgs() != 4u ||
 		    handle.Uses().empty()) return false;
-		if (handle.Uses().front().user->GetOpcode() != ValueOpcode::StoreBufferU32) return false;
+		if (!IndirectBufferAccess(*handle.Uses().front().user)) return false;
+		const bool legacy_store = handle.Uses().front().user->GetOpcode() == ValueOpcode::StoreBufferU32;
 		const auto indexed = m_program.memory_info[handle.Uses().front().user->Flags<MemoryFlags>().index].idxen;
 		for (const auto& use: handle.Uses()) {
 			const auto& inst = *use.user;
-			if (inst.GetOpcode() != ValueOpcode::StoreBufferU32) return false;
+			if (!IndirectBufferAccess(inst)) return false;
 			const auto& memory = m_program.memory_info[inst.Flags<MemoryFlags>().index];
-			uint32_t scalar = 1u;
-			if (memory.kind != ResourceKind::Buffer || memory.typed ||
-			    (memory.formatted && (!indexed || memory.data_bits != 32u || memory.data_dwords != 1u)) ||
-			    memory.idxen != indexed || memory.offen == indexed ||
-			    (indexed && (memory.offset != 0u || inst.Arg(2).Resolve() != Value(0u))) ||
-			    !ImmediateU32(inst.Arg(3), scalar) || scalar != 0u)
-				return false;
+			if (memory.idxen != indexed) return false;
 		}
 		DescriptorSource table_source;
 		Value key;
@@ -2800,12 +2840,25 @@ private:
 		uint32_t table_stride = 0;
 		uint32_t table_immediate = 0;
 		if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset,
-		                          table_stride, table_immediate) || table_stride != 16u) return false;
+		                          table_stride, table_immediate) ||
+		    table_stride < 16u || (table_stride & 3u) != 0u) return false;
 		DescriptorSource::IndirectDescriptor indirect;
 		indirect.table_stride = table_stride;
 		indirect.table_offset = table_offset;
 		indirect.table_immediate = table_immediate;
-		indirect.indexed_store = indexed;
+		indirect.indexed = indexed;
+		for (uint32_t word = 0; word < 4u; ++word) {
+			const auto* anchor = plan.reads[word];
+			const auto* key_inst = key.Resolve().TryInstruction();
+			if (AvailableBefore(*anchor, handle) &&
+			    (key_inst == nullptr || (key_inst != anchor && AvailableBefore(*key_inst, *anchor))) &&
+			    std::all_of(plan.reads.begin(), plan.reads.begin() + 4u,
+			                [&](const Inst* read) { return AvailableBefore(*read, *anchor); })) {
+				plan.selection_anchor = anchor;
+				break;
+			}
+		}
+		if (!legacy_store && plan.selection_anchor == nullptr) return false;
 		const auto finish = [&] {
 			indirect.table_source = InternSource(table_source);
 			DescriptorSource source;
@@ -2817,19 +2870,20 @@ private:
 			plan.key = key;
 			return true;
 		};
-		if (table_source.dword_count == 2u && MatchPackedScalarKey(key, indirect)) return finish();
+		if (const auto* bound = BoundedLoop(key, handle.Parent()); bound != nullptr) {
+			indirect.key_count = bound->Arg(1);
+			indirect.signed_count = bound->GetOpcode() == ValueOpcode::SLessThan32;
+			return finish();
+		}
+		if (!legacy_store || table_stride != 16u) return false;
 		if (table_source.dword_count == 2u) {
-			const auto* bound = BoundedLoop(key, handle.Parent());
-			if (bound != nullptr) {
-				indirect.key_count = bound->Arg(1);
-				indirect.signed_count = bound->GetOpcode() == ValueOpcode::SLessThan32;
-			} else {
+			if (!MatchPackedScalarKey(key, indirect)) {
 				// EXEC implications must describe the same iteration as the selected value.
 				if (m_native_cfg.irreducible || !m_native_cfg.back_edges.empty()) return false;
 				const auto* lane = key.Resolve().TryInstruction();
 				if (lane == nullptr || lane->GetOpcode() != ValueOpcode::ReadFirstLane ||
 				    !std::ranges::all_of(handle.Uses(), [&](const Use& use) {
-					    return Implies(use.user->Arg(5), lane->Arg(1));
+					    return Implies(use.user->Arg(use.user->NumArgs() - 1u), lane->Arg(1));
 				    })) return false;
 				const auto maximum = SelectorMaximum(lane->Arg(0), {}, 0u, lane->Arg(1));
 				if (maximum >= MaxIndirectDescriptorProbes) return false;
@@ -3112,7 +3166,7 @@ private:
 		for (auto* block: m_program.blocks) {
 			for (auto& inst: *block) {
 				if ((ImageOpcodeInfoOf(inst.GetOpcode()).access == ImageAccess::None &&
-				     inst.GetOpcode() != ValueOpcode::StoreBufferU32) ||
+				     !IndirectBufferAccess(inst)) ||
 				    inst.NumArgs() == 0u) {
 					continue;
 				}
@@ -3138,7 +3192,7 @@ private:
 						m_indirect_descriptors.push_back(std::move(plan));
 					}
 				};
-				plan_handle(inst.Arg(0), inst.GetOpcode() == ValueOpcode::StoreBufferU32 ? 4u : 8u,
+				plan_handle(inst.Arg(0), IndirectBufferAccess(inst) ? 4u : 8u,
 				            memory.resource * 4u, false);
 				if (inst.GetOpcode() == ValueOpcode::ImageSampleRaw &&
 				    (memory.image_sample_flags & Decoder::ImageSampleFlagAdjust) == 0u)

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

@@ -343,9 +343,16 @@ void Visit(Program& program, ShaderStageInputInfo input_info, InputUsage& inputs
 				if (memory.resource >= info.buffers.size()) {
 					return Fail("buffer operation has invalid resource metadata");
 				}
-				const auto bits = StorageBufferElementBits(program, memory);
-				info.buffer_u8 |= bits == 8u;
-				info.buffer_u16 |= bits == 16u;
+				const auto collect_element = [&](uint32_t resource) {
+					auto selected = memory;
+					selected.resource = resource;
+					const auto bits = StorageBufferElementBits(program, selected);
+					info.buffer_u8 |= bits == 8u;
+					info.buffer_u16 |= bits == 16u;
+				};
+				collect_element(memory.resource);
+				for (const auto child: info.buffers[memory.resource].indirect_resources)
+					collect_element(child);
 				if ((info.buffers[memory.resource].packed_stride & (1u << 20u)) != 0u) {
 					if (program.stage != ShaderType::Compute) {
 						return Fail("buffer ADD_TID is only valid for compute shaders");