shader: preserve affine image heap candidate domains

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

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

@@ -410,8 +410,10 @@ 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 bool image_heap = dword_count == 8u && !indirect.pointer_offset;
 			const uint64_t step = indirect.table_maximum ? indirect.table_stride
-			    : affine ? uint64_t {1} << std::countr_zero(indirect.table_stride) : 4u;
+			    : affine ? std::max(uint64_t {1} << std::countr_zero(indirect.table_stride), uint64_t {4})
+			             : 4u;
 			const uint64_t first = ((indirect.table_maximum ? indirect.table_offset
 			    : indirect.table_offset % step) & ~uint64_t {3}) +
 			                       (indirect.table_immediate & ~3u);
@@ -420,10 +422,10 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 			const auto count = table_size >= first + 4u
 			    ? std::min((table_size - first - 4u) / step + 1u, domain)
 			    : 0u;
-			if (affine && count > MaxIndirectDescriptorProbes) return false;
+			if (!image_heap && count > MaxIndirectDescriptorProbes) return false;
 			// Affine SMEM offsets wrap in U32 before adding the aligned immediate.
-			// Read only their congruence class; unrestricted image tables retain a bulk read.
+			// Image heaps retain one bulk snapshot and their absolute byte-offset mapping.
-			if (!affine) {
+			if (image_heap) {
 				keys.resize(table_size / 4u);
 				if (!ReadScalarTable(table_base, table_size, 0u, runtime, keys)) return false;
 			}
@@ -432,7 +434,7 @@ 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};
-				const auto table_offset = affine || indirect.pointer_offset
+				const auto table_offset = affine
 				    ? first + word * step : uint64_t {word} * 4u;
 				if (indirect.pointer_offset) {
 					std::array<uint32_t, 2> address {};
@@ -440,11 +442,13 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 					    !ReadScalarTable((uint64_t {address[1]} << 32u) | address[0], UINT64_MAX,
 					                     *indirect.pointer_offset, runtime,
 					                     std::span(candidate.dwords).first(read_dwords))) return false;
-				} else if (affine) {
+				} else if (!image_heap) {
 					if (!ReadScalarTable(table_base, table_size, table_offset, runtime,
 					                     std::span(candidate.dwords).first(read_dwords))) return false;
 				} else {
-					std::copy_n(keys.begin() + word, std::min<size_t>(read_dwords, keys.size() - word),
+					const auto table_word = table_offset / 4u;
+					std::copy_n(keys.begin() + table_word,
+					            std::min<size_t>(read_dwords, keys.size() - table_word),
 					            candidate.dwords.begin());
 				}
 				if (!normalize(candidate))
@@ -453,8 +457,9 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 				uint32_t ordinal;
 				if (!intern_candidate(candidate, ordinal)) return false;
 				if (ordinal == 0u) continue;
-				snapshot.flattened_srt.resize(mapping_offset + 2u + word, 0u);
+				const auto mapping_word = image_heap ? table_offset / 4u : word;
-				snapshot.flattened_srt[mapping_offset + 1u + word] = ordinal;
+				snapshot.flattened_srt.resize(mapping_offset + 2u + mapping_word, 0u);
+				snapshot.flattened_srt[mapping_offset + 1u + mapping_word] = ordinal;
 			}
 			if (descriptors.size() == children_begin) {
 				snapshot.flattened_srt.resize(mapping_offset);
@@ -846,7 +851,7 @@ static bool BuildResourceSpecialization(const ResourcePlan& program, ResourceSna
 			if (i >= program.info.images.size() && source != nullptr &&
 			    source->indirect_descriptor &&
 			    IsBoundedDescriptorTable(program, *source->indirect_descriptor) &&
-			    source->indirect_descriptor->table_stride == 0u) {
+			    !source->indirect_descriptor->pointer_offset) {
 				const auto& indirect = *source->indirect_descriptor;
 				const size_t mapping = specialization.images[base_index].indirect_mapping_offset;
 				const auto& flat = snapshot.flattened_srt;

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

@@ -972,6 +972,7 @@ private:
 	uint32_t InternSource(const DescriptorSource& descriptor) {
 		for (uint32_t candidate = 0; candidate < m_sources.size(); candidate++) {
 			auto& current = m_sources[candidate];
+			bool shared_image_heap = false;
 			if (current.dword_count != descriptor.dword_count ||
 			    current.indirect_descriptor.has_value() != descriptor.indirect_descriptor.has_value()) {
 				continue;
@@ -979,11 +980,16 @@ private:
 			if (current.indirect_descriptor.has_value()) {
 				const auto& a = *current.indirect_descriptor;
 				const auto& b = *descriptor.indirect_descriptor;
+				// An indexed image heap already covers every DWORD start, so its domain
+				// subsumes an affine scalar use of the same table and shares its snapshot.
+				shared_image_heap = descriptor.dword_count == 8u &&
+				    (a.table_record_bytes != 0u || b.table_record_bytes != 0u);
 				if (a.selector != b.selector || a.table_maximum != b.table_maximum ||
 				    a.pointer_offset != b.pointer_offset ||
 				    a.table_source != b.table_source ||
-				    a.table_offset != b.table_offset || a.table_immediate != b.table_immediate ||
+				    (!shared_image_heap && (a.table_offset != b.table_offset ||
-				    a.table_stride != b.table_stride || a.indexed != b.indexed ||
+				                           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 ||
 				    a.signed_count != b.signed_count ||
 				    !EquivalentValue(m_program, a.key_count, b.key_count) ||
@@ -1000,6 +1006,11 @@ private:
 			}
 			if (same) {
 				if (current.indirect_descriptor) {
+					if (shared_image_heap) {
+						current.indirect_descriptor->table_stride = 0u;
+						current.indirect_descriptor->table_offset = 0u;
+						current.indirect_descriptor->table_immediate = 0u;
+					}
 					current.indirect_descriptor->table_scalar |= descriptor.indirect_descriptor->table_scalar;
 					current.indirect_descriptor->table_record_bytes = std::max(
 					    current.indirect_descriptor->table_record_bytes,
@@ -2920,16 +2931,19 @@ private:
 			}
 			if ((table_offset & 3u) != 0u) return false;
 		} else {
-			// A bounded V# supplies the complete image table. Leave every GPU selector
+			// Keep the GPU selector and descriptor reads, and preserve the proven scalar
-			// and descriptor read in the shader; the host only translates table bytes.
+			// offset domain within the bounded image table.
 			if (sampler || plan.split_offsets) return false;
-			indirect.table_stride = 0u;
-			indirect.table_offset = 0u;
-			indirect.table_immediate = 0u;
 			indirect.workgroup_axis = UINT32_MAX;
 			plan.table_read = plan.reads[0];
-			if (plan.table_indexed) indirect.table_record_bytes = table_immediate + plan.read_dwords * 4u;
+			if (plan.table_indexed) {
-			else indirect.table_scalar = true;
+				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;
 		}
 

tests/ResourceTrackingTests.cpp

@@ -336,6 +336,10 @@ void TestInvariantIndirectImageMaterialization(bool nested) {
   };
   put(16u, 0u);
   put(64u, 1u);
+  // Captured sampler/T# overlap: valid-looking format149 at an unreachable DWORD start.
+  const std::array<uint32_t, 8> unreachable{
+      0x00fff000u, 0x09500000u, 0u, 0x900672b0u, 0u, 0u, 0u, 0u};
+  std::copy(unreachable.begin(), unreachable.end(), memory.words.begin() + 0x1064u / 4u);
   SrtRuntime runtime{.user_data = data, .userdata = &memory,
                      .read_specialization_memory = ReadLinearTestMemory};
   ResourceSnapshot snapshot;
@@ -352,6 +356,22 @@ void TestInvariantIndirectImageMaterialization(bool nested) {
             snapshot.flattened_srt[mapping + 17u] == 2u &&
             snapshot.specialization_reads == std::vector<std::pair<uint64_t, uint64_t>>{{0x2000u, 144u}},
         "dense image mapping lost its null slot, trimmed length or coherent read dependency");
+  for (const uint32_t stride : {1u, 2u}) {
+    auto subdword = MakeIndirectImageFixture(false, stride, 0u, 0u, nested);
+    subdword->PlanAndTrack();
+    const auto subdword_plan = ExtractResourcePlan(subdword->program);
+    auto subdword_memory = memory;
+    auto subdword_runtime = runtime;
+    subdword_runtime.userdata = &subdword_memory;
+    ResourceSnapshot subdword_snapshot;
+    ResourceSpecialization subdword_specialization;
+    Check(!MaterializeResources(subdword_plan, subdword_runtime, subdword_snapshot,
+                               subdword_specialization) &&
+              std::ranges::any_of(subdword_snapshot.images, [&](const auto &image) {
+                return image.dwords == unreachable;
+              }),
+          "sub-DWORD affine stride skipped a reachable unsupported descriptor");
+  }
   const auto original_specialization = specialization;
   std::fill_n(memory.words.begin() + 0x1040u / 4u, 8u, 0u);
   put(112u, 1u);
@@ -577,7 +597,8 @@ void TestExactUnfilteredImageCandidates() {
 
 void TestWaterfallImageTable() {
   namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
-  const auto make = [](bool equality, bool scalar) {
+  enum class Sources { ScalarFirst, IndexedFirst, DistinctScalars };
+  const auto make = [](bool equality, bool scalar, Sources sources = Sources::ScalarFirst) {
     auto fixture = std::make_unique<Fixture>();
     auto *entry = fixture->block;
     auto *body = fixture->AddBlock();
@@ -628,24 +649,33 @@ void TestWaterfallImageTable() {
       fixture->Emit(ValueOpcode::ReferenceU32,
           {fixture->Emit(ValueOpcode::CompositeExtractU32x4, {value, Value(0u)}, 0, block)}, 0, block);
     };
+    Value delayed_scalar;
     if (scalar) {
       const auto key = fixture->Emit(ValueOpcode::ReadFirstLane, {lane_value, active});
-      const auto offset = fixture->Emit(ValueOpcode::IAdd32,
+      for (const uint32_t member : {16u, 20u}) {
-          {fixture->Emit(ValueOpcode::IMul32, {key, Value(48u)}), Value(16u)});
+        if (member == 20u && sources != Sources::DistinctScalars) break;
-      std::array<Value, 8> scalar_words;
+        const auto offset = fixture->Emit(ValueOpcode::IAdd32,
-      for (uint32_t word = 0; word < 8u; ++word) {
+            {fixture->Emit(ValueOpcode::IMul32, {key, Value(48u)}), Value(member)});
-        MemoryInfo memory;
+        std::array<Value, 8> scalar_words;
-        memory.kind = ResourceKind::ScalarBuffer;
+        for (uint32_t word = 0; word < 8u; ++word) {
-        memory.offset = word * 4u;
+          MemoryInfo memory;
-        scalar_words[word] = fixture->Emit(ValueOpcode::ReadConstBuffer, {table, offset},
+          memory.kind = ResourceKind::ScalarBuffer;
-                                           fixture->AddMemory(memory, 0x7bcu));
+          memory.offset = word * 4u;
+          scalar_words[word] = fixture->Emit(ValueOpcode::ReadConstBuffer, {table, offset},
+                                             fixture->AddMemory(memory, 0x7bcu));
+        }
+        const auto scalar_image = fixture->Image(scalar_words);
+        if (sources == Sources::IndexedFirst) delayed_scalar = scalar_image;
+        else sample(scalar_image, entry);
       }
-      sample(fixture->Image(scalar_words), entry);
     }
-    const auto image = fixture->Emit(ValueOpcode::GetImageResource,
+    if (sources != Sources::DistinctScalars) {
-        {image_words[0], image_words[1], image_words[2], image_words[3],
+      const auto image = fixture->Emit(ValueOpcode::GetImageResource,
-         image_words[4], image_words[5], image_words[6], image_words[7]}, 0, body);
+          {image_words[0], image_words[1], image_words[2], image_words[3],
-    sample(image, body);
+           image_words[4], image_words[5], image_words[6], image_words[7]}, 0, body);
+      sample(image, body);
+    }
+    if (!delayed_scalar.IsEmpty()) sample(delayed_scalar, body);
     return fixture;
   };
   auto fixture = make(true, true);
@@ -666,6 +696,34 @@ void TestWaterfallImageTable() {
   Check(MaterializeResources(plan, runtime, snapshot, specialization) && memory.reads == 1u &&
             snapshot.images.size() == 2u,
         "shared scalar/waterfall table repeated its coherent bulk read");
+  auto reversed = make(true, true, Sources::IndexedFirst);
+  reversed->PlanAndTrack();
+  const auto reversed_plan = ExtractResourcePlan(reversed->program);
+  Check(reversed_plan.info.images.size() == 1u,
+        "indexed-first image table did not share its scalar use");
+  const auto &shared = *reversed_plan.descriptor_sources[
+      reversed_plan.info.images[0].source].indirect_descriptor;
+  Check(shared.table_scalar && shared.table_record_bytes == 48u &&
+            shared.table_stride == 0u && shared.table_offset == 0u && shared.table_immediate == 0u,
+        "indexed-first union lost scalar access or the unrestricted byte domain");
+  const auto previous_reads = memory.reads;
+  ResourceSnapshot reversed_snapshot;
+  ResourceSpecialization reversed_specialization;
+  Check(MaterializeResources(reversed_plan, runtime, reversed_snapshot, reversed_specialization) &&
+            memory.reads == previous_reads + 1u && reversed_snapshot.images == snapshot.images,
+        "indexed-first table repeated its bulk read or changed descriptor candidates");
+  auto distinct = make(true, true, Sources::DistinctScalars);
+  distinct->PlanAndTrack();
+  const auto distinct_plan = ExtractResourcePlan(distinct->program);
+  Check(distinct_plan.info.images.size() == 2u,
+        "distinct scalar image domains were merged into an unrestricted heap");
+  for (uint32_t image = 0; image < 2u; ++image) {
+    const auto &restricted = *distinct_plan.descriptor_sources[
+        distinct_plan.info.images[image].source].indirect_descriptor;
+    Check(restricted.table_record_bytes == 0u && restricted.table_stride == 48u &&
+              restricted.table_offset == 16u + image * 4u && restricted.table_immediate == 0u,
+          "scalar image union discarded a proven affine residue");
+  }
   const auto handle = std::ranges::find_if(*fixture->program.blocks[1], [](const Inst &inst) {
     return inst.GetOpcode() == ValueOpcode::GetImageResource;
   });

tests/ShaderRecompilerComputeTests.cpp

@@ -38724,6 +38724,11 @@ void CheckNestedPointerImageTable(VulkanHarness &vulkan) {
       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.
+      const std::array unreachable{0x00fff000u, 0x09500000u, 0u, 0x900672b0u};
+      std::copy(unreachable.begin(), unreachable.end(), test.initial.begin() + (table + 0xbcu) / 4u);
+    }
     auto records = MakeStructuredStorageBufferData(compact ? 0u : stride,
                                                     compact ? table_bytes : keys.size());
     records[0] = table;