From de3ed905c877229d5ad3f007fc3ebe450dc7d4a2 Mon Sep 17 00:00:00 2001 From: Claxten Date: Mon, 26 Oct 2026 16:02:19 +0100 Subject: [PATCH] shader: track image table keys divided from a runtime base plus the workgroup ID --- src/graphics/shader/recompiler/ir/Program.cpp | 2 + src/graphics/shader/recompiler/ir/Serialize.h | 3 +- src/graphics/shader/recompiler/ir/ShaderIR.h | 6 +- .../ir/passes/ResourceMaterialization.cpp | 6 ++ .../recompiler/ir/passes/ResourceTracking.cpp | 57 +++++++++++----- tests/ShaderRecompilerComputeTests.cpp | 67 +++++++++++++++++++ 6 files changed, 121 insertions(+), 20 deletions(-) diff --git a/src/graphics/shader/recompiler/ir/Program.cpp b/src/graphics/shader/recompiler/ir/Program.cpp index 92a85b706..0d94cd4aa 100644 --- a/src/graphics/shader/recompiler/ir/Program.cpp +++ b/src/graphics/shader/recompiler/ir/Program.cpp @@ -169,6 +169,8 @@ bool EquivalentSelectorDomain(const ResourcePlan& program, const DescriptorSource::IndirectDescriptor& right) { return left.selector == right.selector && left.signed_count == right.signed_count && left.workgroup_axis == right.workgroup_axis && + left.workgroup_multiplier == right.workgroup_multiplier && + left.workgroup_shift == right.workgroup_shift && EquivalentValue(program, left.key_count, right.key_count) && EquivalentValue(program, left.selector_first, right.selector_first) && EquivalentValue(program, left.selector_mask, right.selector_mask); diff --git a/src/graphics/shader/recompiler/ir/Serialize.h b/src/graphics/shader/recompiler/ir/Serialize.h index 33b204a02..e333df534 100644 --- a/src/graphics/shader/recompiler/ir/Serialize.h +++ b/src/graphics/shader/recompiler/ir/Serialize.h @@ -159,7 +159,8 @@ KYTY_IO_FIELDS(ShaderInfo, a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, 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::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) +KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor, a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, + q, r) KYTY_IO_FIELDS(DescriptorSource, a, b, c) KYTY_IO_FIELDS(SrtRead, a, b) KYTY_IO_FIELDS(ResourceBlock, a, b, c, d) diff --git a/src/graphics/shader/recompiler/ir/ShaderIR.h b/src/graphics/shader/recompiler/ir/ShaderIR.h index 034a17975..943452a08 100644 --- a/src/graphics/shader/recompiler/ir/ShaderIR.h +++ b/src/graphics/shader/recompiler/ir/ShaderIR.h @@ -606,8 +606,12 @@ struct DescriptorSource { bool indexed = false; bool signed_count = false; uint32_t workgroup_axis = UINT32_MAX; + // Nonzero shift: workgroup keys are (selector_first + ID) * multiplier >> shift. + uint32_t workgroup_multiplier = 0; + uint32_t workgroup_shift = 0; Value key_count; - // First selector index, or first descriptor key for a direct or bitmap range. + // First selector index, workgroup key base, or first descriptor key for a direct or bitmap + // range. Value selector_first; Value selector_mask; std::vector sources; diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp index 34330a729..e7312525d 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp +++ b/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp @@ -756,6 +756,12 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p std::iota(keys.begin(), keys.end(), first_key); if (!keys.empty() && keys.back() < keys.front()) std::rotate(keys.begin(), keys.begin() + (UINT32_MAX - first_key + 1ull), keys.end()); + if (indirect.workgroup_shift != 0u) { + for (auto& key: keys) + key = static_cast(uint64_t {key} * indirect.workgroup_multiplier >> + indirect.workgroup_shift); + keys.erase(std::unique(keys.begin(), keys.end()), keys.end()); + } } else { if (!indirect.selector_mask.IsEmpty() && table_value.dword_count != 2u) return false; if (!selector_keys->resolved) { diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp index 5c688867f..207c74d95 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp +++ b/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp @@ -1014,7 +1014,9 @@ private: (!shared_image_heap && (a.table_offset != b.table_offset || 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.workgroup_axis != b.workgroup_axis || + a.workgroup_multiplier != b.workgroup_multiplier || + a.workgroup_shift != b.workgroup_shift || a.sources != b.sources || a.signed_count != b.signed_count || !EquivalentValue(m_program, a.key_count, b.key_count) || a.selector_first.IsEmpty() != b.selector_first.IsEmpty() || @@ -1158,9 +1160,21 @@ private: return guarded && witnessed; } - uint32_t WorkgroupAxis(Value key) const { + // 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(); uint32_t axis = UINT32_MAX; + if (first != nullptr && builtin != nullptr && builtin->GetOpcode() == ValueOpcode::IAdd32) { + for (uint32_t arm = 0; arm < 2u; ++arm) { + axis = WorkgroupAxis(builtin->Arg(arm)); + const auto base = builtin->Arg(arm ^ 1u); + if (axis == UINT32_MAX || + !ValidateRuntimeValue(m_program, base, RuntimeValueType::Integer)) continue; + *first = base; + return axis; + } + return UINT32_MAX; + } if (m_program.stage != ShaderType::Compute || builtin == nullptr || builtin->GetOpcode() != ValueOpcode::GetBuiltin || builtin->NumArgs() != 2u || builtin->Arg(0) != Value(static_cast(StageInputKind::WorkgroupId)) || @@ -1168,6 +1182,27 @@ private: return axis; } + // Matches key = (first + workgroup ID) * multiplier >> shift with shift >= 32; first and the + // division are optional. + bool MatchWorkgroupKey(Value key, DescriptorSource::IndirectDescriptor& indirect) const { + uint32_t bits = 0, multiplier = 0; + const auto* shift = key.Resolve().TryInstruction(); + if (shift != nullptr && shift->GetOpcode() == ValueOpcode::ShiftRightLogical32) { + const auto* high = shift->Arg(0).Resolve().TryInstruction(); + if (high != nullptr && high->GetOpcode() == ValueOpcode::UMulHi && + ImmediateU32(shift->Arg(1), bits) && bits < 32u && + ImmediateU32(high->Arg(1), multiplier)) { + key = high->Arg(0); + bits += 32u; + } + } + indirect.workgroup_axis = WorkgroupAxis(key, &indirect.selector_first); + if (indirect.workgroup_axis == UINT32_MAX) return false; + indirect.workgroup_multiplier = multiplier; + indirect.workgroup_shift = bits; + return true; + } + bool MatchTableOffset(Value value, Value& key, uint32_t& offset, uint32_t& stride) const { offset = 0; for (;;) { @@ -2304,21 +2339,7 @@ private: ((offset | stride | memory->offset) & 3u) == 0u && uint64_t {offset} + memory->offset <= UINT32_MAX; if (affine) { - workgroup_axis = WorkgroupAxis(index); - if (workgroup_axis == UINT32_MAX) { - const auto* sum = index.Resolve().TryInstruction(); - if (sum != nullptr && sum->GetOpcode() == ValueOpcode::IAdd32) { - for (uint32_t arm = 0; arm < 2u; ++arm) { - const auto axis = WorkgroupAxis(sum->Arg(arm)); - const auto base = sum->Arg(arm ^ 1u); - if (axis == UINT32_MAX || - !ValidateRuntimeValue(m_program, base, RuntimeValueType::Integer)) continue; - workgroup_axis = axis; - first = base; - break; - } - } - } + workgroup_axis = WorkgroupAxis(index, &first); const auto* bound = BoundedLoop(index, read->Parent()); if (bound != nullptr && bound->GetOpcode() == ValueOpcode::ULessThan32) { count = bound->Arg(1); @@ -3080,7 +3101,7 @@ private: if (indirect.key_count.IsEmpty() && !MatchUniformizedMaterialKey(key, handle, indirect, material_source, false) && !MatchUniformizedMaterialKey(key, handle, indirect, material_source, true) && - !MatchReversedImageRange(key, indirect)) { + !MatchReversedImageRange(key, indirect) && !MatchWorkgroupKey(key, indirect)) { const auto maximum = SelectorMaximum(key); if (maximum >= MaxIndirectDescriptorProbes) return false; indirect.key_count = Value(maximum + 1u); diff --git a/tests/ShaderRecompilerComputeTests.cpp b/tests/ShaderRecompilerComputeTests.cpp index 244fb22e4..8bc538fa5 100644 --- a/tests/ShaderRecompilerComputeTests.cpp +++ b/tests/ShaderRecompilerComputeTests.cpp @@ -41842,6 +41842,71 @@ void CheckPointerImageTable(VulkanHarness &vulkan) { } } +void CheckWorkgroupDividedImageTable(VulkanHarness &vulkan) { + using O = ShaderOpcode; + constexpr u32 base = 256u, table_offset = 0x100u, first = 3u; + TestCase test; + test.name = "WorkgroupDividedImageTable"; + test.dispatch_z = 4u; + test.bda_mappings = {{0u, 0u}}; + test.initial.resize((base + table_offset + 4u * 32u) / 4u); + test.initial[(base + 0x10u) / 4u] = first; + for (u32 row = 0; row < 4u; ++row) { + auto image = MakeSampledTextureData(Prospero::BufferFormat::k32Float); + image[0] = row + 1u; + image[3] = DstSel(4, 0, 0, 1) | (static_cast(Prospero::ImageType::kColor2D) << 28u); + std::copy_n(image.begin(), 8, test.initial.begin() + (base + table_offset + row * 32u) / 4u); + } + test.user_data[0] = base; + const auto output = MakeStructuredStorageBufferData(0u, test.dispatch_z * 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; + test.compute_info.threads_num[0] = test.compute_info.threads_num[1] = + test.compute_info.threads_num[2] = 1u; + test.compute_info.group_id[2] = true; + test.compute_info.workgroup_register = 44u; + // Row (first + workgroup ID z) / 3, divided by a multiply-high and shift. + auto &code = test.code; + AppendVMovLiteral(&code, 20, 0u); + AppendVMovLiteral(&code, 21, 0u); + code.insert(code.end(), {EncodeSmem0(0x00, 21), EncodeSmem1(0x10, 125), + EncodeSop2(0x02, 23, 21, 44), EncodeSop2(0x35, 23, 23, 255), 0xaaaaaaabu, + EncodeSop2(0x20, 23, 23, InlineU32(1)), EncodeSop2(0x1e, 23, 23, InlineU32(5)), + EncodeSop2(0x02, 23, 23, 255), table_offset, + EncodeSmem0(0x03, 8), EncodeSmem1(0, 23), + EncodeMimg0(0x00, 1), EncodeMimg1(4, 20, 2, 10), + EncodeSop2(0x1e, 25, 44, InlineU32(2)), EncodeVop1(0x01, 30, 25)}); + AppendBufferStoreDword(&code, 4, 30); + AppendEnd(&code); + test.opcodes = {O::S_LOAD_DWORD, O::S_ADD_I32, O::S_MUL_HI_U32, O::S_LSHR_B32, O::S_LSHL_B32, + O::S_LOAD_DWORDX8, O::IMAGE_LOAD, O::V_MOV_B32, O::BUFFER_STORE_DWORD, + O::S_ENDPGM}; + const auto compiled = CompileCase(test, vulkan.SubgroupSize()); + Require(test.name, "reachable rows", compiled.resources.images.size() == 2u, + "descriptor planning enumerated rows no workgroup reaches"); + // The key mapping ends the flattened SRT: its count, then a (key, ordinal) pair per row. + const auto &flat = compiled.resources.flattened_srt; + CompareWords(test, "distinct row keys", {2u, 1u, 0u, 2u, 1u}, + std::vector(flat.end() - std::min(flat.size(), 5u), flat.end())); + std::vector images; + for (const auto &descriptor : compiled.resources.images) + images.push_back(vulkan.CreateImage2D(test.name, 1, 1, vk::Format::eR32Sfloat, + vk::ImageUsageFlagBits::eSampled, + {std::bit_cast(static_cast(descriptor.dwords[0]))}, 1, + vk::ImageLayout::eShaderReadOnlyOptimal)); + auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size(), true); + vulkan.Dispatch(test, compiled, buffer, nullptr, nullptr, nullptr, nullptr, {}, images); + std::vector expected; + for (u32 group = 0; group < test.dispatch_z; ++group) + expected.push_back(std::bit_cast(static_cast((first + group) / 3u + 1u))); + CompareWords(test, "divided workgroup rows", expected, + vulkan.ReadBuffer(test.name, buffer, expected.size())); + vulkan.DestroyBuffer(&buffer); + for (auto &image : images) vulkan.DestroyImage(&image); + std::printf("[compute] %s ok\n", test.name); +} + void CheckNestedPointerImageTable(VulkanHarness &vulkan) { using O = ShaderOpcode; constexpr u32 list_address = 160; @@ -52305,6 +52370,7 @@ int main(int argc, char **argv) { vulkan.CheckNullTextureDescriptors(); CheckCountedRowImageSelector(vulkan); CheckPointerImageTable(vulkan); + CheckWorkgroupDividedImageTable(vulkan); CheckNestedPointerImageTable(vulkan); CheckGuardedRecordImageTable(vulkan); CheckIndirectImageKeySwitch(vulkan); @@ -52483,6 +52549,7 @@ int main(int argc, char **argv) { CheckComputeThreadDimensions(vulkan); CheckCountedRowImageSelector(vulkan); CheckPointerImageTable(vulkan); + CheckWorkgroupDividedImageTable(vulkan); CheckNestedPointerImageTable(vulkan); CheckGuardedRecordImageTable(vulkan); CheckIndirectImageKeySwitch(vulkan);