shifathosn
/
AnyPS5
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #3300 from lendinglibrary/fix/3212-dispatch-split 

fix(agc): split a compute dispatch over the workgroup count limit
This commit is contained in:
Adrià 2026-10-11 13:46:20 +02:00 committed by GitHub
parent cafc074e1a 288ee89d2d
commit e38b5e95d2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 236 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
core/libs/prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp
View File

@ -37,6 +37,8 @@ public:
std::string DeviceName() const;
ShaderRecompiler::SpirvTarget Target() const;
std::optional<ShaderRecompiler::GeometryStageLimits> GeometryLimits() const;
const std::uint32_t* WorkGroupCountLimit() const;
std::uint32_t VendorId() const;
ShaderRecompiler::SpirvTarget ComputeTarget(std::uint32_t waveSize) const;
// Distinguishes this device from every earlier one in the process (a memo keyed by the device
// cannot rely on the pointer, which a replacement may reuse).

62
core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp
View File

@ -69,12 +69,6 @@ void require(bool condition, const char* reason) {
if (!condition) throw std::runtime_error(std::string("Vulkan presentation: ") + reason);
}
std::string workgroupLimitError(std::uint32_t x, std::uint32_t y, std::uint32_t z, const std::uint32_t* limit) {
char text[160];
std::snprintf(text, sizeof(text), "Vulkan dispatch: workgroup count %ux%ux%u exceeds device limits %ux%ux%u", x, y, z, limit[0], limit[1], limit[2]);
return text;
}
// The ShaderResources content cache dispatches share with recorded draws: Graphics::ResourceCache,
// one process-wide instance (see SharedResourceCache) that the State references so this file keeps
// its Find/Insert/Remove/Clear calls; the device clears it at teardown before its descriptor caches
@ -2650,6 +2644,14 @@ std::string VulkanDevice::DeviceName() const {
return state->properties.deviceName;
}
const std::uint32_t* VulkanDevice::WorkGroupCountLimit() const {
return state->properties.limits.maxComputeWorkGroupCount;
}
std::uint32_t VulkanDevice::VendorId() const {
return state->properties.vendorID;
}
ShaderRecompiler::SpirvTarget VulkanDevice::Target() const {
require(state->shaderProfile != nullptr, "shader device profile is unavailable");
return state->shaderProfile->Target();
@ -3421,7 +3423,6 @@ void VulkanDevice::decideIndirect(RecordedDispatch& record, IndirectOutcome& out
auto& timer = *record.timer;
const auto& context = *record.context;
const auto arguments = record.arguments;
const auto* limit = state->properties.limits.maxComputeWorkGroupCount;
// Decided here, after the resource build (which may retire imports and wait for recorded
// work) and before the dispatch itself is recorded, so the import cannot be dropped before
// the record. The build already recorded the copy-ins of its misaligned and staged regions
@ -3477,11 +3478,6 @@ void VulkanDevice::decideIndirect(RecordedDispatch& record, IndirectOutcome& out
record.arguments = 0;
timer.indirect = false;
std::snprintf(groupsText, 40, "%ux%ux%u", record.x, record.y, record.z);
if (record.x > limit[0] || record.y > limit[1] || record.z > limit[2]) {
char where[96];
std::snprintf(where, sizeof(where), " (indirect arguments at 0x%llx read by the CPU, reason %d)", static_cast<unsigned long long>(arguments), outcome.cpuReason);
throw std::runtime_error(workgroupLimitError(record.x, record.y, record.z, limit) + where);
}
}
void VulkanDevice::recordDispatch(RecordedDispatch& record) {
@ -3583,8 +3579,36 @@ void VulkanDevice::recordDispatch(RecordedDispatch& record) {
context.Resolved(&Graphics::DeviceFunctions::cmdPushConstants, "vkCmdPushConstants")(commands, record.objects->layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, Graphics::ComputePushConstantBytes, record.pushBytes->data());
}
const auto gpuTiming = recorder.BeginGpuTiming(record.programAddress != 0 ? record.programAddress : record.shader->program->variantId);
if (argumentImport != nullptr) context.Resolved(&Graphics::DeviceFunctions::cmdDispatchIndirect, "vkCmdDispatchIndirect")(commands, argumentImport->buffer, arguments - argumentImport->base);
else context.Resolved(&Graphics::DeviceFunctions::cmdDispatch, "vkCmdDispatch")(commands, record.x, record.y, record.z);
if (argumentImport != nullptr) {
context.Resolved(&Graphics::DeviceFunctions::cmdDispatchIndirect, "vkCmdDispatchIndirect")(commands, argumentImport->buffer, arguments - argumentImport->base);
} else {
const auto* limit = state->properties.limits.maxComputeWorkGroupCount;
if (record.x <= limit[0] && record.y <= limit[1] && record.z <= limit[2]) {
context.Resolved(&Graphics::DeviceFunctions::cmdDispatch, "vkCmdDispatch")(commands, record.x, record.y, record.z);
} else {
if (record.shader->program->workgroupMemoryDwords != 0) {
char text[256];
std::snprintf(text, sizeof(text), "Vulkan dispatch: workgroup count %ux%ux%u exceeds device limits %ux%ux%u and a split aliases the device-memory LDS slots, whose base is computed from gl_NumWorkGroups of a single chunk", record.x, record.y, record.z, limit[0], limit[1], limit[2]);
throw std::runtime_error(text);
}
if (state->properties.vendorID == 0x10de) {
char text[192];
std::snprintf(text, sizeof(text), "Vulkan dispatch: workgroup count %ux%ux%u exceeds device limits %ux%ux%u and NVIDIA drops workgroups of a vkCmdDispatchBase split", record.x, record.y, record.z, limit[0], limit[1], limit[2]);
throw std::runtime_error(text);
}
const auto dispatchBase = context.Function<PFN_vkCmdDispatchBase>("vkCmdDispatchBase");
bool firstChunk = true;
for (std::uint32_t zBase = 0; zBase < record.z; zBase += limit[2]) {
for (std::uint32_t yBase = 0; yBase < record.y; yBase += limit[1]) {
for (std::uint32_t xBase = 0; xBase < record.x; xBase += limit[0]) {
if (!firstChunk) Graphics::RecordMemoryBarrier(context, commands, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT);
firstChunk = false;
dispatchBase(commands, xBase, yBase, zBase, std::min(limit[0], record.x - xBase), std::min(limit[1], record.y - yBase), std::min(limit[2], record.z - zBase));
}
}
}
}
}
recorder.EndGpuTiming(gpuTiming);
const auto trailingTiming = recorder.BeginGpuTiming(CommandClass::DispatchTrailing);
constexpr VkAccessFlags dispatchedAccess = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT | VK_ACCESS_TRANSFER_WRITE_BIT | VK_ACCESS_HOST_READ_BIT;
@ -3635,10 +3659,6 @@ VulkanDevice::IndirectOutcome VulkanDevice::dispatch(const ShaderRecompiler::Rec
}
auto pushBytes = Graphics::AssemblePushConstants(shaders);
const auto context = graphicsContext();
const auto* limit = state->properties.limits.maxComputeWorkGroupCount;
if (arguments == 0 && (x > limit[0] || y > limit[1] || z > limit[2])) {
throw std::runtime_error(workgroupLimitError(x, y, z, limit));
}
// Pipelines are shared by dispatches of one compiled variant; descriptor set layouts built from the
// same bindings are compatible, so the pipeline layout of the first dispatch serves them all.
const std::uint64_t pipelineKey = shader.PipelineVariantId() != 0 ? (shader.PipelineVariantId() << 1u) | (pushStages != 0 ? 1u : 0u) : 0u;
@ -3805,7 +3825,7 @@ VulkanDevice::IndirectOutcome VulkanDevice::dispatch(const ShaderRecompiler::Rec
layoutInfo.pPushConstantRanges = pushStages != 0 ? &push : nullptr;
check(state->DeviceFunction<PFN_vkCreatePipelineLayout>("vkCreatePipelineLayout")(state->device, &layoutInfo, nullptr, &objects->layout), "vkCreatePipelineLayout");
VkComputePipelineCreateInfo pipelineInfo{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO};
pipelineInfo.flags = context.pipelineExecutableInfo ? VK_PIPELINE_CREATE_CAPTURE_STATISTICS_BIT_KHR : 0;
pipelineInfo.flags = (context.pipelineExecutableInfo ? VK_PIPELINE_CREATE_CAPTURE_STATISTICS_BIT_KHR : 0) | VK_PIPELINE_CREATE_DISPATCH_BASE_BIT;
pipelineInfo.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
pipelineInfo.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
pipelineInfo.stage.module = objects->module;
@ -3923,10 +3943,6 @@ RecipeOutcome VulkanDevice::DispatchRecipe(const ShaderRecompiler::RecompileResu
auto& d = Dispatches();
const std::array<Graphics::CompiledShader, 1> shaders{{{ShaderRecompiler::ShaderStage::Compute, &shader, 0}}};
const auto context = graphicsContext();
const auto* limit = state->properties.limits.maxComputeWorkGroupCount;
if (arguments == 0 && (x > limit[0] || y > limit[1] || z > limit[2])) {
throw std::runtime_error(workgroupLimitError(x, y, z, limit));
}
if (recipe.needsCompletion || recipe.holdsLease || !hit->resources->Reusable()) throw std::runtime_error("Vulkan dispatch: recipe over a non-reusable template");
// Without the refresh the template holds the words it was keyed with: a data-only hit takes
// the ordinary path, whose word-keyed lookup serves or builds the object for the live words.

195
core/libs/prx/libSceAgcDriver/tests/execution/OversizedDispatch.cpp Normal file
View File

@ -0,0 +1,195 @@
#include "prx/libSceAgcDriver/Execution/include/Driver.hpp"
#include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
#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"
#ifdef _WIN32
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#endif
#include <array>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <algorithm>
#include <iostream>
#include <span>
#include <string>
#include <vector>
namespace {
using AgcDriver::Graphics::Require;
constexpr std::uint32_t Threads = 64;
constexpr std::uint32_t Period = 4;
constexpr std::uint32_t Groups = 129199u;
constexpr std::uint32_t Stores = Groups * Threads;
alignas(256) constexpr std::array<std::uint32_t, 69> Code{
0xbfa00003u, 0xd7460002u, 0x04010c0au, 0x7da80408u, 0xbf88003fu, 0x7e000c09u, 0xbf070980u, 0x858a807eu,
0x7e005700u, 0x100000ffu, 0x4f800000u, 0x7e060f00u, 0xd5766a00u, 0x02020609u, 0x7d8a0280u, 0x4c020080u,
0x02000101u, 0xd56a0001u, 0x00020700u, 0x4c000303u, 0x4a020303u, 0x02000101u, 0xd56a0000u, 0x00020500u,
0xd5690001u, 0x00020009u, 0x4c060302u, 0x7d860609u, 0x7d8c02f9u, 0x06068c02u, 0x87ea6a0cu, 0x50000080u,
0xd5286a00u, 0x003200c1u, 0xd5010000u, 0x002a00c1u, 0xd5690000u, 0x00020009u, 0x4c000102u, 0x7d0a0080u,
0xbe88246au, 0xbf880015u, 0x7d0a0081u, 0xbe8a246au, 0xbf88000cu, 0x7d0a0082u, 0xbeea246au, 0xbf880003u,
0x7e000207u, 0xe0102000u, 0x80000002u, 0x8afe7e6au, 0xbf880003u, 0x7e000206u, 0xe0102000u, 0x80000002u,
0xbefe046au, 0x8afe7e0au, 0xbf880003u, 0x7e000205u, 0xe0102000u, 0x80000002u, 0xbefe040au, 0x8afe7e08u,
0xbf880003u, 0x7e000204u, 0xe0102000u, 0x80000002u, 0xbf810000u};
constexpr std::array<std::uint32_t, 4> Pattern{0x11111111u, 0x22222222u, 0x33333333u, 0x44444444u};
class GuestBlock {
public:
explicit GuestBlock(std::size_t dwords) : Bytes(dwords * sizeof(std::uint32_t)) {
#ifdef _WIN32
block = static_cast<std::uint8_t*>(VirtualAlloc(nullptr, Bytes, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE));
#else
block = static_cast<std::uint8_t*>(std::aligned_alloc(256, Bytes));
#endif
Require(block != nullptr, "oversized dispatch: cannot allocate the guest block");
Clear();
GuestAllocations::Mutation().Add(block, Bytes, true, true, true);
}
~GuestBlock() {
GuestAllocations::Mutation().Remove(block);
#ifdef _WIN32
VirtualFree(block, 0, MEM_RELEASE);
#else
std::free(block);
#endif
}
GuestBlock(const GuestBlock&) = delete;
GuestBlock& operator=(const GuestBlock&) = delete;
void Clear() const { std::memset(block, 0, Bytes); }
auto* Data() const { return reinterpret_cast<std::uint32_t*>(block); }
const std::size_t Bytes;
private:
std::uint8_t* block = nullptr;
};
void CheckSplitCoverage(AgcDriver::VulkanDevice& device) {
const GuestBlock guest(Stores);
const auto address = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(guest.Data()));
const std::vector<std::uint32_t> userData{
static_cast<std::uint32_t>(address), static_cast<std::uint32_t>((address >> 32u) & 0xffffu) | (4u << 16u),
Stores, 0x00014004u, Pattern[0], Pattern[1], Pattern[2], Pattern[3], Stores, Period};
const std::span<const std::uint32_t> code(Code);
const std::array<ShaderRecompiler::MemoryRegion, 1> memory{{{reinterpret_cast<std::uintptr_t>(code.data()), std::as_bytes(code)}}};
const ShaderRecompiler::ShaderComputeStageInfo compute{{Threads, 1, 1}, 0, {true, false, false}, false, 1};
ShaderRecompiler::RecompileRequest request{
{ShaderRecompiler::ShaderStage::Compute, reinterpret_cast<std::uintptr_t>(code.data()), code, 0, {}},
{64, 0, userData, compute, std::nullopt, std::nullopt, memory},
device.Target(), {0, 0, 0, 128}, std::nullopt, false};
const auto result = ShaderRecompiler::Recompile(request);
device.Dispatch(result, Groups, 1, 1, {}, reinterpret_cast<std::uintptr_t>(code.data()));
device.WaitIdle();
AgcDriver::GuestMemory::FlushGpuWrites(reinterpret_cast<std::uintptr_t>(guest.Data()), guest.Bytes);
for (std::uint32_t i = 0; i < Stores; ++i) {
Require(guest.Data()[i] == Pattern[i % Period], "oversized dispatch: dword " + std::to_string(i) + " is 0x" + std::to_string(guest.Data()[i]) + ", expected 0x" + std::to_string(Pattern[i % Period]));
}
}
constexpr std::array<std::uint32_t, 69> LdsCode{
0x34020084, 0xe0381000, 0x80000401, 0x7e100208, 0x7e3e0209, 0x343e3e82, 0x4a103f08, 0x7e40020a,
0x34404083, 0x4a104108, 0x34121090, 0x4a121300, 0x4a1412ff, 0x10000000, 0x7e160280, 0x7e180280,
0x7e1a0281, 0x7e1e0281, 0x7e200280, 0xbf8c0000, 0xd9340000, 0x00000b06, 0xd834000c, 0x00000b06,
0xd8340000, 0x00000904, 0xd8340000, 0x00000a05, 0xd8340000, 0x00000907, 0xbf8c0000, 0xbf8a0000,
0xd880000c, 0x28000d06, 0xd9800000, 0x29000f06, 0x3a260884, 0x3a2a0aff, 0x00000080, 0x363608ff,
0xfffffff8, 0xbf8c0000, 0xbf8a0000, 0xd8d80000, 0x2b000013, 0xd8d80000, 0x2c000015, 0xd8d80000,
0x2d000007, 0xd8d8000c, 0x2e000006, 0xd9d80000, 0x2f000006, 0xd9d80000, 0x3100001b, 0xd8d80000,
0x33000005, 0x343c1086, 0x4a3c3d00, 0x343c3c86, 0xbf8c0000, 0xe0781000, 0x8001281e, 0xe0781010,
0x80012c1e, 0xe0781020, 0x8001301e, 0xbf8c0000, 0xbf810000,
};
alignas(256) std::array<std::uint32_t, Threads * 4> LdsInput{};
alignas(256) std::array<std::uint32_t, Threads * 4> LdsOutput{};
void CheckLdsSplitThrows(AgcDriver::VulkanDevice& device) {
const auto target = device.Target();
const auto ldsBytes = target.maxWorkgroupSharedMemoryBytes + 4096u;
std::vector<std::uint32_t> userData(8, 0u);
const auto descriptor = [&](const void* data, std::uint32_t bytes) {
const auto address = reinterpret_cast<std::uintptr_t>(data);
return std::array<std::uint32_t, 4>{static_cast<std::uint32_t>(address), static_cast<std::uint32_t>((address >> 32u) & 0xffffu), bytes, 0x31016facu};
};
const auto input = descriptor(LdsInput.data(), static_cast<std::uint32_t>(LdsInput.size() * 4u));
const auto output = descriptor(LdsOutput.data(), static_cast<std::uint32_t>(LdsOutput.size() * 4u));
std::copy(input.begin(), input.end(), userData.begin());
std::copy(output.begin(), output.end(), userData.begin() + 4);
const std::span<const std::uint32_t> code(LdsCode);
const std::array<ShaderRecompiler::MemoryRegion, 1> memory{{{reinterpret_cast<std::uintptr_t>(code.data()), std::as_bytes(code)}}};
const ShaderRecompiler::ShaderComputeStageInfo compute{{Threads, 1, 1}, ldsBytes / 4u, {true, true, true}, false, 1};
ShaderRecompiler::RecompileRequest request{
{ShaderRecompiler::ShaderStage::Compute, reinterpret_cast<std::uintptr_t>(code.data()), code, 0, {}},
{64, 0, userData, compute, std::nullopt, std::nullopt, memory},
target, {0, 0, 0, 128}, std::nullopt, false};
request.useCache = false;
const auto result = ShaderRecompiler::Recompile(request);
Require(result.workgroupMemoryDwords != 0u, "oversized dispatch: the LDS shader did not take the device-memory path");
bool threw = false;
try {
device.Dispatch(result, 1, device.WorkGroupCountLimit()[1] + 1u, 1, {}, reinterpret_cast<std::uintptr_t>(code.data()));
} catch (const std::runtime_error&) {
threw = true;
}
device.WaitIdle();
Require(threw, "oversized dispatch: a Y count over the limit with device-memory LDS did not throw");
}
void CheckNvidiaSplitThrows(AgcDriver::VulkanDevice& device) {
const GuestBlock guest(Threads);
const auto address = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(guest.Data()));
const std::vector<std::uint32_t> userData{
static_cast<std::uint32_t>(address), static_cast<std::uint32_t>((address >> 32u) & 0xffffu) | (4u << 16u),
Threads, 0x00014004u, Pattern[0], Pattern[1], Pattern[2], Pattern[3], Threads, Period};
const std::span<const std::uint32_t> code(Code);
const std::array<ShaderRecompiler::MemoryRegion, 1> memory{{{reinterpret_cast<std::uintptr_t>(code.data()), std::as_bytes(code)}}};
const ShaderRecompiler::ShaderComputeStageInfo compute{{Threads, 1, 1}, 0, {true, false, false}, false, 1};
ShaderRecompiler::RecompileRequest request{
{ShaderRecompiler::ShaderStage::Compute, reinterpret_cast<std::uintptr_t>(code.data()), code, 0, {}},
{64, 0, userData, compute, std::nullopt, std::nullopt, memory},
device.Target(), {0, 0, 0, 128}, std::nullopt, false};
const auto result = ShaderRecompiler::Recompile(request);
bool threw = false;
try {
device.Dispatch(result, 1, device.WorkGroupCountLimit()[1] + 1u, 1, {}, reinterpret_cast<std::uintptr_t>(code.data()));
} catch (const std::runtime_error&) {
threw = true;
}
device.WaitIdle();
Require(threw, "oversized dispatch: NVIDIA did not throw for a count over the workgroup limit");
}
}
int main() {
try {
const auto device = OpenVulkanTestDevice();
if (!device) return VulkanTestSkipped;
CheckLdsSplitThrows(*device);
if (device->VendorId() == 0x10de) {
CheckNvidiaSplitThrows(*device);
} else if (device->WorkGroupCountLimit()[0] < Groups) {
CheckSplitCoverage(*device);
} else {
std::printf("oversized dispatch split not tested: the device's %u workgroup limit never splits a %ux1x1 dispatch\n", device->WorkGroupCountLimit()[0], Groups);
return VulkanTestSkipped;
}
AgcDriverShutdown_nid_postfix();
std::puts("oversized dispatch tests passed");
return 0;
} catch (const std::exception& error) {
std::cerr << error.what() << '\n';
AgcDriverShutdown_nid_postfix();
return 1;
}
}