shader: share proven selector domains for paired images and samplers

src/graphics/shader/recompiler/ir/Program.cpp

@@ -164,6 +164,16 @@ bool EquivalentValue(const ResourcePlan& program, Value left, Value right) {
 	return EquivalentValue(program, left, right, visited);
 }
 
+bool EquivalentSelectorDomain(const ResourcePlan& program,
+                              const DescriptorSource::IndirectDescriptor& left,
+                              const DescriptorSource::IndirectDescriptor& right) {
+	return left.selector == right.selector && left.signed_count == right.signed_count &&
+	       left.workgroup_axis == right.workgroup_axis &&
+	       EquivalentValue(program, left.key_count, right.key_count) &&
+	       EquivalentValue(program, left.selector_first, right.selector_first) &&
+	       EquivalentValue(program, left.selector_mask, right.selector_mask);
+}
+
 Value ResolveInvariantPhi(const ResourcePlan& program, Value value) {
 	value            = value.Resolve();
 	const auto* root = value.TryInstruction();

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

@@ -707,6 +707,9 @@ bool        HasShaderMemoryWrites(const Program& program);
 void  ValidateProgram(const Program& program, bool require_ssa);
 void  ResolveControlFlowIdentities(Program& program);
 bool  EquivalentValue(const ResourcePlan& program, Value left, Value right);
+bool  EquivalentSelectorDomain(const ResourcePlan& program,
+                               const DescriptorSource::IndirectDescriptor& left,
+                               const DescriptorSource::IndirectDescriptor& right);
 Value ResolveInvariantPhi(const ResourcePlan& program, Value value);
 Value ResolveActiveU32(Value value, Value active);
 

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

@@ -1300,11 +1300,7 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
 		uint32_t group = 0;
 		for (; group < plan.selector_keys.size(); ++group) {
 			const auto& b = *program.descriptor_sources[plan.selector_keys[group].source].indirect_descriptor;
-			if (a.selector == b.selector && a.signed_count == b.signed_count &&
+			if (EquivalentSelectorDomain(program, a, b)) break;
-			    a.workgroup_axis == b.workgroup_axis &&
-			    EquivalentValue(program, a.key_count, b.key_count) &&
-			    EquivalentValue(program, a.selector_first, b.selector_first) &&
-			    EquivalentValue(program, a.selector_mask, b.selector_mask)) break;
 		}
 		if (group == plan.selector_keys.size()) plan.selector_keys.push_back({.source = source});
 		plan.selector_groups[source] = group;

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

@@ -2851,7 +2851,8 @@ private:
 	}
 
 	bool TryMakeTableDescriptor(Inst& handle, const DescriptorSource& descriptor,
-	                            IndirectDescriptorPlan& plan) {
+	                            IndirectDescriptorPlan& plan,
+	                            const IndirectDescriptorPlan* paired_sampler) {
 		const bool sampler = handle.GetOpcode() == ValueOpcode::GetSamplerResource;
 		if ((!sampler && handle.GetOpcode() != ValueOpcode::GetImageResource) ||
 		    handle.NumArgs() != (sampler ? 4u : 8u)) return false;
@@ -2934,15 +2935,39 @@ private:
 			// Keep the GPU selector and descriptor reads, and preserve the proven scalar
 			// offset domain within the bounded image table.
 			if (sampler || plan.split_offsets) return false;
-			indirect.workgroup_axis = UINT32_MAX;
+			bool shared_selector = false;
-			plan.table_read = plan.reads[0];
+			if (!plan.table_indexed && paired_sampler != nullptr &&
-			if (plan.table_indexed) {
+			    m_sources[paired_sampler->source].indirect_descriptor->selector &&
-				indirect.table_stride = 0u;
+			    EquivalentValue(m_program, key, paired_sampler->key) &&
-				indirect.table_offset = 0u;
+			    std::ranges::all_of(handle.Uses(), [&](const Use& use) {
-				indirect.table_immediate = 0u;
+				    return use.user->GetOpcode() == ValueOpcode::ImageSampleRaw &&
-				indirect.table_record_bytes = table_immediate + plan.read_dwords * 4u;
+				        use.user->Arg(1).Resolve().TryInstruction() == paired_sampler->handle &&
-			} else {
+				        (m_program.memory_info[use.user->Flags<MemoryFlags>().index].image_sample_flags &
-				indirect.table_scalar = true;
+				         Decoder::ImageSampleFlagAdjust) == 0u;
+			    })) {
+				// This sampler already requires these exact keys. Prove the pairing before
+				// interning, so image-only handles keep their independent heap domain.
+				auto selected = indirect;
+				DescriptorSource selected_source;
+				if (MatchScalarMaterialKey(key, selected, selected_source, false, true) &&
+				    EquivalentSelectorDomain(m_program, selected,
+				        *m_sources[paired_sampler->source].indirect_descriptor)) {
+					indirect = std::move(selected);
+					material_source = std::move(selected_source);
+					shared_selector = true;
+				}
+			}
+			if (!shared_selector) {
+				indirect.workgroup_axis = UINT32_MAX;
+				plan.table_read = plan.reads[0];
+				if (plan.table_indexed) {
+					indirect.table_stride = 0u;
+					indirect.table_offset = 0u;
+					indirect.table_immediate = 0u;
+					indirect.table_record_bytes = table_immediate + plan.read_dwords * 4u;
+				} else {
+					indirect.table_scalar = true;
+				}
 			}
 			plan.retain_reads = true;
 		}
@@ -3415,6 +3440,7 @@ private:
 				}
 				const auto flags = inst.Flags<MemoryFlags>();
 				const auto& memory = m_program.memory_info[flags.index];
+				const IndirectDescriptorPlan* paired_sampler = nullptr;
 				const auto plan_handle = [&](Value value, uint32_t width, uint32_t base, bool sampler) {
 					auto* handle = value.Resolve().TryInstruction();
 					if (handle == nullptr || FindIndirectDescriptor(*handle) != nullptr) return;
@@ -3437,7 +3463,7 @@ private:
 					}
 					if (TryMakeNestedDescriptor(*handle, descriptor, plan) ||
 					    (sampler && TryMakeFiniteDescriptor(*handle, plan)) ||
-					    TryMakeTableDescriptor(*handle, descriptor, plan) ||
+					    TryMakeTableDescriptor(*handle, descriptor, plan, paired_sampler) ||
 					    (!sampler && (TryMakeFiniteDescriptor(*handle, plan) ||
 					                  TryMakeIndirectBuffer(*handle, descriptor, plan)))) {
 						for (const auto& previous: m_indirect_descriptors) {
@@ -3448,11 +3474,14 @@ private:
 						m_indirect_descriptors.push_back(std::move(plan));
 					}
 				};
+				if (inst.GetOpcode() == ValueOpcode::ImageSampleRaw &&
+				    (memory.image_sample_flags & Decoder::ImageSampleFlagAdjust) == 0u) {
+					plan_handle(inst.Arg(1), 4u, memory.sampler * 4u, true);
+					if (const auto* handle = inst.Arg(1).Resolve().TryInstruction())
+						paired_sampler = FindIndirectDescriptor(*handle);
+				}
 				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)
-					plan_handle(inst.Arg(1), 4u, memory.sampler * 4u, true);
 			}
 		}
 	}

tests/ResourceTrackingTests.cpp

@@ -2995,8 +2995,10 @@ void TestCompactScalarImageTable() {
   using F = Libs::Graphics::Prospero::BufferFormat;
   constexpr uint32_t stride = 0x330u, image_offset = 0x98u, sampler_offset = 0x88u;
   enum class Domain { Unbounded, SharedPrefix, SeparatePrefix, SeparateMask };
+  enum class Pairing { SameKey, DifferentKey, PairedFirst, UnpairedFirst, MixedUses };
   const auto make_plan = [&](bool compact, uint32_t invalid = 0u,
-                             Domain domain = Domain::Unbounded) {
+                             Domain domain = Domain::Unbounded, Pairing pairing = Pairing::SameKey,
+                             uint32_t member = 0x98u) {
     Fixture fixture;
     const auto list = invalid == 1u
         ? fixture.Address(fixture.UserData(0), fixture.UserData(1))
@@ -3033,8 +3035,18 @@ void TestCompactScalarImageTable() {
       if (invalid == 3u) words[7] = Value(1u);
       return fixture.Image(words, 0x59cu);
     };
-    const auto image = make_image(key, image_offset);
+    const auto image = make_image(key, member);
-    const auto record = fixture.Emit(ValueOpcode::IMul32, {key, Value(stride)});
+    auto sampler_key = key;
+    Value unpaired;
+    if (pairing == Pairing::DifferentKey || pairing == Pairing::PairedFirst ||
+        pairing == Pairing::UnpairedFirst) {
+      scalar.offset = 4u;
+      const auto other_key = fixture.Emit(ValueOpcode::ReadConstBuffer, {list, offset},
+                                           fixture.AddMemory(scalar, 0x578u));
+      if (pairing == Pairing::DifferentKey) sampler_key = other_key;
+      else unpaired = make_image(other_key, member);
+    } else if (pairing == Pairing::MixedUses) unpaired = image;
+    const auto record = fixture.Emit(ValueOpcode::IMul32, {sampler_key, Value(stride)});
     std::array<Value, 4> sampler_words;
     for (uint32_t word = 0; word < 4u; ++word) {
       scalar.offset = sampler_offset + word * 4u;
@@ -3046,8 +3058,14 @@ void TestCompactScalarImageTable() {
     sample.image_dimension = Decoder::ImageDimension::Dim2D;
     sample.image_r128 = compact;
     const auto sampler = fixture.Sampler(sampler_words, 0x59cu);
+    const auto read_unpaired = [&] {
+      fixture.Emit(ValueOpcode::ImageRead, {unpaired, fixture.ImageAddress(), Value(true)},
+                   fixture.AddMemory(sample, 0x580u));
+    };
+    if (pairing == Pairing::UnpairedFirst) read_unpaired();
     fixture.Emit(ValueOpcode::ImageSampleRaw,
         {image, sampler, fixture.ImageAddress(), Value(true)}, fixture.AddMemory(sample, 0x59cu));
+    if (pairing == Pairing::PairedFirst || pairing == Pairing::MixedUses) read_unpaired();
     if (domain != Domain::Unbounded) {
       // Distinct SSA expressions with equivalent runtime first values share a domain.
       const auto second_offset = fixture.Emit(ValueOpcode::ShiftLeftLogical32,
@@ -3077,12 +3095,34 @@ void TestCompactScalarImageTable() {
               source.indirect_descriptor->table_stride == stride &&
               source.indirect_descriptor->table_immediate == sampler_offset,
           "compact image sampler lost its complete bounded index-list domain");
+    if (domain == Domain::Unbounded) {
+      const auto &image_source = fixture.program.descriptor_sources[
+          fixture.program.info.images[image.Instruction()->Flags<uint32_t>()].source];
+      const bool paired = pairing != Pairing::DifferentKey && pairing != Pairing::MixedUses;
+      Check(image_source.indirect_descriptor &&
+                image_source.indirect_descriptor->selector.has_value() == paired,
+            "paired image selector proof ignored key identity or an unpaired handle use");
+      if (!unpaired.IsEmpty() && unpaired != image) {
+        const auto &other_source = fixture.program.descriptor_sources[
+            fixture.program.info.images[unpaired.Instruction()->Flags<uint32_t>()].source];
+        Check(&other_source != &image_source && other_source.indirect_descriptor &&
+                  !other_source.indirect_descriptor->selector &&
+                  other_source.indirect_descriptor->table_stride == stride &&
+                  other_source.indirect_descriptor->table_immediate == member,
+              "paired image narrowed an existing unpaired heap source");
+      }
+    }
     return ExtractResourcePlan(fixture.program);
   };
-  for (const bool compact : {false, true}) {
+  struct Case { bool compact; uint32_t member; Pairing pairing; };
-    auto plan = make_plan(compact);
+  for (const auto [compact, member, pairing] : {
+           Case{false, image_offset, Pairing::SameKey}, Case{true, image_offset, Pairing::SameKey},
+           Case{true, 0x16cu, Pairing::SameKey}, Case{true, 0x16cu, Pairing::DifferentKey},
+           Case{true, 0x16cu, Pairing::PairedFirst}, Case{true, 0x16cu, Pairing::UnpairedFirst},
+           Case{true, 0x16cu, Pairing::MixedUses}}) {
+    auto plan = make_plan(compact, 0u, Domain::Unbounded, pairing, member);
     const uint32_t width = compact ? 4u : 8u;
-    const auto bytes = 2u * stride + image_offset + width * 4u;
+    const auto bytes = 2u * stride + member + width * 4u;
     std::array<uint32_t, 8> data{0x1000u, 0u, 8u, 0u, 0x2000u, 0u, bytes, 0u};
     struct Memory : LinearTestMemory { uint32_t selector_attempts = 0; } memory;
     memory.words.resize((0x1000u + bytes) / 4u, 0xdeadbeefu);
@@ -3092,7 +3132,7 @@ void TestCompactScalarImageTable() {
       auto at = memory.words.begin() + (0x1000u + row * stride + sampler_offset) / 4u;
       std::fill_n(at, 4u, 0u);
       at[0] = row;
-      at += 4;
+      at = memory.words.begin() + (0x1000u + row * stride + member) / 4u;
       std::fill_n(at, width, 0u);
       at[0] = 0x100u + row;
       at[1] = static_cast<uint32_t>(F::k32_32_32_32Float) << 20u;
@@ -3100,6 +3140,9 @@ void TestCompactScalarImageTable() {
           (static_cast<uint32_t>(Libs::Graphics::Prospero::ImageType::kColor2D) << 28u);
       if (!compact) at[5] = 0x00700000u;
     }
+    const std::array poison{0x00fff000u, 0x09500000u, 0u, 0x900672b0u};
+    if (member == 0x16cu)
+      std::copy(poison.begin(), poison.end(), memory.words.begin() + (0x1000u + 0x17cu) / 4u);
     memory.fail_address = 0x2000u + bytes; // Last descriptor ends at the readable boundary.
     SrtRuntime runtime{.user_data = data, .userdata = &memory,
         .read_specialization_memory = [](void *userdata, uint64_t address,
@@ -3110,10 +3153,20 @@ void TestCompactScalarImageTable() {
         }};
     ResourceSnapshot snapshot;
     ResourceSpecialization specialization;
+    if (pairing != Pairing::SameKey) {
+      Check(!MaterializeResources(plan, runtime, snapshot, specialization) &&
+                std::ranges::any_of(snapshot.images, [&](const auto &image) {
+                  return std::equal(poison.begin(), poison.end(), image.dwords.begin());
+                }),
+            "unpaired image use silently inherited a narrower sampler domain");
+      continue;
+    }
     const auto materialize = [&] {
+      memory.selector_attempts = 0u;
       Check(MaterializeResources(plan, runtime, snapshot, specialization),
             "paired compact image/sampler table did not materialize");
-      Check(snapshot.images.size() == 4u && snapshot.samplers.size() == 3u,
+      Check(snapshot.images.size() == 3u && snapshot.samplers.size() == 3u &&
+                memory.selector_attempts == 1u,
             "paired table omitted an image, index-list value or OOB zero key");
       for (const auto &image : snapshot.images) {
         Check(image.dword_count == 8u, "compact image escaped normalized eight-DWORD ABI");
@@ -3124,16 +3177,19 @@ void TestCompactScalarImageTable() {
     };
     materialize();
     const auto original = specialization;
-    memory.words[(0x1000u + stride + image_offset) / 4u] = 0x201u;
+    memory.words[(0x1000u + stride + member) / 4u] = 0x201u;
     materialize();
     Check(specialization == original && std::ranges::any_of(snapshot.images,
               [](const auto &image) { return image.dwords[0] == 0x201u; }),
           "compact descriptor payload refresh changed specialization or stayed stale");
     data[2] = 4u;
-    Check(MaterializeResources(plan, runtime, snapshot, specialization) && snapshot.samplers.size() == 2u,
+    Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+              snapshot.images.size() == 2u && snapshot.samplers.size() == 2u,
           "index-list count refresh retained a stale sampler candidate");
     data[2] = 0u;
-    Check(MaterializeResources(plan, runtime, snapshot, specialization) && snapshot.samplers.size() == 1u &&
+    Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+              snapshot.images.size() == 1u && snapshot.images[0].dwords[0] == 0x100u &&
+              snapshot.samplers.size() == 1u &&
               snapshot.samplers[0].dwords[0] == 0u,
           "empty index list lost the scalar OOB zero key");
     data[2] = 8u;

tests/ShaderRecompilerComputeTests.cpp

@@ -38699,12 +38699,14 @@ void CheckNestedPointerImageTable(VulkanHarness &vulkan) {
   constexpr u32 keys_address = 128, list_address = 160, table = 256, objects = 4096;
   constexpr std::array keys{5u, 0u, 3u, 1u, 4u, 2u};
   constexpr std::array list{2u, 1u, 4u, 0u, 5u, 3u};
-  for (const bool compact : {false, true})
+  for (const u32 member : {0u, 0x98u, 0x16cu})
   for (const u32 wave : {32u, 64u}) {
+    const bool compact = member != 0u;
     const u32 stride = compact ? 0x330u : 592u;
-    const u32 table_bytes = (keys.size() - 1u) * stride + 0x98u + 16u;
+    const u32 table_bytes = (keys.size() - 1u) * stride + member + 16u;
     TestCase test;
-    test.name = compact ? "CompactRecordImageSamplerTable" : "NestedPointerImageSamplerTable";
+    test.name = member == 0x16cu ? "CompactPairedImageSamplerTable"
+              : compact ? "CompactRecordImageSamplerTable" : "NestedPointerImageSamplerTable";
     test.initial.resize((compact ? table + table_bytes : objects + keys.size() * 0x430u) / 4u,
                         compact ? 0xdeadbeefu : 0u);
     std::fill_n(test.initial.begin(), keys.size(), 0xdeadbeefu);
@@ -38719,15 +38721,16 @@ void CheckNestedPointerImageTable(VulkanHarness &vulkan) {
       image[3] = DstSel(4, 0, 0, 1) |
           (static_cast<u32>(Prospero::ImageType::kColor2D) << 28u);
       std::copy_n(image.begin(), compact ? 4u : 8u,
-                  test.initial.begin() + (object + (compact ? 0x98u : 0x3d0u)) / 4u);
+                  test.initial.begin() + (object + (compact ? member : 0x3d0u)) / 4u);
       const std::array sampler{0u, 0x00fff000u, (row & 1u) ? 0x0a500000u : 0u, 0u};
       std::copy(sampler.begin(), sampler.end(),
                 test.initial.begin() + (object + (compact ? 0x88u : 0x3f8u)) / 4u);
     }
     if (compact) {
-      // PC59c can only start at residue8 mod16; this captured tuple starts at residue12.
+      // Member 0x16c admits this tuple's residue; only the paired selector list excludes it.
       const std::array unreachable{0x00fff000u, 0x09500000u, 0u, 0x900672b0u};
-      std::copy(unreachable.begin(), unreachable.end(), test.initial.begin() + (table + 0xbcu) / 4u);
+      std::copy(unreachable.begin(), unreachable.end(),
+                test.initial.begin() + (table + (member == 0x16cu ? 0x17cu : 0xbcu)) / 4u);
     }
     auto records = MakeStructuredStorageBufferData(compact ? 0u : stride,
                                                     compact ? table_bytes : keys.size());
@@ -38760,8 +38763,11 @@ void CheckNestedPointerImageTable(VulkanHarness &vulkan) {
     if (compact) {
       code.insert(code.end(), {EncodeSop2(0x1e, 23, 22, InlineU32(2)),
                               EncodeSmem0(0x08, 22, 18), EncodeSmem1(0, 23),
-                              EncodeSop2(0x26, 23, 22, 255), stride,
+                              EncodeSop2(0x26, 23, 22, 255), stride});
-                              EncodeSmem0(0x0b, 8, 16), EncodeSmem1(0x88u, 23)});
+      if (member == 0x98u)
+        code.insert(code.end(), {EncodeSmem0(0x0b, 8, 16), EncodeSmem1(0x88u, 23)});
+      else code.insert(code.end(), {EncodeSmem0(0x0a, 8, 16), EncodeSmem1(0x88u, 23),
+                                   EncodeSmem0(0x0a, 12, 16), EncodeSmem1(member, 23)});
       for (u32 reg = 16u; reg < 20u; ++reg) AppendSMovLiteral(&code, reg, 0xdeadbeefu);
     } else code.insert(code.end(), {EncodeSop2(0x26, 23, 22, 255), stride,
                             EncodeSmem0(0x09, 20, 16), EncodeSmem1(0, 23),
@@ -38778,18 +38784,23 @@ void CheckNestedPointerImageTable(VulkanHarness &vulkan) {
         O::IMAGE_SAMPLE, O::BUFFER_STORE_DWORD, O::S_ENDPGM};
     if (compact) {
       test.opcodes.insert(test.opcodes.end(), {O::S_BUFFER_LOAD_DWORD,
-          O::S_BUFFER_LOAD_DWORDX8, O::S_MOV_B32});
+          member == 0x98u ? O::S_BUFFER_LOAD_DWORDX8 : O::S_BUFFER_LOAD_DWORDX4, O::S_MOV_B32});
       test.decoded_counts = {{"r128=1", 1}};
     } else test.opcodes.insert(test.opcodes.end(),
         {O::S_BUFFER_LOAD_DWORDX2, O::S_LOAD_DWORDX8, O::S_LOAD_DWORDX4});
     const auto compiled = CompileCase(test, vulkan.SubgroupSize());
-    if (compact) Require(test.name, "physical R128 descriptor boundary",
+    if (compact) {
-        compiled.resources.images.size() == keys.size() + 1u &&
+      std::vector<u32> candidates;
+      for (const auto &image : compiled.resources.images) candidates.push_back(image.dwords[0]);
+      std::ranges::sort(candidates);
+      Require(test.name, "paired R128 descriptor boundary",
+        candidates == std::vector<u32>{10u, 11u, 12u, 13u, 14u, 15u} &&
             std::ranges::all_of(compiled.resources.images, [](const auto &image) {
               return image.dword_count == 8u &&
                   std::ranges::all_of(std::span(image.dwords).subspan(4u),
                                       [](u32 word) { return word == 0u; });
-            }), "compact descriptor copied adjacent poison or omitted the trailing 16-byte image");
+            }), "paired selector admitted poison or omitted the trailing 16-byte image");
+    }
     for (const auto &value : compiled.resources.buffers) {
       ShaderBufferResource descriptor{};
       std::copy_n(value.dwords.begin(), 4, descriptor.fields);