From e4aca7a60b4b90f86d868b8d83597797768da8fe Mon Sep 17 00:00:00 2001 From: nmzik Date: Sun, 4 Oct 2026 01:15:39 +0400 Subject: [PATCH] shader: initialize EXEC from requested dispatch thread extents Co-authored-by: TheCruZ Co-authored-by: Mac --- .../host_gpu/renderer/renderCompute.cpp | 4 + .../recompiler/backend/spirv/SpirvEmitter.cpp | 6 +- .../backend/spirv/spirvEmitterFlow.cpp | 8 ++ .../backend/spirv/spirvEmitterInstructions.h | 1 + .../frontend/translate/Translate.cpp | 10 ++ src/graphics/shader/recompiler/ir/ShaderIR.h | 1 + .../recompiler/ir/opcodes/ValueOpcodes.inc | 1 + .../recompiler/ir/passes/BindingLayout.cpp | 21 ++-- tests/ShaderRecompilerComputeTests.cpp | 107 +++++++++++++++++- 9 files changed, 147 insertions(+), 12 deletions(-) diff --git a/src/graphics/host_gpu/renderer/renderCompute.cpp b/src/graphics/host_gpu/renderer/renderCompute.cpp index e7e2f0ada..4f9c0673e 100644 --- a/src/graphics/host_gpu/renderer/renderCompute.cpp +++ b/src/graphics/host_gpu/renderer/renderCompute.cpp @@ -302,6 +302,10 @@ void RenderExecutor::DispatchDirect(uint64_t submit_id, CommandBuffer& buffer, m_context.GetPipelineCache().GetComputePipeline(input_info, compute_program); auto& bindings = m_compute_bindings; PrepareBindings(input_info.stage, bindings); + if (program.bindings.dispatch_thread_dword != ShaderRecompiler::IR::PushData::NoStart) { + std::copy(std::begin(input_info.dispatch_threads_num), std::end(input_info.dispatch_threads_num), + bindings.shader_data.begin() + program.bindings.dispatch_thread_dword); + } FindBuffers(bindings); if (program.info.uses_dma) { m_context.PrepareBda(); diff --git a/src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp b/src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp index ab6e666e1..67f9d6cb8 100644 --- a/src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp @@ -85,9 +85,13 @@ void ValidateNativeProgram(const IR::Program& program) { } const auto has_shader_data_storage = present[static_cast(Kind::ShaderData)]; const auto shader_data_dwords = program.bindings.ShaderDataDwords(); + const auto user_data_dwords = program.bindings.user_data_registers.size(); + const bool has_dispatch_threads = program.bindings.dispatch_thread_dword != IR::PushData::NoStart; if ((program.bindings.UsesPushData() && !IR::PushData::CanFit(program.bindings.push_data_start_dword, shader_data_dwords)) || - program.bindings.memory_offset_dword != program.bindings.user_data_registers.size() || + (has_dispatch_threads && (program.stage != ShaderType::Compute || + program.bindings.dispatch_thread_dword != user_data_dwords)) || + program.bindings.memory_offset_dword != user_data_dwords + (has_dispatch_threads ? 3u : 0u) || program.bindings.memory_offset_count != buffers.size() || has_shader_data_storage != (shader_data_dwords != 0 && !program.bindings.UsesPushData()) || !std::is_sorted(program.bindings.user_data_registers.begin(), diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp index 68d23aea5..a7b167e47 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp @@ -631,6 +631,14 @@ uint32_t EmitGetBuiltin(ValueEmitContext& ctx, IR::Value kind, IR::Value index) return EmitBuiltinU32(ctx.state, static_cast(kind.U32()), index.U32()); } +uint32_t EmitGetDispatchThreadExtent(ValueEmitContext& ctx, const IR::Inst& inst) { + const auto start = ctx.state.program.bindings.dispatch_thread_dword; + if (start == IR::PushData::NoStart || !inst.Arg(0).IsImmediate() || inst.Arg(0).U32() >= 3u) { + ctx.Fail(inst, "invalid dispatch thread extent"); + } + return EmitShaderDataDwordLoad(ctx.state, start + inst.Arg(0).U32()); +} + uint32_t EmitUndefU1(EmitterState& state, const IR::Inst& inst) { const auto result = state.builder.AllocateId(); state.builder.AddFunction(spv::OpUndef, TypeId(state, inst.GetType()), result); diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterInstructions.h b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterInstructions.h index 35a7b7813..4a2618df7 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterInstructions.h +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterInstructions.h @@ -198,6 +198,7 @@ uint32_t EmitMeshDrawParameter(ValueEmitContext& ctx, const IR::Ins uint32_t EmitGetTessellationAttribute(ValueEmitContext& ctx, const IR::Inst& inst); void EmitSetTessellationAttribute(ValueEmitContext& ctx, const IR::Inst& inst); uint32_t EmitGetUserData(EmitterState& state, IR::ScalarReg reg); +uint32_t EmitGetDispatchThreadExtent(ValueEmitContext& ctx, const IR::Inst& inst); uint32_t EmitGetBuiltin(ValueEmitContext& ctx, IR::Value kind, IR::Value index); uint32_t EmitUndefU1(EmitterState& state, const IR::Inst& inst); inline constexpr auto EmitUndefU8 = EmitUndefU1; diff --git a/src/graphics/shader/recompiler/frontend/translate/Translate.cpp b/src/graphics/shader/recompiler/frontend/translate/Translate.cpp index b06c3a6f3..ea14959e0 100644 --- a/src/graphics/shader/recompiler/frontend/translate/Translate.cpp +++ b/src/graphics/shader/recompiler/frontend/translate/Translate.cpp @@ -1104,6 +1104,16 @@ IR::Program TranslateProgram(const Decoder::Program& decoded, const CFG::Graph& IR::U32(IR::Value(total_threads))); } } + if (options.stage == ShaderType::Compute && + options.input_info.compute->dispatch_thread_dimensions) { + for (uint32_t axis = 0; axis < 3u; axis++) { + const auto extent = IR::U32(entry_ir.Emit( + IR::ValueOpcode::GetDispatchThreadExtent, {IR::Value(axis)})); + initial_exec = entry_ir.LogicalAnd( + initial_exec, + entry_ir.ULessThan(builtin(IR::StageInputKind::GlobalInvocationId, axis), extent)); + } + } entry_ir.SetExec(initial_exec); const auto initial_mask = entry_ir.Emit(IR::ValueOpcode::Ballot, {initial_exec}); entry_ir.SetExecLo(entry_ir.CompositeExtract(initial_mask, 0)); diff --git a/src/graphics/shader/recompiler/ir/ShaderIR.h b/src/graphics/shader/recompiler/ir/ShaderIR.h index 73bade90f..5df80941d 100644 --- a/src/graphics/shader/recompiler/ir/ShaderIR.h +++ b/src/graphics/shader/recompiler/ir/ShaderIR.h @@ -426,6 +426,7 @@ struct DescriptorBinding { struct BindingLayout { uint32_t push_data_start_dword = PushData::NoStart; + uint32_t dispatch_thread_dword = PushData::NoStart; uint32_t memory_offset_dword = 0; uint32_t memory_offset_count = 0; std::vector user_data_registers; diff --git a/src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.inc b/src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.inc index a50d2695f..0dc701989 100644 --- a/src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.inc +++ b/src/graphics/shader/recompiler/ir/opcodes/ValueOpcodes.inc @@ -8,6 +8,7 @@ VALUE_OPCODE(ConditionRef, U1, U1) // Translation and register-state pseudo operations. VALUE_OPCODE(GetUserData, U32, ScalarReg) +VALUE_OPCODE(GetDispatchThreadExtent, U32, U32) VALUE_OPCODE(GetShaderBase, U64) VALUE_OPCODE(MeshDrawParameter, U32, U32) VALUE_OPCODE(MeshAllocate, Void, U32) diff --git a/src/graphics/shader/recompiler/ir/passes/BindingLayout.cpp b/src/graphics/shader/recompiler/ir/passes/BindingLayout.cpp index 08aaed1f1..38dc23910 100644 --- a/src/graphics/shader/recompiler/ir/passes/BindingLayout.cpp +++ b/src/graphics/shader/recompiler/ir/passes/BindingLayout.cpp @@ -15,11 +15,16 @@ namespace { std::abort(); } -std::vector CollectUserData(const Program& program) { +void CollectShaderData(const Program& program, BindingLayout& layout) { std::array registers {}; + bool uses_dispatch_threads = false; for (const auto* block: program.blocks) { for (const auto& inst: *block) { - if (inst.GetOpcode() != ValueOpcode::GetUserData || !inst.HasUses()) { + if (!inst.HasUses()) { + continue; + } + uses_dispatch_threads |= inst.GetOpcode() == ValueOpcode::GetDispatchThreadExtent; + if (inst.GetOpcode() != ValueOpcode::GetUserData) { continue; } if (inst.Arg(0).GetType() != Type::ScalarReg) { @@ -32,13 +37,16 @@ std::vector CollectUserData(const Program& program) { registers[index] = true; } } - std::vector result; for (uint32_t index = 0; index < registers.size(); index++) { if (registers[index]) { - result.push_back(index); + layout.user_data_registers.push_back(index); } } - return result; + layout.memory_offset_dword = static_cast(layout.user_data_registers.size()); + if (uses_dispatch_threads) { + layout.dispatch_thread_dword = layout.memory_offset_dword; + layout.memory_offset_dword += 3u; + } } void AddBinding(BindingLayout& layout, DescriptorBindingKind kind, @@ -109,8 +117,7 @@ void AllocateBindings(Program& program, uint32_t push_data_start_dword) { BindingLayout next; std::vector buffers; const bool uses_gds = CollectMemoryResources(program, buffers); - next.user_data_registers = CollectUserData(program); - next.memory_offset_dword = static_cast(next.user_data_registers.size()); + CollectShaderData(program, next); next.memory_offset_count = static_cast(buffers.size()); next.push_data_start_dword = PushData::StartFor(push_data_start_dword, next.ShaderDataDwords()); diff --git a/tests/ShaderRecompilerComputeTests.cpp b/tests/ShaderRecompilerComputeTests.cpp index 7ae5d77ab..173b2a1ae 100644 --- a/tests/ShaderRecompilerComputeTests.cpp +++ b/tests/ShaderRecompilerComputeTests.cpp @@ -14152,6 +14152,13 @@ public: vk::Sampler sampler = nullptr) { using Kind = ShaderRecompiler::IR::DescriptorBindingKind; const auto &layout = compiled.program.bindings; + auto shader_data = compiled.packed_user_data; + if (layout.dispatch_thread_dword != ShaderRecompiler::IR::PushData::NoStart) { + Require(test.name, "dispatch dimensions", test.has_compute_info, + "runtime thread extents are required by the compiled shader"); + std::copy_n(test.compute_info.dispatch_threads_num, 3, + shader_data.begin() + layout.dispatch_thread_dword); + } auto Binding = [&](Kind kind) { return ShaderRecompiler::IR::FindBinding(layout, kind); }; @@ -14378,8 +14385,7 @@ public: } if (const auto *user = Binding(Kind::ShaderData); user != nullptr) { user_data_buffer = - CreateStorageBuffer(test.name, compiled.packed_user_data, - compiled.packed_user_data.size()); + CreateStorageBuffer(test.name, shader_data, shader_data.size()); user_data_info = {user_data_buffer.buffer, 0, user_data_buffer.size}; vk::WriteDescriptorSet write{}; write.sType = vk::StructureType::eWriteDescriptorSet; @@ -14574,8 +14580,7 @@ public: 1, &descriptor_set, 0, nullptr); if (layout.UsesPushData()) { ShaderRecompiler::IR::PushData push_data; - std::copy(compiled.packed_user_data.begin(), - compiled.packed_user_data.end(), + std::copy(shader_data.begin(), shader_data.end(), push_data.dwords.begin() + layout.push_data_start_dword); cmd.pushConstants(pipeline_layout, vk::ShaderStageFlagBits::eCompute, 0, sizeof(push_data), push_data.dwords.data()); @@ -30601,6 +30606,94 @@ TestCase ImageD16StoreUnpacksHalfPairs() { return test; } +void CheckComputeThreadDimensions(VulkanHarness &vulkan) { + using O = ShaderOpcode; + constexpr u32 grid_size = 16 * 8 * 4; + const std::array extents = {std::array{9u, 5u, 3u}, std::array{13u, 7u, 4u}, + std::array{3u, 2u, 1u}, std::array{16u, 8u, 4u}}; + for (const u32 wave_size : {32u, 64u}) { + for (const bool thread_dimensions : {false, true}) { + TestCase test; + test.name = thread_dimensions ? (wave_size == 32 ? "ThreadDimensions32" : "ThreadDimensions64") + : (wave_size == 32 ? "GroupDimensions32" : "GroupDimensions64"); + auto &code = test.code; + code.push_back(EncodeSop1(0x04, 20, 126)); // Save the SPI-provided initial EXEC. + AppendSMovLiteral(&code, 24, grid_size * 4); + for (u32 axis = 0; axis < 3; ++axis) { + code.push_back(EncodeVop1(0x01, 3 + axis, 4 + axis)); + code.push_back(EncodeVop2(0x1a, 3 + axis, InlineU32(3 - axis), 3 + axis)); + code.push_back(EncodeVop2(0x25, 3 + axis, Vgpr(axis), 3 + axis)); + } + code.push_back(EncodeVop2(0x1a, 4, InlineU32(4), 4)); + code.push_back(EncodeVop2(0x1a, 5, InlineU32(7), 5)); + code.push_back(EncodeVop2(0x25, 3, Vgpr(4), 3)); + code.push_back(EncodeVop2(0x25, 3, Vgpr(5), 3)); + code.push_back(EncodeVop2(0x25, 20, InlineU32(1), 3)); + code.push_back(EncodeVop2(0x1a, 20, InlineU32(2), 20)); + AppendVMovU32(&code, 6, 1); + AppendVMovU32(&code, 7, 0); + AppendBufferStoreDword(&code, 6, 20); + AppendBufferStoreOpcode(&code, 0x32, 6, 7); + code.push_back(EncodeSop2(0x0f, 126, 20, InlineU32(1))); + AppendVMovU32(&code, 6, 2); // Divergence changes only lane zero of each wave. + code.push_back(EncodeSop1(0x04, 126, 20)); + code.push_back(EncodeVop2(0x25, 20, 24, 20)); + AppendBufferStoreDword(&code, 6, 20); + AppendVMovU32(&code, 8, 1); + AppendBufferStoreOpcode(&code, 0x32, 8, 7); + AppendEnd(&code); + test.initial.resize(1 + 2 * grid_size); + test.opcodes = {O::S_MOV_B64, O::S_MOV_B32, O::S_AND_B64, O::V_MOV_B32, + O::V_LSHLREV_B32, O::V_ADD_NC_U32, O::BUFFER_STORE_DWORD, + O::BUFFER_ATOMIC_ADD, O::S_ENDPGM}; + test.compute_info.wave_size = wave_size; + test.compute_info.host_subgroup_size = vulkan.SubgroupSize(); + test.compute_info.threads_num[0] = 8; + test.compute_info.threads_num[1] = 4; + test.compute_info.threads_num[2] = 2; + std::fill_n(test.compute_info.group_id, 3, true); + test.compute_info.thread_ids_num = 3; + test.compute_info.workgroup_register = 4; + test.compute_info.dispatch_thread_dimensions = thread_dimensions; + test.has_compute_info = true; + test.expand_shader_data_storage = thread_dimensions && wave_size == 32; + const auto compiled = CompileCase(test, vulkan.SubgroupSize()); + std::vector compiled_key; + BuildStageStaticKey(test.compute_info, compiled_key); + for (const auto &extent : extents) { + std::copy(extent.begin(), extent.end(), test.compute_info.dispatch_threads_num); + std::vector runtime_key; + BuildStageStaticKey(test.compute_info, runtime_key); + Require(test.name, "runtime shader reuse", compiled_key == runtime_key, + "dispatch thread counts changed the compiled shader cache key"); + test.dispatch_x = thread_dimensions ? (extent[0] + 7) / 8 : 2; + test.dispatch_y = thread_dimensions ? (extent[1] + 3) / 4 : 2; + test.dispatch_z = thread_dimensions ? (extent[2] + 1) / 2 : 2; + test.expected.assign(test.initial.size(), 0); + for (u32 z = 0; z < 4; ++z) { + for (u32 y = 0; y < 8; ++y) { + for (u32 x = 0; x < 16; ++x) { + if (thread_dimensions && (x >= extent[0] || y >= extent[1] || z >= extent[2])) continue; + const auto index = x + 16 * y + 128 * z; + const auto local_index = x % 8 + 8 * (y % 4) + 32 * (z % 2); + test.expected[0] += 2; + test.expected[1 + index] = 1; + test.expected[1 + grid_size + index] = local_index % wave_size == 0 ? 2 : 1; + } + } + } + auto output = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size()); + vulkan.Dispatch(test, compiled, output); + const auto actual = vulkan.ReadBuffer(test.name, output, test.expected.size()); + vulkan.DestroyBuffer(&output); + CompareWords(test, "initial EXEC, restored EXEC and padded-lane atomics", test.expected, actual); + if (!thread_dimensions) break; + } + std::printf("[compute] %-32s ok\n", test.name); + } + } +} + void CheckIndirectBufferStore(VulkanHarness &vulkan) { using namespace ShaderRecompiler::IR; TestCase test; @@ -37246,6 +37339,11 @@ int main(int argc, char **argv) { std::setvbuf(stdout, nullptr, _IONBF, 0); EnsureConfigInitialized(); CheckLeastRecentlyUsedCacheOrdering(); + if (argc == 2 && std::strcmp(argv[1], "--thread-dimensions-only") == 0) { + VulkanHarness vulkan; + CheckComputeThreadDimensions(vulkan); + return 0; + } if (argc == 2 && std::strcmp(argv[1], "--lds-limit-only") == 0) { VulkanHarness vulkan; CheckComputeLdsLimit(vulkan); @@ -38006,6 +38104,7 @@ int main(int argc, char **argv) { CheckPixelParameterAliases(); CheckRectListShaders(); CheckIndirectBufferStore(vulkan); + CheckComputeThreadDimensions(vulkan); CheckIndirectImageKeySwitch(); CheckWave64WholeWaveResults(); CheckPs5GameExampleImageClearRuntimeShape();