shader: resolve packed buffer selectors with shared indexed stores

Co-authored-by: Ali Almohaya <AliAlmohaya@gmail.com>

src/graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h

@@ -398,6 +398,8 @@ uint32_t StorageBufferPackedStride(const EmitterState& state, const IR::MemoryIn
 
 Prospero::BufferFormat StorageBufferFormat(const EmitterState& state, const IR::MemoryInfo& mem);
 
+bool UsesSharedIndirectBufferStore(const EmitterState& state, uint32_t resource);
+
 void EmitMemoryOffsets(EmitterState& state);
 
 uint32_t LdsDwordCount(const EmitterState& state);
@@ -409,6 +411,7 @@ struct MemoryResourceAccess {
 	uint32_t              length           = 0;
 	uint32_t              byte_offset      = 0;
 	uint32_t              element_bits     = 32;
+	bool                  non_uniform      = false;
 	spv::MemoryAccessMask memory_access    = spv::MemoryAccessMaskNone;
 };
 
@@ -418,7 +421,8 @@ MemoryResourceAccess PrepareMemoryResourceAccess(EmitterState& state, const IR::
 
 MemoryResourceAccess PrepareStorageBufferResourceAccess(EmitterState&         state,
                                                         const IR::MemoryInfo& mem,
-                                                        BufferDefinition& buffer);
+                                                        BufferDefinition& buffer,
+                                                        uint32_t array_index = 0);
 
 uint32_t EmitMemoryElementInBounds(EmitterState& state, const MemoryResourceAccess& access,
                                    uint32_t index);

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

@@ -249,10 +249,11 @@ uint32_t LoadBda(ValueEmitContext& ctx, uint32_t address, uint32_t active, uint3
 	});
 }
 
-uint32_t ByteAddress(ValueEmitContext& ctx, const IR::Inst& inst, const IR::MemoryInfo& mem) {
+uint32_t ByteAddress(ValueEmitContext& ctx, const IR::Inst& inst, const IR::MemoryInfo& mem,
+                     uint32_t prefix = 0) {
 	if (mem.kind == IR::ResourceKind::Buffer) {
 		const auto resource = ctx.state.program.info.buffers.at(mem.resource).descriptor_index;
-		const auto prefix = ctx.state.memory_byte_offsets[resource];
+		if (prefix == 0u) prefix = ctx.state.memory_byte_offsets[resource];
 		const auto address = Binary(ctx.state, spv::OpIAdd, TypeU32(ctx.state),
 		                            BufferByteAddress(ctx, inst, mem), prefix);
 		// The host binding prefix must not wrap an out-of-range guest address into the buffer.
@@ -273,8 +274,9 @@ uint32_t ByteAddress(ValueEmitContext& ctx, const IR::Inst& inst, const IR::Memo
 	return ScratchByteAddress(ctx, mem, ctx.Arg(inst, 1), ctx.Arg(inst, 2));
 }
 
-uint32_t DwordIndex(ValueEmitContext& ctx, const IR::Inst& inst, const IR::MemoryInfo& mem) {
-	auto address = ByteAddress(ctx, inst, mem);
+uint32_t DwordIndex(ValueEmitContext& ctx, const IR::Inst& inst, const IR::MemoryInfo& mem,
+                    uint32_t prefix = 0) {
+	auto address = ByteAddress(ctx, inst, mem, prefix);
 	if (mem.kind == IR::ResourceKind::Lds || mem.kind == IR::ResourceKind::Gds) {
 		// RDNA2 DS region addresses retain bits [15:2] after adding the byte offset.
 		address = Binary(ctx.state, spv::OpBitwiseAnd, TypeU32(ctx.state), address,
@@ -540,8 +542,30 @@ void StoreLocalFlat(ValueEmitContext& ctx, const IR::Inst& inst) {
 
 void StoreWordPrepared(ValueEmitContext& ctx, const IR::Inst& inst, const IR::MemoryInfo& mem,
                        const MemoryResourceAccess& resource, uint32_t data) {
-	const auto index = DwordIndex(ctx, inst, mem);
-	EmitIfCondition(ctx.state, EmitMemoryElementInBounds(ctx.state, resource, index), [&]() {
+	auto& state = ctx.state;
+	const auto source = mem.kind == IR::ResourceKind::Buffer
+	                        ? state.program.info.buffers[mem.resource].source : UINT32_MAX;
+	const auto& sources = state.program.descriptor_sources;
+	const auto* indirect = source < sources.size() && sources[source].indirect_descriptor
+	                           ? &*sources[source].indirect_descriptor : nullptr;
+	uint32_t index = 0;
+	uint32_t in_bounds = 0;
+	if (indirect != nullptr && indirect->indexed_store) {
+		// The resource plan proves stride-four, index-only stores. Check the original
+		// index before any byte multiplication can wrap it into the buffer.
+		const auto prefix = Binary(state, spv::OpShiftRightLogical, TypeU32(state),
+		                           resource.byte_offset, ConstantU32(state, 2));
+		const auto guest_index = ctx.Arg(inst, 1);
+		const auto records = Binary(state, spv::OpISub, TypeU32(state), resource.length, prefix);
+		in_bounds = AndCondition(state,
+		    Binary(state, spv::OpULessThanEqual, TypeBool(state), prefix, resource.length),
+		    Binary(state, spv::OpULessThan, TypeBool(state), guest_index, records));
+		index = EmitAddU32(state, guest_index, prefix);
+	} else {
+		index = DwordIndex(ctx, inst, mem, resource.byte_offset);
+		in_bounds = EmitMemoryElementInBounds(state, resource, index);
+	}
+	EmitIfCondition(state, in_bounds, [&]() {
 		StoreWordInBounds(ctx, resource, index, data);
 	});
 }
@@ -576,6 +600,15 @@ void StoreWord(ValueEmitContext& ctx, const IR::Inst& inst, IR::MemoryInfo mem)
 		const auto selected = EmitIndirectResourceIndex(
 		    state, ctx.Arg(*handle, 0), buffer.indirect_mapping_offset,
 		    buffer.indirect_search_iterations, UINT32_MAX);
+		if (UsesSharedIndirectBufferStore(state, mem.resource)) {
+			EmitIfCondition(state, Binary(state, spv::OpULessThan, TypeBool(state), selected,
+			                             ConstantU32(state, buffer.indirect_resources.size())), [&]() {
+				const auto slot = EmitAddU32(state, selected, ConstantU32(state, buffer.descriptor_index));
+				const auto access = PrepareStorageBufferResourceAccess(state, mem, state.storage_buffers[2], slot);
+				StoreWordPrepared(ctx, inst, mem, access, data);
+			});
+			return;
+		}
 		const auto merge = state.builder.AllocateId();
 		std::vector<uint32_t> labels(buffer.indirect_resources.size());
 		std::vector<uint32_t> branches {spv::OpSwitch, selected, merge};

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

@@ -2,14 +2,13 @@
 
 namespace Libs::Graphics::ShaderRecompiler::Spirv::Emitter {
 
-uint32_t EmitShaderDataDwordLoad(EmitterState& state, uint32_t dword_index) {
+static uint32_t EmitShaderDataDwordLoadAt(EmitterState& state, uint32_t dword_index) {
 	if (state.program.bindings.UsesPushData()) {
-		dword_index += state.program.bindings.push_data_start_dword;
 		const auto pointer = state.builder.AllocateId();
 		const auto value   = state.builder.AllocateId();
 		state.builder.AddFunction(spv::OpAccessChain, TypePushConstantElementPointer(state),
 		                          pointer, state.push_constant_variable, ConstantU32(state, 0),
-		                          ConstantU32(state, dword_index));
+		                          dword_index);
 		state.builder.AddFunction(spv::OpLoad, TypeU32(state), value, pointer);
 		return value;
 	}
@@ -18,13 +17,20 @@ uint32_t EmitShaderDataDwordLoad(EmitterState& state, uint32_t dword_index) {
 		const auto value   = state.builder.AllocateId();
 		state.builder.AddFunction(spv::OpAccessChain, TypeStorageBufferElementPointer(state),
 		                          pointer, state.shader_data_storage_variable,
-		                          ConstantU32(state, 0), ConstantU32(state, dword_index));
+		                          ConstantU32(state, 0), dword_index);
 		state.builder.AddFunction(spv::OpLoad, TypeU32(state), value, pointer);
 		return value;
 	}
 	return ConstantU32(state, 0);
 }
 
+uint32_t EmitShaderDataDwordLoad(EmitterState& state, uint32_t dword_index) {
+	if (state.program.bindings.UsesPushData()) {
+		dword_index += state.program.bindings.push_data_start_dword;
+	}
+	return EmitShaderDataDwordLoadAt(state, ConstantU32(state, dword_index));
+}
+
 uint32_t EmitBinaryU32(EmitterState& state, spv::Op opcode, uint32_t lhs, uint32_t rhs) {
 	const auto ret = state.builder.AllocateId();
 	state.builder.AddFunction(opcode, TypeU32(state), ret, lhs, rhs);
@@ -102,12 +108,38 @@ Prospero::BufferFormat StorageBufferFormat(const EmitterState& state, const IR::
 	return state.program.info.buffers[mem.resource].descriptor_format;
 }
 
+bool UsesSharedIndirectBufferStore(const EmitterState& state, uint32_t resource) {
+	const auto& buffers = state.program.info.buffers;
+	const auto& root = buffers[resource];
+	if (root.indirect_root != resource || root.descriptor_index == UINT32_MAX ||
+	    root.packed_stride == 0u || root.indirect_resources.empty()) {
+		return false;
+	}
+	uint32_t slot = root.descriptor_index;
+	for (const auto candidate: root.indirect_resources) {
+		if (buffers[candidate].packed_stride != root.packed_stride ||
+		    buffers[candidate].descriptor_index != slot++) return false;
+	}
+	return true;
+}
+
 void EmitMemoryOffsets(EmitterState& state) {
+	std::array<bool, IR::ShaderInfo::MaxBuffers> dynamic_offsets {};
+	for (uint32_t resource = 0; resource < state.program.info.buffers.size(); ++resource) {
+		if (!UsesSharedIndirectBufferStore(state, resource)) continue;
+		for (const auto candidate: state.program.info.buffers[resource].indirect_resources) {
+			dynamic_offsets[state.program.info.buffers[candidate].descriptor_index] = true;
+		}
+	}
 	const auto count = state.program.bindings.descriptor_counts[static_cast<size_t>(IR::DescriptorBindingKind::Buffers)];
 	for (uint32_t i = 0; i < count; i += 4u) {
+		const auto lanes = std::min(4u, count - i);
+		if (std::all_of(dynamic_offsets.begin() + i, dynamic_offsets.begin() + i + lanes,
+		                [](bool dynamic) { return dynamic; })) continue;
 		const auto word =
 		    EmitShaderDataDwordLoad(state, state.program.bindings.memory_offset_dword + i / 4u);
-		for (uint32_t lane = 0; lane < std::min(4u, count - i); ++lane) {
+		for (uint32_t lane = 0; lane < lanes; ++lane) {
+			if (dynamic_offsets[i + lane]) continue;
 			state.memory_byte_offsets[i + lane] = EmitNative<spv::OpBitFieldUExtract, IR::Type::U32>(
 			    state, word, ConstantU32(state, lane * 8u), ConstantU32(state, 8u));
 		}
@@ -152,7 +184,8 @@ void EnsureLdsStorage(EmitterState& state) {
 
 MemoryResourceAccess PrepareStorageBufferResourceAccess(EmitterState& state,
                                                          const IR::MemoryInfo& mem,
-                                                         BufferDefinition& buffer) {
+                                                         BufferDefinition& buffer,
+                                                         uint32_t array_index) {
 	if (buffer.variable == 0) {
 		ExitDescriptorBindingFailure(state, IR::DescriptorBindingKind::Buffers, mem.resource,
 		                             "storage buffer descriptor array was not emitted");
@@ -160,15 +193,34 @@ MemoryResourceAccess PrepareStorageBufferResourceAccess(EmitterState& state,
 	if (buffer.pointer_type == 0) {
 		buffer.pointer_type = TypePointer(state, spv::StorageClassStorageBuffer, buffer.type);
 	}
-	const auto array_index =
-	    state.program.info.buffers.at(mem.resource).descriptor_index;
+	const bool non_uniform = array_index != 0u;
+	const auto static_index = state.program.info.buffers.at(mem.resource).descriptor_index;
+	if (!non_uniform) array_index = ConstantU32(state, static_index);
 	MemoryResourceAccess access {
 	    .kind = mem.kind,
+	    .non_uniform = non_uniform,
 	    .memory_access = mem.coherent ? spv::MemoryAccessVolatileMask : spv::MemoryAccessMaskNone};
 	access.object_pointer = state.builder.AllocateId();
 	state.builder.AddFunction(spv::OpAccessChain, buffer.pointer_type, access.object_pointer,
-	                          buffer.variable, ConstantU32(state, array_index));
-	access.byte_offset = state.memory_byte_offsets[array_index];
+	                          buffer.variable, array_index);
+	if (non_uniform) {
+		state.builder.RequireExtension("SPV_EXT_descriptor_indexing");
+		state.builder.RequireCapability(spv::CapabilityShaderNonUniform);
+		state.builder.RequireCapability(spv::CapabilityStorageBufferArrayNonUniformIndexing);
+		state.builder.AddAnnotation(spv::OpDecorate, access.object_pointer, spv::DecorationNonUniform);
+		const auto offsets_start = state.program.bindings.memory_offset_dword +
+		    (state.program.bindings.UsesPushData() ? state.program.bindings.push_data_start_dword : 0u);
+		const auto word = EmitShaderDataDwordLoadAt(state, EmitAddU32(
+		    state, ConstantU32(state, offsets_start),
+		    Binary(state, spv::OpShiftRightLogical, TypeU32(state), array_index, ConstantU32(state, 2))));
+		const auto shift = Binary(state, spv::OpShiftLeftLogical, TypeU32(state),
+		    Binary(state, spv::OpBitwiseAnd, TypeU32(state), array_index, ConstantU32(state, 3)),
+		    ConstantU32(state, 3));
+		access.byte_offset = EmitNative<spv::OpBitFieldUExtract, IR::Type::U32>(
+		    state, word, shift, ConstantU32(state, 8));
+	} else {
+		access.byte_offset = state.memory_byte_offsets[static_index];
+	}
 	access.length      = state.builder.AllocateId();
 	state.builder.AddFunction(spv::OpArrayLength, TypeU32(state), access.length,
 	                          access.object_pointer, 0);
@@ -260,6 +312,9 @@ uint32_t EmitStorageBufferElementPointer(EmitterState& state,
 	const auto pointer = state.builder.AllocateId();
 	state.builder.AddFunction(spv::OpAccessChain, pointer_type, pointer, access.object_pointer,
 	                          ConstantU32(state, 0), index);
+	if (access.non_uniform) {
+		state.builder.AddAnnotation(spv::OpDecorate, pointer, spv::DecorationNonUniform);
+	}
 	return pointer;
 }
 

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

@@ -504,6 +504,7 @@ struct DescriptorSource {
 			uint32_t source = UINT32_MAX;
 			uint32_t stride = 0;
 			uint32_t offset = 0;
+			std::vector<std::pair<uint32_t, uint32_t>> byte_ranges;
 
 			bool operator==(const SelectorRead& other) const = default;
 		};
@@ -514,6 +515,7 @@ struct DescriptorSource {
 		uint32_t table_stride    = 0;
 		uint32_t table_record_bytes = 0;
 		bool     table_scalar = false;
+		bool     indexed_store = false;
 		uint32_t workgroup_axis  = UINT32_MAX;
 		Value    key_count;
 		Value                 selector_first;

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

@@ -195,15 +195,15 @@ bool ReadScalarTable(uint64_t base, uint64_t size, uint64_t dynamic_offset,
 	       runtime.read_specialization_memory(runtime.userdata, address, prefix);
 }
 
-bool NormalizeIndirectStoreBuffer(DescriptorValue& value) {
+bool NormalizeIndirectStoreBuffer(DescriptorValue& value, bool indexed) {
 	ShaderBufferResource descriptor;
 	if (!DecodeBufferDescriptor(value, descriptor)) return false;
 	if (descriptor.Type() != 0u || descriptor.Format() == Prospero::BufferFormat::kInvalid ||
 	    descriptor.Base48() == 0u || descriptor.GetSize() == 0u) {
 		value.dwords.fill(0u);
-	} else if (descriptor.Stride() != 1u || descriptor.SwizzleEnabled() || descriptor.AddTid() ||
-	           descriptor.OutOfBounds() != 3u) {
-		return SpecializationFail("indirect buffer stores require linear byte bounds");
+	} else if (descriptor.Stride() != (indexed ? 4u : 1u) || descriptor.SwizzleEnabled() ||
+	           descriptor.AddTid() || descriptor.OutOfBounds() != (indexed ? 0u : 3u)) {
+		return SpecializationFail("indirect buffer stores require exact linear DWORD or byte bounds");
 	}
 	return true;
 }
@@ -337,6 +337,21 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 			}
 			keys.resize(key_count);
 			std::iota(keys.begin(), keys.end(), 0u);
+		} else if (!selector->byte_ranges.empty()) {
+			if (material_value.dword_count != 2u) return false;
+			const auto base = (uint64_t {material_value.dwords[1]} << 32u) | material_value.dwords[0];
+			std::array<bool, 256> selected {};
+			for (const auto [offset, size]: selector->byte_ranges) {
+				if (size > MaxIndirectDescriptorProbes || ((offset | size) & 3u) != 0u) return false;
+				keys.resize(size / 4u);
+				if (!ReadScalarTable(base, UINT64_MAX, offset, runtime, keys)) return false;
+				for (const auto word: keys)
+					for (uint32_t shift = 0; shift < 32u; shift += 8u)
+						selected[(word >> shift) & 255u] = true;
+			}
+			keys.clear();
+			for (uint32_t key = 0; key < selected.size(); ++key)
+				if (selected[key]) keys.push_back(key);
 		} else if (!indirect.selector_first.IsEmpty()) {
 			ShaderBufferResource material;
 			uint32_t             first = 0, count = 0;
@@ -380,7 +395,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 				mask &= mask - 1u;
 			}
 		} else return false;
-		if (selector != nullptr) {
+		if (selector != nullptr && selector->byte_ranges.empty()) {
 			std::ranges::sort(keys);
 			keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
 		}
@@ -1132,7 +1147,9 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 					if (!MaterializeIndirectDescriptor(
 					        program, *source->indirect_descriptor, i, 4u, observed, clean, snapshot,
 					        snapshot.buffers, specialization.buffers, ShaderInfo::MaxBuffers,
-					        NormalizeIndirectStoreBuffer))
+					        [&](DescriptorValue& value) {
+						        return NormalizeIndirectStoreBuffer(value, source->indirect_descriptor->indexed_store);
+					        }))
 						return false;
 				}
 			} else if (!evaluate(base.source, snapshot.buffers[i], base.written)) {

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

@@ -909,7 +909,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.table_stride != b.table_stride || a.indexed_store != b.indexed_store ||
 				    a.workgroup_axis != b.workgroup_axis || a.sources != b.sources ||
 				    !EquivalentValue(m_program, a.key_count, b.key_count) ||
 				    a.selector_first.IsEmpty() != b.selector_first.IsEmpty() ||
@@ -1821,17 +1821,71 @@ private:
 		return true;
 	}
 
+	bool MatchPackedScalarKey(Value key, DescriptorSource::IndirectDescriptor& indirect) {
+		const auto* mask = key.Resolve().TryInstruction();
+		uint32_t bits = 0;
+		if (mask == nullptr || mask->GetOpcode() != ValueOpcode::BitwiseAnd32 ||
+		    !ImmediateU32(mask->Arg(1), bits) || bits != 255u) return false;
+		const auto* shift = mask->Arg(0).Resolve().TryInstruction();
+		if (shift == nullptr || shift->GetOpcode() != ValueOpcode::ShiftRightLogical32 ||
+		    (PossibleU32Bits(shift->Arg(1)) & ~24u) != 0u) return false;
+		DescriptorSource::IndirectDescriptor::SelectorRead selector;
+		std::vector<const Inst*> pending {shift->Arg(0).Resolve().TryInstruction()};
+		std::vector<const Inst*> visited;
+		while (!pending.empty()) {
+			const auto* read = pending.back();
+			pending.pop_back();
+			if (read == nullptr) return false;
+			if (std::ranges::find(visited, read) != visited.end()) continue;
+			visited.push_back(read);
+			if (read->GetOpcode() == ValueOpcode::Phi) {
+				for (size_t arm = 0; arm < read->NumArgs(); ++arm)
+					pending.push_back(read->Arg(arm).Resolve().TryInstruction());
+				continue;
+			}
+			uint32_t index = 0;
+			const auto* memory = ScalarReadMemory(*read, index);
+			DescriptorSource source;
+			if (memory == nullptr || memory->kind != ResourceKind::ScalarAddress ||
+			    memory->offset > INT32_MAX || !MakeRuntimeTableSource(*read, source)) return false;
+			const auto maximum = PossibleU32Bits(read->Arg(1)) & ~3u;
+			const auto offset = memory->offset & ~3u;
+			if (uint64_t {offset} + maximum + 4u > UINT32_MAX) return false;
+			const auto source_index = InternSource(source);
+			if (selector.source != UINT32_MAX && selector.source != source_index) return false;
+			selector.source = source_index;
+			selector.byte_ranges.emplace_back(offset, maximum + 4u);
+		}
+		if (selector.byte_ranges.empty()) return false;
+		// Scalar reads align both offsets independently. Merge the possible read spans;
+		// the GPU retains the actual byte selection and the host only discovers V#s.
+		std::ranges::sort(selector.byte_ranges);
+		for (size_t i = 1; i < selector.byte_ranges.size();) {
+			auto& previous = selector.byte_ranges[i - 1u];
+			const auto current = selector.byte_ranges[i];
+			if (current.first > previous.first + previous.second) { ++i; continue; }
+			previous.second = std::max(previous.second, current.first + current.second - previous.first);
+			selector.byte_ranges.erase(selector.byte_ranges.begin() + i);
+		}
+		indirect.selector = std::move(selector);
+		return true;
+	}
+
 	bool TryMakeIndirectBuffer(Inst& handle, const DescriptorSource& descriptor,
 	                           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;
+		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;
 			const auto& memory = m_program.memory_info[inst.Flags<MemoryFlags>().index];
 			uint32_t scalar = 1u;
 			if (memory.kind != ResourceKind::Buffer || memory.typed || memory.formatted ||
-			    memory.idxen || !memory.offen || !ImmediateU32(inst.Arg(3), scalar) || scalar != 0u)
+			    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;
 		}
 		DescriptorSource table_source;
@@ -1840,8 +1894,26 @@ 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 ||
-		    table_source.dword_count != 4u || table_offset != 0u || table_immediate != 0u) return false;
+		                          table_stride, table_immediate) || table_stride != 16u) return false;
+		DescriptorSource::IndirectDescriptor indirect;
+		indirect.table_stride = table_stride;
+		indirect.table_offset = table_offset;
+		indirect.table_immediate = table_immediate;
+		indirect.indexed_store = indexed;
+		const auto finish = [&] {
+			indirect.table_source = InternSource(table_source);
+			DescriptorSource source;
+			source.dword_count = 4u;
+			source.dwords.fill(Value(0u));
+			source.indirect_descriptor = std::move(indirect);
+			plan.handle = &handle;
+			plan.source = InternSource(source);
+			plan.key = key;
+			return true;
+		};
+		if (table_source.dword_count == 2u && MatchPackedScalarKey(key, indirect)) return finish();
+		if (table_source.dword_count != 4u || table_offset != 0u || table_immediate != 0u || indexed)
+			return false;
 		const auto* lane_read = key.Resolve().TryInstruction();
 		if (lane_read == nullptr || lane_read->GetOpcode() != ValueOpcode::ReadFirstLane)
 			return false;
@@ -1869,21 +1941,11 @@ private:
 		    !ImmediateU32(read->Arg(3), scalar) || scalar != 0u ||
 		    !EquivalentValue(m_program, selected->Arg(0), read->Arg(4))) return false;
 		DescriptorSource material_source;
-		DescriptorSource::IndirectDescriptor indirect;
-		indirect.table_stride = table_stride;
 		if (!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});
-		indirect.table_source = InternSource(table_source);
-		DescriptorSource source;
-		source.dword_count = 4u;
-		source.dwords.fill(Value(0u));
-		source.indirect_descriptor = indirect;
-		plan.handle = &handle;
-		plan.source = InternSource(source);
-		plan.key = key;
-		return true;
+		return finish();
 	}
 
 	using SelectorBounds = std::span<const std::pair<const Inst*, uint32_t>>;