diff --git a/src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp b/src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp index 6f60780d0..2eca0762a 100644 --- a/src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp +++ b/src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp @@ -299,11 +299,14 @@ struct PipelineCache::ProgramCache { lookup_key.code_size = static_cast(params.code.size()); BuildStageStaticKey(input_info, lookup_key.static_state); auto entry = programs.find(lookup_key); - const ShaderRecompiler::IR::SrtRuntime runtime { + ShaderRecompiler::IR::SrtRuntime runtime { .user_data = user_data, .shader_base = params.Base(), .read_specialization_memory = ReadShaderGuestMemory, }; + if constexpr (std::is_same_v) { + runtime.workgroup_counts = input_info.workgroup_counts; + } if (entry != programs.end()) { EXIT_IF(!ShaderRecompiler::IR::MaterializeResources( entry->second.resource_plan, runtime, entry->second.resources, diff --git a/src/graphics/host_gpu/renderer/renderCompute.cpp b/src/graphics/host_gpu/renderer/renderCompute.cpp index 4f9c0673e..ef8ab03a5 100644 --- a/src/graphics/host_gpu/renderer/renderCompute.cpp +++ b/src/graphics/host_gpu/renderer/renderCompute.cpp @@ -254,6 +254,18 @@ void RenderExecutor::DispatchDirect(uint64_t submit_id, CommandBuffer& buffer, ShaderComputeInputInfo input_info {}; const bool use_thread_dimensions = (mode & DISPATCH_INITIATOR_USE_THREAD_DIMENSIONS) != 0; input_info.dispatch_thread_dimensions = use_thread_dimensions; + input_info.workgroup_counts[0] = thread_group_x; + input_info.workgroup_counts[1] = thread_group_y; + input_info.workgroup_counts[2] = thread_group_z; + if (use_thread_dimensions) { + const uint32_t group_sizes[] = {cs_regs.cs_regs.num_thread_x, cs_regs.cs_regs.num_thread_y, + cs_regs.cs_regs.num_thread_z}; + for (uint32_t axis = 0; axis < 3u; ++axis) { + const auto size = std::max(group_sizes[axis], 1u); + const auto threads = input_info.workgroup_counts[axis]; + input_info.workgroup_counts[axis] = threads / size + (threads % size != 0u); + } + } const auto compute_program = m_context.GetPipelineCache().GetComputeProgram(cs_regs, sh_regs, input_info); if (use_thread_dimensions) { @@ -273,18 +285,12 @@ void RenderExecutor::DispatchDirect(uint64_t submit_id, CommandBuffer& buffer, } if (use_thread_dimensions) { - auto groups_from_threads = [](uint32_t threads, uint32_t group_size) { - return (threads == 0 - ? 0u - : (threads + std::max(group_size, 1u) - 1u) / std::max(group_size, 1u)); - }; - const uint32_t old_x = thread_group_x; const uint32_t old_y = thread_group_y; const uint32_t old_z = thread_group_z; - thread_group_x = groups_from_threads(thread_group_x, cs_regs.cs_regs.num_thread_x); - thread_group_y = groups_from_threads(thread_group_y, cs_regs.cs_regs.num_thread_y); - thread_group_z = groups_from_threads(thread_group_z, cs_regs.cs_regs.num_thread_z); + thread_group_x = input_info.workgroup_counts[0]; + thread_group_y = input_info.workgroup_counts[1]; + thread_group_z = input_info.workgroup_counts[2]; static std::atomic log_count {0}; if (log_count.fetch_add(1, std::memory_order_relaxed) < 32) { diff --git a/src/graphics/shader/recompiler/ir/ShaderIR.h b/src/graphics/shader/recompiler/ir/ShaderIR.h index 56a6233f4..625b877aa 100644 --- a/src/graphics/shader/recompiler/ir/ShaderIR.h +++ b/src/graphics/shader/recompiler/ir/ShaderIR.h @@ -488,6 +488,8 @@ struct DescriptorSource { uint32_t selector_stride = 0; uint32_t selector_offset = 0; uint32_t table_offset = 0; + uint32_t table_stride = 0; + uint32_t workgroup_axis = UINT32_MAX; uint32_t selector_shift = 0; Value key_count; Value selector_first; diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp index a7b919c13..ec8762f66 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp +++ b/src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp @@ -285,12 +285,19 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p } if (indirect.material_source == UINT32_MAX) { uint32_t key_count = 0; - const bool evaluated = clean.Evaluate(indirect.key_count, key_count); - if (std::bit_cast(key_count) <= 0) key_count = 0; - if (table_value.dword_count != 2u || !evaluated || - key_count > MaxIndirectDescriptorProbes || - uint64_t {indirect.table_offset} + uint64_t {key_count} * descriptor_bytes > - UINT32_MAX + 1ull) { + if (indirect.workgroup_axis != UINT32_MAX) { + if (indirect.workgroup_axis >= runtime.workgroup_counts.size()) return false; + key_count = runtime.workgroup_counts[indirect.workgroup_axis]; + if (key_count == 0u) return false; + } else { + if (table_value.dword_count != 2u || !clean.Evaluate(indirect.key_count, key_count)) + return false; + if (std::bit_cast(key_count) <= 0) key_count = 0; + } + if (indirect.table_stride == 0u || key_count > MaxIndirectDescriptorProbes || + (key_count != 0u && uint64_t {indirect.table_offset} + + uint64_t {key_count - 1u} * indirect.table_stride + descriptor_bytes > + UINT32_MAX + 1ull)) { return false; } keys.resize(key_count); @@ -376,7 +383,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p candidate.dword_count = dword_count; if (sources.empty()) { const auto table_offset = - static_cast(key * descriptor_bytes) + indirect.table_offset; + static_cast(key * indirect.table_stride) + indirect.table_offset; if (!ReadScalarTable(table_base, table_size, table_offset, runtime, std::span(candidate.dwords).first(dword_count))) { return false; diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp index 7fde52b2c..1284aa6d6 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp +++ b/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp @@ -834,7 +834,8 @@ private: if (a.material_source != b.material_source || a.table_source != b.table_source || a.selector_stride != b.selector_stride || a.selector_offset != b.selector_offset || a.selector_shift != b.selector_shift || - a.table_offset != b.table_offset || a.sources != b.sources || + a.table_offset != b.table_offset || a.table_stride != b.table_stride || + a.workgroup_axis != b.workgroup_axis || a.sources != b.sources || !EquivalentValue(m_program, a.key_count, b.key_count) || a.selector_first.IsEmpty() != b.selector_first.IsEmpty() || (!a.selector_first.IsEmpty() && @@ -1010,7 +1011,17 @@ private: return false; } - bool MatchTableOffset(Value value, Value& key, uint32_t& offset, uint32_t shift = 5u) const { + uint32_t WorkgroupAxis(Value key) const { + const auto* builtin = key.Resolve().TryInstruction(); + uint32_t axis = 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)) || + !ImmediateU32(builtin->Arg(1), axis) || axis >= 3u) return UINT32_MAX; + return axis; + } + + bool MatchTableOffset(Value value, Value& key, uint32_t& offset, uint32_t& stride) const { offset = 0; for (;;) { const auto* inst = value.Resolve().TryInstruction(); @@ -1019,10 +1030,17 @@ private: } uint32_t immediate; if (inst->GetOpcode() == ValueOpcode::ShiftLeftLogical32 && - ImmediateU32(inst->Arg(1), immediate) && immediate == shift) { + ImmediateU32(inst->Arg(1), immediate) && immediate < 32u) { key = inst->Arg(0).Resolve(); + stride = 1u << immediate; return key.GetType() == Type::U32; } + if (inst->GetOpcode() == ValueOpcode::IMul32) { + if (ImmediateU32(inst->Arg(0), stride)) key = inst->Arg(1).Resolve(); + else if (ImmediateU32(inst->Arg(1), stride)) key = inst->Arg(0).Resolve(); + else return false; + return stride != 0u && key.GetType() == Type::U32; + } if (inst->GetOpcode() != ValueOpcode::IAdd32) { return false; } @@ -1563,7 +1581,8 @@ private: bool MatchDescriptorTable(Inst& handle, const DescriptorSource& descriptor, IndirectDescriptorPlan& plan, - DescriptorSource& table_source, Value& key, uint32_t& table_offset) { + DescriptorSource& table_source, Value& key, uint32_t& table_offset, + uint32_t& table_stride) { Inst* table_handle = nullptr; for (uint32_t dword = 0; dword < handle.NumArgs(); ++dword) { auto* read = descriptor.dwords[dword].Resolve().TryInstruction(); @@ -1579,6 +1598,7 @@ private: auto* current_handle = read->Arg(0).Resolve().TryInstruction(); Value current_key; uint32_t offset = 0; + uint32_t stride = 0; if (current_handle == nullptr || current_handle->GetOpcode() != (memory->kind == ResourceKind::ScalarAddress ? ValueOpcode::GetAddressResource @@ -1586,7 +1606,7 @@ private: (memory->kind == ResourceKind::ScalarAddress && read->Parent() != handle.Parent()) || (table_handle != nullptr && !EquivalentValue(m_program, Value(table_handle), Value(current_handle))) || - !MatchTableOffset(read->Arg(1), current_key, offset, handle.NumArgs() == 4u ? 4u : 5u) || + !MatchTableOffset(read->Arg(1), current_key, offset, stride) || memory->offset > UINT32_MAX - offset) { return false; } @@ -1594,12 +1614,14 @@ private: if (dword == 0u) { key = current_key; table_offset = offset; - } else if (!EquivalentValue(m_program, key, current_key) || + table_stride = stride; + } else if (table_stride != stride || !EquivalentValue(m_program, key, current_key) || static_cast(table_offset) + dword * sizeof(uint32_t) != offset) { return false; } table_handle = current_handle; - if (handle.NumArgs() == 8u ? !UsesOnlyImageDescriptors(*read) : + if (handle.NumArgs() == 8u + ? WorkgroupAxis(key) == UINT32_MAX && !UsesOnlyImageDescriptors(*read) : std::ranges::any_of(read->Uses(), [](const Use& use) { return use.user->GetOpcode() != ValueOpcode::GetBufferResource || std::ranges::any_of(use.user->Uses(), [](const Use& consumer) { @@ -1654,11 +1676,21 @@ private: DescriptorSource table_source; Value key; uint32_t table_offset = 0; - if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset)) return false; + uint32_t table_stride = 0; + if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset, + table_stride)) return false; DescriptorSource material_source; DescriptorSource::IndirectDescriptor indirect; indirect.table_offset = table_offset; - if (table_source.dword_count == 2u) { + indirect.table_stride = table_stride; + indirect.workgroup_axis = WorkgroupAxis(key); + if (indirect.workgroup_axis != UINT32_MAX) { + // The descriptor also supplies dimensions to shader arithmetic. Keep its reads; + // only the image handle is projected onto the bounded workgroup key. + plan.reads.fill(nullptr); + } else if (table_stride != 32u) { + return false; + } else if (table_source.dword_count == 2u) { const auto* selector = key.Resolve().TryInstruction(); const bool bitscan = selector != nullptr && selector->GetOpcode() == ValueOpcode::FindILsb32 && selector->NumArgs() == 1u && !m_shader_writes && @@ -1733,7 +1765,9 @@ private: DescriptorSource table_source; Value key; uint32_t table_offset = 0; - if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset) || + uint32_t table_stride = 0; + if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset, + table_stride) || table_stride != 16u || table_source.dword_count != 4u || table_offset != 0u) return false; const auto* lane_read = key.Resolve().TryInstruction(); if (lane_read == nullptr || lane_read->GetOpcode() != ValueOpcode::ReadFirstLane) @@ -1763,6 +1797,7 @@ private: !EquivalentValue(m_program, selected->Arg(0), read->Arg(4))) return false; DescriptorSource material_source; DescriptorSource::IndirectDescriptor indirect; + indirect.table_stride = table_stride; if (!MakeRuntimeTableSource(*read, material_source) || !BoundedSelectedIndex(read->Arg(1), read->Arg(4), indirect)) return false; indirect.material_source = InternSource(material_source); diff --git a/src/graphics/shader/recompiler/ir/passes/SrtWalker.h b/src/graphics/shader/recompiler/ir/passes/SrtWalker.h index 4903d11bb..50cb67bce 100644 --- a/src/graphics/shader/recompiler/ir/passes/SrtWalker.h +++ b/src/graphics/shader/recompiler/ir/passes/SrtWalker.h @@ -17,6 +17,7 @@ struct SrtRuntime { SrtMemoryReader read_memory = nullptr; void* userdata = nullptr; SrtMemoryReader read_specialization_memory = nullptr; + std::span workgroup_counts; }; enum class RuntimeValueType { Any, Integer }; diff --git a/src/graphics/shader/shader.cpp b/src/graphics/shader/shader.cpp index ba849a7a7..766c06c66 100644 --- a/src/graphics/shader/shader.cpp +++ b/src/graphics/shader/shader.cpp @@ -571,9 +571,11 @@ static void ShaderGetStaticInputInfoCS(const HW::ComputeShaderInfo& regs, const ShaderMappedData& data, ShaderComputeInputInfo& info) { const bool dispatch_thread_dimensions = info.dispatch_thread_dimensions; const auto host_subgroup_size = info.host_subgroup_size; + const auto workgroup_counts = std::to_array(info.workgroup_counts); info = {}; info.dispatch_thread_dimensions = dispatch_thread_dimensions; info.host_subgroup_size = host_subgroup_size; + std::ranges::copy(workgroup_counts, info.workgroup_counts); info.threads_num[0] = regs.cs_regs.num_thread_x; info.threads_num[1] = regs.cs_regs.num_thread_y; info.threads_num[2] = regs.cs_regs.num_thread_z; diff --git a/src/graphics/shader/shader.h b/src/graphics/shader/shader.h index 08c041226..3c54a28e3 100644 --- a/src/graphics/shader/shader.h +++ b/src/graphics/shader/shader.h @@ -148,6 +148,7 @@ struct ShaderVertexInputInfo { struct ShaderComputeInputInfo: ShaderWorkgroupInputInfo { uint8_t float_mode = 0xc0; uint32_t dispatch_threads_num[3] = {0, 0, 0}; + uint32_t workgroup_counts[3] = {0, 0, 0}; bool group_id[3] = {false, false, false}; bool dispatch_thread_dimensions = false; int thread_ids_num = 0; diff --git a/tests/ShaderRecompilerComputeTests.cpp b/tests/ShaderRecompilerComputeTests.cpp index 2bf6003c1..9a4c152f7 100644 --- a/tests/ShaderRecompilerComputeTests.cpp +++ b/tests/ShaderRecompilerComputeTests.cpp @@ -9921,6 +9921,101 @@ public: std::vector(8, image_value + count), "raw buffer descriptor materialization overwrote the newer native image"); } + + // PPSA24156 loads an image sharp from a 148-byte record selected by + // WorkGroupID.z, then uses the same descriptor words for dimensions. + constexpr auto parameters = base + 0x80000; + constexpr auto table_address = parameters + 0x100; + std::vector table_shader{ + EncodeSop2(0x26, 8, 4, 255), 148, + EncodeSmem0(0x02, 16), EncodeSmem1(0, 125), + EncodeSmem0(0x0b, 32, 8), EncodeSmem1(0, 8), + EncodeSop2(0x1e, 8, 34, InlineU32(2)), + EncodeSop2(0x20, 9, 33, InlineU32(30)), + EncodeSop2(0x0e, 8, 8, 255), 0xfffcu, + EncodeSop2(0x10, 8, 8, 9), + EncodeSop2(0x27, 9, 34, 255), 0x000e000eu, + EncodeSop2(0x02, 8, 8, 9), + EncodeSop2(0x02, 8, 8, InlineU32(2)), + EncodeVop1(0x01, 0, 4), EncodeVop1(0x01, 1, InlineU32(0)), + EncodeVop1(0x01, 2, 8), + EncodeMimg0(0x08, 1), EncodeMimg1(2, 0, 8)}; + AppendEnd(&table_shader); + ShaderMapUserData(reinterpret_cast(table_shader.data()), + {.type = Prospero::ShaderBinaryType::kCs, + .code_size_bytes = static_cast(table_shader.size() * sizeof(u32))}); + shaders.SetCsShader({.data_addr = reinterpret_cast(table_shader.data()), + .num_thread_x = 1, .num_thread_y = 1, .num_thread_z = 1, + .wave_size = 64, .user_sgpr = 2, .tgid_x_en = true, + .tgid_y_en = true, .tgid_z_en = true}); + shaders.SetCsUserSgpr(0, static_cast(parameters), HW::UserSgprType::Unknown); + shaders.SetCsUserSgpr(1, static_cast(parameters >> 32u), HW::UserSgprType::Unknown); + const ShaderBufferResource table_buffer{{static_cast(table_address), + static_cast(table_address >> 32u) | (148u << 16u), 2, 0x5204u}}; + LibKernel::Memory::WriteBacking(parameters, table_buffer.fields, sizeof(table_buffer.fields)); + struct TableCase { u32 groups, width, height, image; }; + uint64_t table_program_id = 0; + for (const auto test : {TableCase{1, 4, 2, 0}, TableCase{2, 4, 2, 0}, + TableCase{2, 8, 3, 1}, TableCase{1, 4, 2, 0}}) { + const auto address = base + 0x90000 + test.image * 0x10000; + const ShaderTextureResource sharp{{static_cast(address >> 8u), + (static_cast(Prospero::BufferFormat::k32UInt) << 20u) | + (((test.width - 1u) & 3u) << 30u), + ((test.width - 1u) >> 2u) | ((test.height - 1u) << 14u), + DstSel(4, 5, 6, 7) | + (static_cast(Prospero::TileMode::kLinear) << 20u) | + (static_cast(Prospero::ImageType::kColor2D) << 28u), + 0, 0x00700000u, 0, 0}}; + std::array table{}; + for (u32 group = 0; group < test.groups; group++) { + std::copy_n(sharp.fields, 8, table.begin() + group * 37); + } + Require(name, "descriptor table refresh", context.InvalidateMemory(table_address, sizeof(table)), + "the descriptor table escaped its mapped allocation"); + LibKernel::Memory::WriteBacking(table_address, table.data(), sizeof(table)); + ShaderComputeInputInfo input{}; + input.workgroup_counts[0] = input.workgroup_counts[1] = 1; + input.workgroup_counts[2] = test.groups; + const auto program = context.GetPipelineCache().GetComputeProgram( + shaders.GetCs(), processor.GetCtx().GetShaderRegisters(), input); + Require(name, "workgroup table refresh and coalescing", + input.workgroup_counts[2] == test.groups && + input.stage.program->info.images.size() == 1 && + input.stage.program->info.images[0].indirect_resources.empty() && + input.stage.resources->images.size() == 1 && + std::equal(std::begin(sharp.fields), std::end(sharp.fields), + input.stage.resources->images[0].dwords.begin()) && + !input.stage.program->info.buffers.empty(), + "workgroup bounds, refreshed sharp, or live descriptor dimension reads were lost"); + Require(name, "workgroup table program reuse", + table_program_id == 0 || program.id == table_program_id, + "changing the workgroup count or identical sharp values recompiled the shader"); + table_program_id = program.id; + const auto *snapshot = input.stage.resources; + const auto *image_storage = snapshot->images.data(); + const auto flattened_capacity = snapshot->flattened_srt.capacity(); + const auto repeated = context.GetPipelineCache().GetComputeProgram( + shaders.GetCs(), processor.GetCtx().GetShaderRegisters(), input); + Require(name, "workgroup table cache-hit storage reuse", + repeated.id == program.id && input.stage.resources == snapshot && + snapshot->images.data() == image_storage && + snapshot->flattened_srt.capacity() == flattened_capacity, + "a repeated descriptor table hit replaced its snapshot or storage"); + auto binding = RenderExecutorTestAccess::ResolveTexture( + context.GetRenderExecutor(), input.stage.program->info.images[0], + input.stage.resources->images[0]); + const auto target = textures.FindImage(binding.desc); + TextureCacheTestAccess::ClearImage(textures, scheduler.Current(), target, + {vk::ImageAspectFlagBits::eColor, 0, 1, 0, 1}, {}); + processor.DispatchDirect(1, 1, test.groups, 0x41u); + const auto pixels = ReadCachedTexel(name, context, target, {}, + {test.width, test.height, 1}); + for (u32 pixel = 0; pixel < pixels.size(); pixel++) { + const auto expected = pixel < test.groups ? test.width + test.height : 0u; + Require(name, "workgroup table GPU dimension reads", pixels[pixel] == expected, + "a retained descriptor load used stale dimensions or the wrong record stride"); + } + } context.UnmapMemory(base, allocation_size); scheduler.Finish(); }); @@ -32161,7 +32256,7 @@ void CheckIndirectBufferStore(VulkanHarness &vulkan) { candidate.indirect_search_iterations = 0; candidate.indirect_resources.clear(); program.descriptor_sources.resize(1); - program.descriptor_sources[0].indirect_descriptor.emplace(); + program.descriptor_sources[0].indirect_descriptor.emplace().table_stride = 16; ShaderComputeInputInfo compute; compute.wave_size = 32; compute.host_subgroup_size = vulkan.SubgroupSize(); @@ -32272,7 +32367,8 @@ void CheckIndirectImageKeySwitch() { program.descriptor_sources.resize(2); program.descriptor_sources[0].dword_count = 8; program.descriptor_sources[0].indirect_descriptor = - DescriptorSource::IndirectDescriptor{0u, 0u, 224u, 12u, 0u}; + DescriptorSource::IndirectDescriptor{.material_source = 0, .table_source = 0, + .selector_stride = 224, .selector_offset = 12, .table_stride = 32}; program.descriptor_sources[1].dword_count = 4; ImageResource root{}; @@ -39609,6 +39705,11 @@ int main(int argc, char **argv) { vulkan.CheckRenderExecutorStencilBindingDiscovery(); return 0; } + if (argc == 2 && std::strcmp(argv[1], "--native-indirect-only") == 0) { + VulkanHarness vulkan; + vulkan.CheckNativeIndirectDispatch(); + return 0; + } if (argc == 2 && std::strcmp(argv[1], "--layered-image-only") == 0) { VulkanHarness vulkan; vulkan.CheckUnifiedImageViewCache();