shader: read bit-scan records only at ranks their mask buffer words set

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

@@ -157,7 +157,9 @@ 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, k)
+KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor::SelectorRead::MaskRanks, a, b, c)
+KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor::SelectorRead, a, b, c, d, e, f, g, h, i, j, k,
+               l)
 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

@@ -572,6 +572,13 @@ struct DescriptorSource {
 
 				bool operator==(const CountedRows& other) const = default;
 			};
+			struct MaskRanks {
+				uint32_t source = UINT32_MAX;
+				uint32_t stride = 0; // Byte spacing of the mask words its load can address.
+				uint32_t offset = 0;
+
+				bool operator==(const MaskRanks& other) const = default;
+			};
 			uint32_t source = UINT32_MAX;
 			uint32_t stride = 0;
 			uint32_t offset = 0;
@@ -579,11 +586,12 @@ struct DescriptorSource {
 			uint32_t bit_offset = 0;
 			uint32_t bit_mask = UINT32_MAX;
 			bool global = false;
-			// Four DWORD bitmap words per record; set-bit ranks subtract from the last key.
+			// Four or eight DWORD bitmap words per record; set ranks subtract from the last key.
 			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.
+			std::optional<MaskRanks> ranks;   // Records are the bits some mask word sets.
 
 			bool operator==(const SelectorRead& other) const = default;
 		};

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

@@ -331,6 +331,18 @@ bool ResolveSelectorKeys(const ResourcePlan& program, const DescriptorSource::In
 		keys.resize(count);
 		return true;
 	};
+	// Reads the key of each record whose rank is set.
+	const auto read_ranks = [&](uint64_t base, uint64_t size, uint64_t first, uint64_t step,
+	                            uint32_t ranks) {
+		keys.reserve(std::popcount(ranks));
+		for (; ranks != 0u; ranks &= ranks - 1u) {
+			const auto offset = first + std::countr_zero(ranks) * step;
+			uint32_t key = 0;
+			if (offset > UINT32_MAX || !read_selector(base, size, offset, {&key, 1})) return false;
+			keys.push_back(key);
+		}
+		return true;
+	};
 	if (selector->rows) {
 		const auto& rows = *selector->rows;
 		DescriptorValue count_value;
@@ -450,8 +462,32 @@ bool ResolveSelectorKeys(const ResourcePlan& program, const DescriptorSource::In
 		// 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) return false;
-		if (!read_keys(base, size, uint64_t {first} * step + selector->offset, step, count))
+		const auto start = uint64_t {first} * step + selector->offset;
+		if (selector->ranks) {
+			// A scan selects the ranks some word of its mask's address class sets.
+			const auto& mask = *selector->ranks;
+			DescriptorValue mask_value;
+			ShaderBufferResource masks;
+			if (count > 32u || !WalkedDescriptor(program, mask.source, mask_value) ||
+			    !DecodeBufferDescriptor(mask_value, masks)) return false;
+			// Each chunk holds whole strides of the power-of-two class, or one of its words.
+			const auto words = mask.stride / 4u;
+			const auto chunk =
+			    words <= MaxIndirectDescriptorProbes ? MaxIndirectDescriptorProbes : 1u;
+			uint32_t ranks = 0;
+			for (auto at = uint64_t {mask.offset}; at + 4u <= masks.GetSize();
+			     at += uint64_t {std::max(chunk, words)} * 4u) {
+				keys.resize(std::min<uint64_t>((masks.GetSize() - at) / 4u, chunk));
+				if (!ReadScalarTable(masks.Base48(), masks.GetSize(), at, runtime, keys))
 					return false;
+				for (size_t word = 0; word < keys.size(); word += words) ranks |= keys[word];
+			}
+			keys.clear();
+			if (count < 32u) ranks &= (1u << count) - 1u;
+			if (!read_ranks(base, size, start, step, ranks)) return false;
+		} else if (!read_keys(base, size, start, step, count)) {
+			return false;
+		}
 	} else if (!indirect.selector_mask.IsEmpty()) {
 		uint32_t mask = 0;
 		uint32_t count = 0;
@@ -463,18 +499,8 @@ bool ResolveSelectorKeys(const ResourcePlan& program, const DescriptorSource::In
 		if (count < 32u) mask &= (1u << count) - 1u;
 		const auto material_base =
 		    (static_cast<uint64_t>(material_value.dwords[1]) << 32u) | material_value.dwords[0];
-		keys.reserve(std::popcount(mask));
-		while (mask != 0u) {
-			const auto index = std::countr_zero(mask);
-			const auto offset = static_cast<uint64_t>(selector->offset) +
-			                    static_cast<uint64_t>(index) * selector->stride;
-			if (offset > UINT32_MAX) return false;
-			uint32_t key = 0;
-			if (!read_selector(material_base, UINT64_MAX, static_cast<uint32_t>(offset),
-			                   {&key, 1})) return false;
-			keys.push_back(key);
-			mask &= mask - 1u;
-		}
+		if (!read_ranks(material_base, UINT64_MAX, selector->offset, selector->stride, mask))
+			return false;
 	} else if (selector->stride == 0u) {
 		ShaderBufferResource material;
 		if (!DecodeBufferDescriptor(material_value, material) ||

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

@@ -2382,8 +2382,10 @@ private:
 			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.
+				// Other bounded indices, such as bit scans, read the records up to their maximum;
+				// a bit scan of a nonzero mask load reads only the ranks its words set.
 				const auto maximum = SelectorMaximum(index);
+				if (maximum == 31u) selector.ranks = BitScanRanks(index);
 				if (maximum < MaxIndirectDescriptorProbes) count = Value(maximum + 1u);
 			}
 			// A record selected from a shared append list is below the list's bound.
@@ -2580,8 +2582,10 @@ private:
 			bool all = true;
 			std::vector<size_t> terms;
 		};
-		const auto load_guard = reads[0]->Arg(4).Resolve();
-		const auto words = static_cast<uint32_t>(reads.size()) * 4u;
+		// A scalar load reads one word without EXEC; four-word loads take their guard.
+		const bool scalar = reads[0]->GetOpcode() == ValueOpcode::ReadConstBuffer;
+		const auto load_guard = scalar ? Value(true) : reads[0]->Arg(4).Resolve();
+		const auto words = static_cast<uint32_t>(reads.size()) * (scalar ? 1u : 4u);
 		std::vector<State> states(2);
 		std::vector<Equation> equations(2);
 		const auto add = [&](State state) {
@@ -2649,6 +2653,9 @@ private:
 							equation.terms = {1u};
 						break;
 					}
+					case ValueOpcode::ReadConstBuffer:
+						if (scalar && inst == reads[state.component]) equation.terms = {1u};
+						break;
 					case ValueOpcode::Phi:
 						// The predecessor may be an earlier outer iteration. Establish
 						// this register's invariant without the current word or EXEC.
@@ -2717,6 +2724,52 @@ private:
 		return std::ranges::all_of(roots, [&](size_t root) { return valid[root]; });
 	}
 
+	// The mask words whose set bits are every record a bit scan of index selects: the scanned
+	// value keeps only bits of one scalar load, which reads a word its address can reach.
+	std::optional<DescriptorSource::IndirectDescriptor::SelectorRead::MaskRanks>
+	BitScanRanks(Value index) {
+		const auto* scan = index.Resolve().TryInstruction();
+		if (scan == nullptr || scan->GetOpcode() != ValueOpcode::FindILsb32) return {};
+		const auto* masked = scan->Arg(0).Resolve().TryInstruction();
+		if (masked == nullptr) return {};
+		std::vector<const Inst*> values {masked};
+		const Inst* load = nullptr;
+		for (size_t i = 0; i < values.size() && load == nullptr; ++i) {
+			const auto* value = values[i];
+			if (value->GetOpcode() == ValueOpcode::ReadConstBuffer) load = value;
+			if (value->GetOpcode() != ValueOpcode::Phi &&
+			    value->GetOpcode() != ValueOpcode::BitwiseAnd32) continue;
+			for (size_t arg = 0; arg < value->NumArgs(); ++arg) {
+				const auto* next = value->Arg(arg).Resolve().TryInstruction();
+				if (next != nullptr && std::ranges::find(values, next) == values.end())
+					values.push_back(next);
+			}
+		}
+		uint32_t memory_index = 0;
+		const auto* memory = load != nullptr ? ScalarReadMemory(*load, memory_index) : nullptr;
+		if (memory == nullptr || memory->offset > INT32_MAX ||
+		    !MemoryIndexBelongsTo(memory_index, *load) ||
+		    !BitsetMaskSubset(scan->Arg(0), Value(true), Value(0u), std::array {load})) return {};
+		auto address = load->Arg(1);
+		auto active = Value(true);
+		const auto* lane = address.Resolve().TryInstruction();
+		if (lane != nullptr && lane->GetOpcode() == ValueOpcode::ReadFirstLane) {
+			if (!HasActiveLane(lane->Arg(1), lane->Parent())) return {};
+			active = lane->Arg(1);
+			address = lane->Arg(0);
+		}
+		// A left shift keeps the address a multiple of its scale; other addresses read any word.
+		const auto* shift = ActiveValue(address, active).TryInstruction();
+		uint32_t bits = 0;
+		if (shift == nullptr || shift->GetOpcode() != ValueOpcode::ShiftLeftLogical32 ||
+		    !ImmediateU32(shift->Arg(1), bits) || bits >= 32u) bits = 0;
+		DescriptorSource source;
+		if (!MakeRuntimeTableSource(*load, source)) return {};
+		return DescriptorSource::IndirectDescriptor::SelectorRead::MaskRanks {
+		    .source = InternSource(source), .stride = std::max(1u << bits, 4u),
+		    .offset = memory->offset & ~3u};
+	}
+
 	bool MatchBitsetSelector(Value rank, Value active, SelectorBounds bounds, uint32_t words,
 	                         DescriptorSource::IndirectDescriptor& indirect) {
 		const auto selected_value = [&](Value value) {

tests/ShaderRecompilerComputeTests.cpp

@@ -42289,13 +42289,17 @@ void CheckClampedSiblingImageTable(VulkanHarness &vulkan) {
 void CheckGuardedSelectorImageTable(VulkanHarness &vulkan) {
   using O = ShaderOpcode;
   constexpr u32 base = 256u, table_offset = 0x100u;
-  // A loop over the set bits of a mask reads each record's row from a four-record buffer. Rows
-  // past it read zero, and a guard keeps row -1, whose descriptor is a 3D image, from the table.
+  // A loop over the set bits of a mask word, which a lane selects from a mask buffer, reads each
+  // record's row from a four-record buffer. Only the words a lane can address select records, so
+  // record 1's float row, set in another word, is never read. Rank 5 is past the records and
+  // reads zero, and a guard keeps row -1, whose descriptor is a 3D image, from the table.
   // A guard on another value keeps row -1, which a table behind a V# reads through the heap.
+  // A loop guard on another value lets an empty mask's scan return -1, which reads key 0 past
+  // the records; the set ranks then miss a key, and no domain holds both, so tracking rejects it.
   constexpr std::array names{"GuardedSelectorImageTable", "GuardedSelectorImageTableOther",
-                             "GuardedSelectorImageTableHeap"};
+                             "GuardedSelectorImageTableHeap", "GuardedSelectorImageTableEmpty"};
   for (u32 variant = 0; variant < names.size(); ++variant) {
-    const bool guarded = variant == 0u, heap = variant == 2u;
+    const bool heap = variant == 2u, empty = variant == 3u, guarded = variant == 0u || empty;
     TestCase test;
     test.name = names[variant];
     test.bda_mappings = {{0u, 0u}};
@@ -42307,18 +42311,25 @@ void CheckGuardedSelectorImageTable(VulkanHarness &vulkan) {
       image[3] = DstSel(4, 0, 0, 1) | (static_cast<u32>(type) << 28u);
       std::copy_n(image.begin(), 8, test.initial.begin() + (base + table_offset + row * 32) / 4u);
     }
-    for (const auto [record, row] : {std::pair{0u, 2u}, {1u, 5u}, {2u, UINT32_MAX}, {3u, 1u}})
+    for (const auto [record, row] :
+         {std::pair{0u, 2u}, {1u, 0x3f800000u}, {2u, UINT32_MAX}, {3u, 1u}})
       test.initial[(record * 16u + 8u) / 4u] = row;
+    // Lanes address the mask words at multiples of 16 bytes, between the record keys.
+    for (const auto [byte, ranks] : {std::pair{0u, 0b1001u}, {4u, 0b10u}, {16u, 0b100100u}})
+      test.initial[byte / 4u] = ranks;
     auto table = MakeStructuredStorageBufferData(0u, table_offset + 8u * 32u);
     table[0] = base;
     std::copy_n(table.begin(), heap ? 4 : 1, test.user_data.begin());
-    test.user_data[8] = 0b1101u;
     auto records = MakeStructuredStorageBufferData(0u, 64u);
     std::copy_n(records.begin(), 4, test.user_data.begin() + 12);
+    const auto masks = MakeStructuredStorageBufferData(0u, 32u);
+    std::copy_n(masks.begin(), 4, test.user_data.begin() + 24);
     const auto output = MakeStructuredStorageBufferData(0u, 4u);
     std::copy_n(output.begin(), 4, test.user_data.begin() + 48);
     test.has_user_data = test.has_compute_info = true;
     test.compute_info.wave_size = 64u;
+    // The mask word's address sums two lane ids.
+    test.compute_info.thread_ids_num = 2u;
     test.compute_info.threads_num[0] = test.compute_info.threads_num[1] =
         test.compute_info.threads_num[2] = 1u;
     auto &code = test.code;
@@ -42326,9 +42337,12 @@ void CheckGuardedSelectorImageTable(VulkanHarness &vulkan) {
     AppendVMovLiteral(&code, 21, 0u);
     AppendVMovLiteral(&code, 6, 0u);
     AppendVMovLiteral(&code, 30, 0u);
-    code.push_back(EncodeSop1(0x03, 40, 8));
+    const auto active_branch = code.size();
+    code.insert(code.end(), {0u, EncodeVop2(0x25, 7, Vgpr(0), 1),
+        EncodeVop2(0x1a, 7, InlineU32(4), 7), EncodeVop1(0x02, 42, Vgpr(7)),
+        EncodeSmem0(0x08, 40, 12), EncodeSmem1(0, 42)});
     const auto header = code.size();
-    code.push_back(EncodeSopc(0x07, InlineU32(0), 40));
+    code.push_back(EncodeSopc(0x07, InlineU32(0), empty ? 42 : 40));
     const auto exit_branch = code.size();
     code.push_back(0u);
     code.insert(code.end(), {EncodeSop1(0x13, 41, 40), EncodeSop2(0x1e, 5, 41, InlineU32(4)),
@@ -42344,20 +42358,28 @@ void CheckGuardedSelectorImageTable(VulkanHarness &vulkan) {
     code.push_back(EncodeSop1(0x1b, 40, 41));
     code.push_back(EncodeSopp(0x02, header - code.size() - 1u));
     code[exit_branch] = EncodeSopp(0x04, code.size() - exit_branch - 1u);
+    code[active_branch] = EncodeSopp(0x08, code.size() - active_branch - 1u);
     AppendBufferStoreDword(&code, 6, 30);
     AppendEnd(&code);
+    if (empty) {
+#if KYTY_PLATFORM != KYTY_PLATFORM_WINDOWS
+      ExpectFatal(test.name, [&] { MaterializeCase(test); });
+#endif
+      continue;
+    }
     if (!guarded) {
       Require(test.name, "excluded row", MaterializeCase(test) == heap,
               "a guard on another value excluded row -1, or a V# table enumerated it");
       continue;
     }
-    test.opcodes = {O::V_MOV_B32, O::S_MOV_B32, O::S_CMP_LG_U32, O::S_CBRANCH_SCC0,
+    test.opcodes = {O::V_MOV_B32, O::S_CBRANCH_EXECZ, O::V_ADD_NC_U32, O::V_LSHLREV_B32,
+                    O::V_READFIRSTLANE_B32, O::S_CMP_LG_U32, O::S_CBRANCH_SCC0,
                     O::S_FF1_I32_B32, O::S_LSHL_B32, O::S_BUFFER_LOAD_DWORD, O::S_ADD_I32,
                     O::S_LOAD_DWORDX8, O::IMAGE_LOAD, O::V_ADD_F32, O::S_BITSET0_B32,
                     O::S_BRANCH, O::BUFFER_STORE_DWORD, O::S_ENDPGM};
     const auto compiled = CompileCase(test, vulkan.SubgroupSize());
-    Require(test.name, "selected rows", compiled.resources.images.size() == 4u,
-            "record keys did not select rows 0, 1, 2 and 5");
+    Require(test.name, "selected rows", compiled.resources.images.size() == 3u,
+            "record keys did not select rows 0, 1 and 2");
     std::vector<VulkanHarness::Image> images;
     for (const auto &descriptor : compiled.resources.images)
       images.push_back(vulkan.CreateImage2D(test.name, 1, 1, vk::Format::eR32Sfloat,