shader: bound bit-scan record keys and drop the key a guard rules out

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

@@ -157,7 +157,7 @@ KYTY_IO_FIELDS(StageOutput, a, b, c, d)
 KYTY_IO_FIELDS(ShaderInfo, a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v, w, x, y, z,
                A, B, C, D, E, F, G, H, I, J, K)
 KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor::SelectorRead::CountedRows, a, b)
-KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor::SelectorRead, a, b, c, d, e, f, g, h, i, j)
+KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor::SelectorRead, a, b, c, d, e, f, g, h, i, j, k)
 KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor::PointerTable, a, b, c, d)
 KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor, a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p,
                q, r)

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

@@ -583,6 +583,7 @@ struct DescriptorSource {
 			bool bitset = false;
 			std::optional<CountedRows> rows;
 			std::vector<std::pair<uint32_t, uint32_t>> byte_ranges;
+			std::optional<uint32_t> excluded; // A guard keeps this key from the read.
 
 			bool operator==(const SelectorRead& other) const = default;
 		};

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

@@ -445,10 +445,9 @@ bool ResolveSelectorKeys(const ResourcePlan& program, const DescriptorSource::In
 		     uint64_t {first} + count > material.NumRecords())) return false;
 		const auto last = count != 0u ? uint64_t {first} + count - 1u : 0u;
 		if (last > UINT32_MAX) return false;
+		// Out-of-bounds scalar reads return zero.
 		const auto end = count != 0u ? last * step + selector->offset + 4u : 0u;
-		if (end > uint64_t {UINT32_MAX} + 1u ||
-		    (indirect.workgroup_axis == UINT32_MAX && end > size))
-			return false;
+		if (end > uint64_t {UINT32_MAX} + 1u) return false;
 		if (!read_keys(base, size, uint64_t {first} * step + selector->offset, step, count))
 			return false;
 	} else if (!indirect.selector_mask.IsEmpty()) {
@@ -503,6 +502,7 @@ bool ResolveSelectorKeys(const ResourcePlan& program, const DescriptorSource::In
 		if (selector->bit_offset != 0u || selector->bit_mask != UINT32_MAX || selector->maximum != UINT32_MAX)
 			for (auto& key: keys)
 				key = std::min((key >> selector->bit_offset) & selector->bit_mask, selector->maximum);
+		if (selector->excluded) std::erase(keys, *selector->excluded);
 		std::ranges::sort(keys);
 		keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
 	}

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

@@ -1139,28 +1139,48 @@ private:
 		LaneQuantifier lanes = LaneQuantifier::All;
 	};
 
+	// The immediate that edge proves value differs from on every lane, if any.
+	std::optional<uint32_t> ExcludedBy(const EdgePredicate& edge, Value value) const {
+		if (edge.lanes != LaneQuantifier::All) return {};
+		const auto* test = edge.condition.TryInstruction();
+		if (test != nullptr && test->NumArgs() == 2u &&
+		    ((test->GetOpcode() == ValueOpcode::INotEqual32 && edge.positive) ||
+		     (test->GetOpcode() == ValueOpcode::IEqual32 && !edge.positive))) {
+			for (uint32_t arg = 0; arg < 2u; ++arg) {
+				uint32_t immediate;
+				if (ImmediateU32(test->Arg(arg), immediate) &&
+				    EquivalentValue(m_program, test->Arg(arg ^ 1u), value)) return immediate;
+			}
+		}
+		return {};
+	}
+
 	bool NonzeroOnEntry(Value value, const Block* block) const {
 		if (block == nullptr) return false;
 		bool witnessed = false;
 		const bool guarded = GuardedOnEntry(block, [](const Block* current) {
 			return current->ImmPredecessors().size() != 1u;
 		}, [&](const EdgePredicate& edge) {
-			if (edge.lanes != LaneQuantifier::All) return false;
-			const auto* test = edge.condition.TryInstruction();
-			if (test != nullptr && test->NumArgs() == 2u &&
-			    ((test->GetOpcode() == ValueOpcode::INotEqual32 && edge.positive) ||
-			     (test->GetOpcode() == ValueOpcode::IEqual32 && !edge.positive))) {
-				for (uint32_t arg = 0; arg < 2u; ++arg) {
-					uint32_t immediate;
-					if (ImmediateU32(test->Arg(arg), immediate) && immediate == 0u &&
-					    EquivalentValue(m_program, test->Arg(arg ^ 1u), value)) return witnessed = true;
-				}
-			}
-			return false;
+			return ExcludedBy(edge, value) == 0u && (witnessed = true);
 		});
 		return guarded && witnessed;
 	}
 
+	// The key that every path from its definition to block rules out, if any.
+	std::optional<uint32_t> ExcludedKeyOnEntry(Value key, const Block* block) const {
+		const auto* definition = key.Resolve().TryInstruction();
+		if (definition == nullptr) return {};
+		std::optional<uint32_t> excluded;
+		const auto defined = [&](const Block* current) { return current == definition->Parent(); };
+		const bool guarded = GuardedOnEntry(block, defined, [&](const EdgePredicate& edge) {
+			const auto immediate = ExcludedBy(edge, key);
+			if (!immediate || (excluded && excluded != immediate)) return false;
+			excluded = immediate;
+			return true;
+		});
+		return guarded ? excluded : std::nullopt;
+	}
+
 	// The workgroup ID axis of key; with first, also of a runtime first + workgroup ID sum.
 	uint32_t WorkgroupAxis(Value key, Value* first = nullptr) const {
 		const auto* builtin = key.Resolve().TryInstruction();
@@ -2344,6 +2364,10 @@ private:
 			const auto* bound = BoundedLoop(index, read->Parent());
 			if (bound != nullptr && bound->GetOpcode() == ValueOpcode::ULessThan32) {
 				count = bound->Arg(1);
+			} else if (workgroup_axis == UINT32_MAX && !require_bound) {
+				// Other bounded indices, such as bit scans, read the records up to their maximum.
+				const auto maximum = SelectorMaximum(index);
+				if (maximum < MaxIndirectDescriptorProbes) count = Value(maximum + 1u);
 			}
 		}
 		// Raw pointers have no V# bounds for enumerating an unbounded selector domain.
@@ -3082,6 +3106,7 @@ private:
 		} else if (!plan.table_indexed &&
 		           MatchScalarMaterialKey(key, indirect, material_source,
 		                                  !sampler && table_source.dword_count == 4u, sampler)) {
+			indirect.selector->excluded = ExcludedKeyOnEntry(key, handle.Parent());
 			// 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;