diff --git a/core/libs/prx/libSceAgcDriver/CMakeLists.txt b/core/libs/prx/libSceAgcDriver/CMakeLists.txt index 051af9c6e5..a37a6fe02e 100644 --- a/core/libs/prx/libSceAgcDriver/CMakeLists.txt +++ b/core/libs/prx/libSceAgcDriver/CMakeLists.txt @@ -83,6 +83,19 @@ add_custom_command( VERBATIM ) +set(agcDispatchArgumentsShaderSource ${CMAKE_CURRENT_SOURCE_DIR}/Graphics/shaders/DispatchArguments.comp) +set(agcDispatchArgumentsShaderSpv ${agcTextureDetileGeneratedDir}/DispatchArguments.spv) +set(agcDispatchArgumentsShaderHeader ${agcTextureDetileGeneratedDir}/DispatchArguments_spv.h) +add_custom_command( + OUTPUT "${agcDispatchArgumentsShaderHeader}" + COMMAND ${CMAKE_COMMAND} -E make_directory "${agcTextureDetileGeneratedDir}" + COMMAND "$" -V --target-env vulkan1.1 -o "${agcDispatchArgumentsShaderSpv}" "${agcDispatchArgumentsShaderSource}" + COMMAND ${CMAKE_COMMAND} -DINPUT=${agcDispatchArgumentsShaderSpv} -DOUTPUT=${agcDispatchArgumentsShaderHeader} + -DSYMBOL=DISPATCH_ARGUMENTS_SPV -P "${CMAKE_CURRENT_SOURCE_DIR}/embed_spirv.cmake" + DEPENDS "${agcDispatchArgumentsShaderSource}" glslang-standalone + VERBATIM +) + set(agcSampleCounterShaderSource ${CMAKE_CURRENT_SOURCE_DIR}/Graphics/shaders/SampleCounter.comp) set(agcSampleCounterShaderSpv ${agcTextureDetileGeneratedDir}/SampleCounter.spv) set(agcSampleCounterShaderHeader ${agcTextureDetileGeneratedDir}/SampleCounter_spv.h) @@ -230,6 +243,7 @@ set(agcDriverSources ${agcCmaskClearShaderHeader} ${agcMeshArgumentsShaderHeader} ${agcTextureDetileImageShaderHeader} + ${agcDispatchArgumentsShaderHeader} ${agcSampleCounterShaderHeader} ${agcSampleDumpsShaderHeader} ${agcPresentationShaderHeaders} diff --git a/core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp b/core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp index bb39fa7beb..5737f40d74 100644 --- a/core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp +++ b/core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp @@ -4364,6 +4364,16 @@ void VulkanDevice::recordDispatch(RecordedDispatch& record) { recorder.EndGpuTiming(timing); ++d.preBarriersRecorded; } + VkBuffer indirectBuffer = argumentImport != nullptr ? argumentImport->buffer : VK_NULL_HANDLE; + VkDeviceSize indirectOffset = argumentImport != nullptr ? arguments - argumentImport->base : 0; + if (argumentImport != nullptr && argumentImport->address != 0) { + auto checked = std::make_shared(context, 12, VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT); + if (recorder.RecordDispatchArguments(commands, argumentImport->address + (arguments - argumentImport->base), checked->DeviceAddress())) { + indirectBuffer = checked->Handle(); + indirectOffset = 0; + recorder.Keep(std::move(checked)); + } + } context.Resolved(&Graphics::DeviceFunctions::cmdBindPipeline, "vkCmdBindPipeline")(commands, VK_PIPELINE_BIND_POINT_COMPUTE, record.objects->pipeline); resources.Bind(commands, VK_PIPELINE_BIND_POINT_COMPUTE, record.objects->layout); if (const auto dwords = record.shader->program->workgroupMemoryDwords; dwords != 0) { @@ -4377,7 +4387,7 @@ void VulkanDevice::recordDispatch(RecordedDispatch& record) { const auto gpuTiming = recorder.BeginGpuTiming(record.programAddress != 0 ? record.programAddress : record.shader->program->variantId); if (state->checkpointMarker != nullptr) state->checkpointMarker(commands, reinterpret_cast(static_cast(record.programAddress))); if (argumentImport != nullptr) { - context.Resolved(&Graphics::DeviceFunctions::cmdDispatchIndirect, "vkCmdDispatchIndirect")(commands, argumentImport->buffer, arguments - argumentImport->base); + context.Resolved(&Graphics::DeviceFunctions::cmdDispatchIndirect, "vkCmdDispatchIndirect")(commands, indirectBuffer, indirectOffset); } else { const auto* limit = state->properties.limits.maxComputeWorkGroupCount; if (record.x <= limit[0] && record.y <= limit[1] && record.z <= limit[2]) { diff --git a/core/libs/prx/libSceAgcDriver/Graphics/include/Recorder.hpp b/core/libs/prx/libSceAgcDriver/Graphics/include/Recorder.hpp index 2470ab5323..8e9f4fc786 100644 --- a/core/libs/prx/libSceAgcDriver/Graphics/include/Recorder.hpp +++ b/core/libs/prx/libSceAgcDriver/Graphics/include/Recorder.hpp @@ -189,6 +189,7 @@ public: // ever submitted); WaitSerial(serial) then covers all recorded work. std::uint64_t SubmitAndEpoch(); bool HasTimeline() const { return timeline != VK_NULL_HANDLE; } + bool RecordDispatchArguments(VkCommandBuffer commands, VkDeviceAddress source, VkDeviceAddress arguments); // Waits until every batch up to `serial` completed on the GPU. Called WITHOUT GuestMemory::GpuMutex: // it touches only the immutable timeline semaphore (the caller keeps the device alive). Requires // HasTimeline(); the debug switch APS5_NO_TIMELINE makes the device create the recorder without one. @@ -703,6 +704,9 @@ private: VkPipelineLayout dumpLayout = VK_NULL_HANDLE; VkPipeline dumpPipeline = VK_NULL_HANDLE; std::shared_ptr segmentPools; + int dispatchArgumentState = 0; + VkPipelineLayout dispatchArgumentLayout = VK_NULL_HANDLE; + VkPipeline dispatchArgumentPipeline = VK_NULL_HANDLE; // BeginGpuTiming on the open batch without Commands() (RecordStore times its own run, which // Commands() would close). std::uint32_t beginTiming(std::uint64_t key); diff --git a/core/libs/prx/libSceAgcDriver/Graphics/shaders/DispatchArguments.comp b/core/libs/prx/libSceAgcDriver/Graphics/shaders/DispatchArguments.comp new file mode 100644 index 0000000000..ed0b252561 --- /dev/null +++ b/core/libs/prx/libSceAgcDriver/Graphics/shaders/DispatchArguments.comp @@ -0,0 +1,23 @@ +#version 450 +#extension GL_EXT_buffer_reference : require + +layout(local_size_x = 1) in; +layout(buffer_reference, std430, buffer_reference_align = 4) buffer Words { uint words[]; }; +layout(push_constant) uniform Parameters { + Words source; + Words arguments; + uint maxX; + uint maxY; + uint maxZ; + uint unused; +} parameters; + +void main() { + uint x = parameters.source.words[0]; + uint y = parameters.source.words[1]; + uint z = parameters.source.words[2]; + bool valid = x <= parameters.maxX && y <= parameters.maxY && z <= parameters.maxZ; + parameters.arguments.words[0] = valid ? x : 0u; + parameters.arguments.words[1] = valid ? y : 0u; + parameters.arguments.words[2] = valid ? z : 0u; +} diff --git a/core/libs/prx/libSceAgcDriver/Graphics/src/Recorder.cpp b/core/libs/prx/libSceAgcDriver/Graphics/src/Recorder.cpp index b494d51ffe..467c0fc27f 100644 --- a/core/libs/prx/libSceAgcDriver/Graphics/src/Recorder.cpp +++ b/core/libs/prx/libSceAgcDriver/Graphics/src/Recorder.cpp @@ -13,6 +13,7 @@ #include "prx/libSceAgcDriver/Execution/include/Pm4Opcodes.hpp" #include "prx/libc/include/CpuTopology.hpp" #include "prx/libc/include/GuestAllocations.hpp" +#include "prx/libSceAgcDriver/Graphics/shaders/DispatchArguments_spv.h" #include #include #include @@ -1168,6 +1169,8 @@ Recorder::~Recorder() { context.Function("vkDestroyPipelineLayout")(context.device, dumpLayout, nullptr); sampleCounter.reset(); } + if (dispatchArgumentPipeline != VK_NULL_HANDLE) context.Function("vkDestroyPipeline")(context.device, dispatchArgumentPipeline, nullptr); + if (dispatchArgumentLayout != VK_NULL_HANDLE) context.Function("vkDestroyPipelineLayout")(context.device, dispatchArgumentLayout, nullptr); for (auto& [commands, fence] : spare) { context.Function("vkFreeCommandBuffers")(context.device, context.pool, 1, &commands); context.Function("vkDestroyFence")(context.device, fence, nullptr); @@ -1946,6 +1949,52 @@ void Recorder::EndGpuTiming(std::uint32_t index, std::uint64_t bytes) { context.Function("vkCmdWriteTimestamp")(open->commands, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT, open->queries, index * 2 + 1); } +bool Recorder::RecordDispatchArguments(VkCommandBuffer commands, VkDeviceAddress source, VkDeviceAddress arguments) { + constexpr std::uint32_t PushBytes = 32; + if (dispatchArgumentState == 0) { + dispatchArgumentState = -1; + VkShaderModule module = VK_NULL_HANDLE; + try { + const VkPushConstantRange push{VK_SHADER_STAGE_COMPUTE_BIT, 0, PushBytes}; + VkPipelineLayoutCreateInfo layoutInfo{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; + layoutInfo.pushConstantRangeCount = 1; + layoutInfo.pPushConstantRanges = &push; + Check(context.Function("vkCreatePipelineLayout")(context.device, &layoutInfo, nullptr, &dispatchArgumentLayout), "vkCreatePipelineLayout dispatch arguments"); + VkShaderModuleCreateInfo moduleInfo{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO}; + moduleInfo.codeSize = sizeof(DISPATCH_ARGUMENTS_SPV); + moduleInfo.pCode = DISPATCH_ARGUMENTS_SPV; + Check(context.Function("vkCreateShaderModule")(context.device, &moduleInfo, nullptr, &module), "vkCreateShaderModule dispatch arguments"); + VkComputePipelineCreateInfo pipelineInfo{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO}; + pipelineInfo.stage = {VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO}; + pipelineInfo.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT; + pipelineInfo.stage.module = module; + pipelineInfo.stage.pName = "main"; + pipelineInfo.layout = dispatchArgumentLayout; + Check(context.Function("vkCreateComputePipelines")(context.device, context.pipelineCache, 1, &pipelineInfo, nullptr, &dispatchArgumentPipeline), "vkCreateComputePipelines dispatch arguments"); + context.Function("vkDestroyShaderModule")(context.device, module, nullptr); + dispatchArgumentState = 1; + } catch (const std::exception& error) { + if (module != VK_NULL_HANDLE) context.Function("vkDestroyShaderModule")(context.device, module, nullptr); + std::fprintf(stderr, "[gpu] indirect dispatches read their group counts unchecked: %s\n", error.what()); + } + } + if (dispatchArgumentState < 0) return false; + recordBarrier(commands, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT | VK_PIPELINE_STAGE_HOST_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT | VK_ACCESS_HOST_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + function(cmdBindPipeline, "vkCmdBindPipeline")(commands, VK_PIPELINE_BIND_POINT_COMPUTE, dispatchArgumentPipeline); + const auto* limit = context.limits.maxComputeWorkGroupCount; + struct { + VkDeviceAddress source; + VkDeviceAddress arguments; + std::array limits; + } parameters{source, arguments, {limit[0], limit[1], limit[2], 0u}}; + static_assert(sizeof(parameters) == PushBytes); + function(cmdPushConstants, "vkCmdPushConstants")(commands, dispatchArgumentLayout, VK_SHADER_STAGE_COMPUTE_BIT, 0, PushBytes, ¶meters); + function(cmdDispatch, "vkCmdDispatch")(commands, 1, 1, 1); + recordBarrier(commands, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_DRAW_INDIRECT_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_INDIRECT_COMMAND_READ_BIT); + CountBarriers(CommandClass::DispatchLeading, 2); + return true; +} + void Recorder::CountSamples() { GuestMemory::AssertGpuLockHeld("Recorder::CountSamples"); gpuSampleCounter(); diff --git a/core/libs/prx/libSceAgcDriver/tests/execution/IndirectDispatchLimit.cpp b/core/libs/prx/libSceAgcDriver/tests/execution/IndirectDispatchLimit.cpp new file mode 100644 index 0000000000..5b4420522e --- /dev/null +++ b/core/libs/prx/libSceAgcDriver/tests/execution/IndirectDispatchLimit.cpp @@ -0,0 +1,88 @@ +#include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp" +#include "prx/libSceAgcDriver/Graphics/include/Draw.hpp" +#include "prx/libc/include/GuestAllocations.hpp" +#include "Recompiler.hpp" +#include "VulkanTestDevice.hpp" +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using AgcDriver::Graphics::Require; +using ShaderRecompiler::ShaderStage; + +constexpr std::uint32_t Threads = 64; +constexpr std::size_t BlockBytes = 65536; +alignas(256) std::array Output{}; + +alignas(256) constexpr std::array StoreCode{0x34020082u, 0x7e040287u, 0xe0701000u, 0x80010201u, 0xbf810000u}; + +std::array BufferDescriptor(const void* data, std::uint32_t count) { + const auto address = reinterpret_cast(data); + return {static_cast(address), static_cast((address >> 32u) & 0xffffu), count * 4u, 0x31016facu}; +} + +int Run(AgcDriver::VulkanDevice& device, std::uint32_t* block, std::array groups) { + Output.fill(0xdeadbeefu); + std::memcpy(block, groups.data(), sizeof(groups)); + std::vector userData(8, 0u); + const auto output = BufferDescriptor(Output.data(), Threads); + std::copy(output.begin(), output.end(), userData.begin() + 4); + const std::span code(StoreCode); + const std::array memory{{{reinterpret_cast(code.data()), std::as_bytes(code)}}}; + const ShaderRecompiler::ShaderComputeStageInfo compute{{Threads, 1, 1}, 0, {false, false, false}, false, 1}; + ShaderRecompiler::RecompileRequest request{ + {ShaderStage::Compute, reinterpret_cast(code.data()), code, 0, {}}, + {64, 0, userData, compute, std::nullopt, std::nullopt, memory}, + device.Target(), + {0, 0, 0, 128} + }; + request.useCache = false; + const auto result = ShaderRecompiler::Recompile(request); + const auto outcome = device.DispatchIndirect(result, reinterpret_cast(block), {}, reinterpret_cast(code.data())); + device.WaitIdle(); + return outcome.cpuReason; +} + +} + +int main() { + try { + const auto device = OpenVulkanTestDevice(); + if (!device) return VulkanTestSkipped; + auto* block = static_cast(::operator new(BlockBytes, std::align_val_t{BlockBytes})); + std::memset(block, 0, BlockBytes); + auto* mutation = GuestAllocations::GuestAllocationsBegin_nid_postfix(); + GuestAllocations::GuestAllocationsAdd_nid_postfix(mutation, block, BlockBytes, true, true, true); + GuestAllocations::GuestAllocationsEnd_nid_postfix(mutation); + int failures = 0; + try { + const int invalidPath = Run(*device, block, {0xffffffffu, 1u, 1u}); + if (invalidPath != 0) { + std::printf("skipped, the group counts were read on the CPU (reason %d): the device cannot import the block\n", invalidPath); + } else { + for (std::uint32_t tid = 0; tid < Threads; ++tid) Require(Output[tid] == 0xdeadbeefu, "indirect dispatch limit: thread " + std::to_string(tid) + " ran with 0xffffffff groups"); + Require(Run(*device, block, {1u, 1u, 1u}) == 0, "indirect dispatch limit: valid counts left the GPU path"); + for (std::uint32_t tid = 0; tid < Threads; ++tid) Require(Output[tid] == 7u, "indirect dispatch limit: thread " + std::to_string(tid) + " stored " + std::to_string(Output[tid]) + " with one group, expected 7"); + } + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + failures = 1; + } + mutation = GuestAllocations::GuestAllocationsBegin_nid_postfix(); + GuestAllocations::GuestAllocationsRemove_nid_postfix(mutation, block); + GuestAllocations::GuestAllocationsEnd_nid_postfix(mutation); + ::operator delete(block, std::align_val_t{BlockBytes}); + if (failures != 0) return 1; + std::puts("indirect dispatch limit tests passed"); + return 0; + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +}