From 99f4eba2230dfb4193b6f16a115ce370216ecdc3 Mon Sep 17 00:00:00 2001 From: Claxten Date: Tue, 27 Oct 2026 01:32:29 +0100 Subject: [PATCH] shader: read bit-scan records only at ranks their mask buffer words set --- src/graphics/shader/recompiler/ir/Serialize.h | 4 +- src/graphics/shader/recompiler/ir/ShaderIR.h | 10 +++- .../ir/passes/ResourceMaterialization.cpp | 52 ++++++++++++---- .../recompiler/ir/passes/ResourceTracking.cpp | 59 ++++++++++++++++++- tests/ShaderRecompilerComputeTests.cpp | 44 ++++++++++---- 5 files changed, 140 insertions(+), 29 deletions(-) diff --git a/src/graphics/shader/recompiler/ir/Serialize.h b/src/graphics/shader/recompiler/ir/Serialize.h index 1cda07589..29427b229 100644 --- a/src/graphics/shader/recompiler/ir/Serialize.h +++ b/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) diff --git a/src/graphics/shader/recompiler/ir/ShaderIR.h b/src/graphics/shader/recompiler/ir/ShaderIR.h index 22f1f8b64..ce4627186 100644 --- a/src/graphics/shader/recompiler/ir/ShaderIR.h +++ b/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 rows; std::vector> byte_ranges; std::optional excluded; // A guard keeps this key from the read. + std::optional ranks; // Records are the bits some mask word sets. bool operator==(const SelectorRead& other) const = default; }; diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp index a5f622a6b..b1626c737 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp +++ b/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((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(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(selector->offset) + - static_cast(index) * selector->stride; - if (offset > UINT32_MAX) return false; - uint32_t key = 0; - if (!read_selector(material_base, UINT64_MAX, static_cast(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) || diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp index 3d144f302..ad093c2ca 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp +++ b/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 terms; }; - const auto load_guard = reads[0]->Arg(4).Resolve(); - const auto words = static_cast(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(reads.size()) * (scalar ? 1u : 4u); std::vector states(2); std::vector 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 + 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 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) { diff --git a/tests/ShaderRecompilerComputeTests.cpp b/tests/ShaderRecompilerComputeTests.cpp index 3eb5df97d..02151a9b9 100644 --- a/tests/ShaderRecompilerComputeTests.cpp +++ b/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(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 images; for (const auto &descriptor : compiled.resources.images) images.push_back(vulkan.CreateImage2D(test.name, 1, 1, vk::Format::eR32Sfloat,