shader: resolve GPU scalar image selectors within buffer bounds

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

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

@@ -394,7 +394,22 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 				keys.push_back(key);
 				mask &= mask - 1u;
 			}
-		} else return false;
+		} else {
+			ShaderBufferResource material;
+			if (!DecodeBufferDescriptor(material_value, material) || selector->stride < 4u ||
+			    !std::has_single_bit(selector->stride) || (selector->offset & 3u) != 0u)
+				return false;
+			// Empty selector bounds denote the complete U32 congruence class. SMEM
+			// adds its aligned immediate after the dynamic offset, without U32 wrap.
+			const auto domain_count = (uint64_t {1} << 32u) / selector->stride;
+			const auto available = material.GetSize() > selector->offset
+			    ? material.GetSize() - selector->offset : 0u;
+			const auto count = std::min(domain_count,
+			    available < 4u ? 0u : (available - 4u) / selector->stride + 1u);
+			if (count > MaxIndirectDescriptorProbes ||
+			    !read_keys(material, selector->offset, selector->stride, count)) return false;
+			if (count < domain_count) keys.push_back(0u); // An out-of-bounds scalar read returns zero.
+		}
 		if (selector != nullptr && selector->byte_ranges.empty()) {
 			std::ranges::sort(keys);
 			keys.erase(std::unique(keys.begin(), keys.end()), keys.end());

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

@@ -8,6 +8,7 @@
 #include <algorithm>
 #include <fmt/format.h>
 #include <map>
+#include <numeric>
 #include <optional>
 #include <span>
 #include <tuple>
@@ -1152,6 +1153,7 @@ private:
 
 	bool HasActiveLane(Value mask, const Block* use) const {
 		mask = mask.Resolve();
+		if (mask.IsImmediate()) return mask.GetType() == Type::U1 && mask.U1();
 		const auto incoming_is_nonempty = [&](const Block* from, const Block* to,
 		                                     Value incoming, const Block* header) {
 			const auto accepts = [&](const EdgePredicate& edge) {
@@ -1534,20 +1536,37 @@ private:
 		uint32_t memory_index = 0;
 		const auto* memory = ScalarReadMemory(*read, memory_index);
 		Value index;
+		Value count;
 		uint32_t offset = 0, stride = 0;
 		if (memory == nullptr || memory->kind != ResourceKind::ScalarBuffer ||
-		    memory->offset > INT32_MAX || !MemoryIndexBelongsTo(memory_index, *read) ||
+		    memory->offset > INT32_MAX || !MemoryIndexBelongsTo(memory_index, *read)) return false;
-		    !MatchTableOffset(read->Arg(1), index, offset, stride) ||
+		if (MatchTableOffset(read->Arg(1), index, offset, stride) &&
-		    ((offset | stride | memory->offset) & 3u) != 0u ||
+		    ((offset | stride | memory->offset) & 3u) == 0u &&
-		    uint64_t {offset} + memory->offset > UINT32_MAX) return false;
+		    uint64_t {offset} + memory->offset <= UINT32_MAX) {
-		const auto* bound = BoundedLoop(index, read->Parent());
+			const auto* bound = BoundedLoop(index, read->Parent());
-		if (bound == nullptr || bound->GetOpcode() != ValueOpcode::ULessThan32 ||
+			if (bound != nullptr && bound->GetOpcode() == ValueOpcode::ULessThan32) {
-		    !MakeRuntimeTableSource(*read, material_source)) return false;
+				count = bound->Arg(1);
+			}
+		}
+		if (count.IsEmpty()) {
+			const auto* lane = read->Arg(1).Resolve().TryInstruction();
+			if (lane == nullptr || lane->GetOpcode() != ValueOpcode::ReadFirstLane ||
+			    !HasActiveLane(lane->Arg(1), lane->Parent())) return false;
+			AffineOffset affine;
+			if (!MatchAffineOffset(lane->Arg(0), lane->Arg(1), affine) ||
+			    affine.index.IsEmpty() || affine.stride == 0u) return false;
+			// U32 multiplication can wrap. Its complete byte-offset domain is one
+			// congruence class, not the source V#'s structured record stride.
+			const auto divisor = std::gcd(affine.stride, uint64_t {1} << 32u);
+			stride = static_cast<uint32_t>(std::max<uint64_t>(divisor, 4u));
+			offset = static_cast<uint32_t>(affine.offset % divisor) & ~3u;
+		}
+		if (!MakeRuntimeTableSource(*read, material_source)) return false;
 		indirect.selector.emplace(DescriptorSource::IndirectDescriptor::SelectorRead {
 		    .source = InternSource(material_source), .stride = stride,
-		    .offset = offset + memory->offset});
+		    .offset = offset + (memory->offset & ~3u)});
-		indirect.selector_first = Value(0u);
+		indirect.key_count = count;
-		indirect.key_count = bound->Arg(1);
+		if (!count.IsEmpty()) indirect.selector_first = Value(0u);
 		return true;
 	}
 
@@ -1768,6 +1787,11 @@ private:
 			// The descriptor also supplies dimensions to shader arithmetic. Keep its reads;
 			// only the image handle is projected onto the bounded workgroup key.
 			plan.retain_reads = true;
+		} else if (!plan.table_indexed &&
+		           MatchScalarMaterialKey(key, indirect, material_source)) {
+			// Keep live descriptor payloads, but project descriptor-only SMEM onto
+			// the GPU key using the same plan as other indirect image reads.
+			plan.retain_reads |= table_source.dword_count == 4u;
 		} else if (table_source.dword_count == 2u) {
 			if (table_stride != 32u) return false;
 			const auto* selector = key.Resolve().TryInstruction();
@@ -1786,10 +1810,6 @@ private:
 			}
 			if ((table_offset & 3u) != 0u ||
 			    (bitscan && table_offset > UINT32_MAX - (32u * 32u - 1u))) return false;
-		} else if (!plan.table_indexed &&
-		           MatchScalarMaterialKey(key, indirect, material_source)) {
-			// The host enumerates the guarded scalar loop; its key and T# payload stay on the GPU.
-			plan.retain_reads = true;
 		} else {
 			// A bounded V# supplies the complete image table. Leave every GPU selector
 			// and descriptor read in the shader; the host only translates table bytes.

tests/ResourceTrackingTests.cpp

@@ -1620,6 +1620,126 @@ void TestUniformScalarBufferImage() {
         "unavailable nested image table was accepted");
 }
 
+void TestScalarBufferSelectorDomain() {
+  const auto make_plan = [](bool active = true, bool sparse = false) {
+    Fixture fixture;
+    const auto material = fixture.Buffer({fixture.UserData(0), fixture.UserData(1),
+                                          fixture.UserData(2), fixture.UserData(3)});
+    MemoryInfo lds;
+    lds.kind = ResourceKind::Lds;
+    const auto index = fixture.Emit(ValueOpcode::LoadSharedU32,
+        {Value(8u), Value(true)}, fixture.AddMemory(lds, 0x10));
+    auto offset = fixture.Emit(ValueOpcode::ShiftLeftLogical32,
+        {index, Value(sparse ? 31u : 5u)});
+    if (!sparse)
+      offset = fixture.Emit(ValueOpcode::IAdd32,
+          {fixture.Emit(ValueOpcode::ShiftLeftLogical32, {offset, Value(2u)}), offset});
+    offset = fixture.Emit(ValueOpcode::IAdd32, {offset, Value(3u)});
+    offset = fixture.Emit(ValueOpcode::ReadFirstLane, {offset, Value(active)});
+    MemoryInfo scalar;
+    scalar.kind = ResourceKind::ScalarBuffer;
+    scalar.offset = 147u;
+    const auto key = fixture.Emit(ValueOpcode::ReadConstBuffer,
+        {material, offset}, fixture.AddMemory(scalar, 0x20));
+    const auto table = fixture.Address(fixture.UserData(4), fixture.UserData(5));
+    const auto record = fixture.Emit(ValueOpcode::ShiftLeftLogical32, {key, Value(5u)});
+    std::array<Value, 8> words;
+    for (uint32_t word = 0; word < words.size(); ++word) {
+      scalar.kind = ResourceKind::ScalarAddress;
+      scalar.offset = word * 4u;
+      words[word] = fixture.Emit(ValueOpcode::LoadAddressU32,
+          {table, record, Value(0u), Value(true)}, fixture.AddMemory(scalar, 0x30));
+    }
+    MemoryInfo sample;
+    sample.kind = ResourceKind::Image;
+    sample.image_dimension = Decoder::ImageDimension::Dim2D;
+    fixture.Emit(ValueOpcode::ImageSampleRaw,
+        {fixture.Image(words), fixture.Sampler({Value(0u), Value(0u), Value(0u), Value(0u)}),
+         fixture.ImageAddress(), Value(true)}, fixture.AddMemory(sample, 0x40));
+    fixture.PlanAndTrack();
+    const auto &indirect = fixture.program.descriptor_sources[
+        fixture.program.info.images[0].source].indirect_descriptor;
+    Check(indirect && indirect->selector && indirect->key_count.IsEmpty() &&
+              indirect->selector_first.IsEmpty() && indirect->selector->offset == 144u &&
+              indirect->selector->stride == (sparse ? 0x80000000u : 32u),
+          "scalar selector assumed an unwrapped 160-byte record or combined unaligned offsets");
+    Check(key.ResolveInstruction()->Parent() != nullptr &&
+              std::ranges::count_if(fixture.program.memory_info, [](const MemoryInfo &memory) {
+                return memory.kind == ResourceKind::ScalarAddress && memory.planning_only;
+              }) == 8u,
+          "scalar selector lost its GPU key or retained descriptor-only address loads");
+    return ExtractResourcePlan(fixture.program);
+  };
+  auto plan = make_plan();
+  CheckFatal([&] { make_plan(false); }, "not a valid runtime value",
+             "scalar selector accepted an empty EXEC lane read");
+  LinearTestMemory memory;
+  memory.words[144u / 4u] = 1u;
+  memory.words[160u / 4u] = 99u; // This field is outside the affine offset domain.
+  memory.words[176u / 4u] = 2u;  // Reachable when the native 160*index multiplication wraps.
+  for (uint32_t key = 0; key < 3u; ++key) {
+    const auto word = (0x1000u + key * 32u) / 4u;
+    memory.words[word] = 0x100u + key;
+    memory.words[word + 1u] = static_cast<uint32_t>(
+        Libs::Graphics::Prospero::BufferFormat::k32_32_32_32Float) << 20u;
+    memory.words[word + 3u] = Libs::Graphics::DstSel(4, 5, 6, 7) |
+        (static_cast<uint32_t>(Libs::Graphics::Prospero::ImageType::kColor2D) << 28u);
+  }
+  std::array<uint32_t, 6> user_data {0x1000u, 0u, 180u, 0u, 0x2000u, 0u};
+  const SrtRuntime runtime {.user_data = user_data, .read_memory = ReadLinearTestMemory,
+      .userdata = &memory, .read_specialization_memory = ReadLinearTestMemory};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  memory.watched_address = 0x1000u + 144u;
+  memory.fail_address = 0x1000u + 160u;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+            snapshot.images.size() == 3u && memory.watched_reads == 1u &&
+            memory.watched_dwords == 1u && snapshot.images[0].dwords[0] == 0x100u &&
+            snapshot.images[2].dwords[0] == 0x102u,
+        "scalar selector read outside its wrapped domain or omitted the OOB zero key");
+  for (const auto size : {0u, 144u, 147u, 148u}) {
+    user_data[2] = size;
+    memory.watched_reads = 0u;
+    Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+              snapshot.images.size() == (size == 148u ? 2u : 1u) &&
+              memory.watched_reads == (size == 148u ? 1u : 0u) &&
+              snapshot.images[0].dwords[0] == 0x100u,
+          "scalar selector read an incomplete DWORD or lost a fully OOB zero key");
+  }
+  user_data[2] = 180u;
+  memory.words[176u / 4u] = 1u;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+            snapshot.images.size() == 2u,
+        "scalar selector reused stale keys after material memory changed");
+
+  plan = make_plan(true, true);
+  struct SparseMemory { std::vector<uint64_t> offsets; } sparse;
+  const auto read_sparse = +[](void *userdata, uint64_t address, std::span<uint32_t> words) {
+    auto &memory = *static_cast<SparseMemory *>(userdata);
+    std::ranges::fill(words, 0u);
+    if (words.size() == 1u) {
+      memory.offsets.push_back(address - 0x1000u);
+      words[0] = address == 0x1090u ? 1u : 2u;
+      return address == 0x1090u || address == 0x80001090u;
+    }
+    if (words.size() != 8u || (address != 0x2020u && address != 0x2040u)) return false;
+    words[0] = static_cast<uint32_t>(address);
+    words[1] = static_cast<uint32_t>(
+        Libs::Graphics::Prospero::BufferFormat::k32_32_32_32Float) << 20u;
+    words[3] = Libs::Graphics::DstSel(4, 5, 6, 7) |
+        (static_cast<uint32_t>(Libs::Graphics::Prospero::ImageType::kColor2D) << 28u);
+    return true;
+  };
+  user_data[1] = 2u << 16u;
+  user_data[2] = UINT32_MAX;
+  Check(MaterializeResources(plan,
+            {.user_data = user_data, .read_memory = read_sparse, .userdata = &sparse,
+             .read_specialization_memory = read_sparse}, snapshot, specialization) &&
+            sparse.offsets == std::vector<uint64_t>{144u, 0x80000090u} &&
+            snapshot.images.size() == 2u,
+        "large V# bounds extended the native U32 domain or invented an OOB zero key");
+}
+
 void TestComputeBufferFill() {
   struct Options {
     bool scalar = false;
@@ -3838,6 +3958,7 @@ int main() {
     Run("uniformized material image keys", TestUniformizedMaterialImageKeys);
     Run("image descriptor fields", TestImageDescriptorFields);
     Run("draw-uniform scalar image", TestUniformScalarBufferImage);
+    Run("scalar buffer selector domain", TestScalarBufferSelectorDomain);
     Run("SRT runtime", TestSrtFlatteningAndRuntimeMemoization);
     Run("dynamic SRT", TestDynamicSrtReadRemainsExplicit);
     Run("writable descriptor phi", TestWritableDescriptorPhi);