From 31af5a9d2232b9950ecd977c59f2ce29ca175b20 Mon Sep 17 00:00:00 2001 From: Dean Galvin Date: Mon, 5 Oct 2026 13:36:44 +1300 Subject: [PATCH] perf(recompiler): index lane-strided per-invocation LDS by constant slots Stages without workgroup memory keep LDS in a per-invocation Function array. With the array sized to the addresses it reaches (#586), PPSA21564's intro pixel shader still takes 864 bytes of local memory per thread on NVIDIA: its accesses are lane * 4 + {0, 256, 512}, so the index stays dynamic and the 192-dword array lives in local memory. When every LDS access's address is lane * stride + constant (lane id, constants, adds, multiplies and left shifts), with one stride, a whole number of dwords, for all of them and no wrap (the bound analysis already proves that), the emitter now addresses each access by its constant plus its instruction offset. The lane id is the same for every access of an invocation, so taking lane * stride out maps the addresses one to one, keeps which accesses alias and keeps every access in bounds; the array sees only constant indices and the driver keeps it in registers. Shaders with any other address form are emitted as before. --- .../libSceAgcDriver/tests/ShaderMemory.cpp | 49 ++++++ .../tests/execution/PixelLdsLaneSlots.cpp | 165 ++++++++++++++++++ .../include/SpirvBackend/SpirvAnalysis.hpp | 3 + .../SpirvBackend/src/SpirvAnalysis.cpp | 74 +++++++- .../SpirvMemory/SpirvMemoryInstructions.cpp | 5 + 5 files changed, 295 insertions(+), 1 deletion(-) create mode 100644 core/libs/prx/libSceAgcDriver/tests/execution/PixelLdsLaneSlots.cpp diff --git a/core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp b/core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp index dad622661..a65c7d9e5 100644 --- a/core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp +++ b/core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp @@ -937,6 +937,54 @@ void verifyFunctionLdsBound() { }); require(laneSlots == 192u, "function LDS: lane-strided slots up to byte 767 must take 192 dwords"); require(laneRequirements.functionLds && laneRequirements.functionLdsDwords == 192u, "function LDS: the requirements do not carry the bounded size"); + { + std::vector rebased; + for (const auto& [inst, address] : laneRequirements.functionLdsAddresses) rebased.push_back(address); + std::sort(rebased.begin(), rebased.end()); + require(rebased == std::vector{0u, 0u, 0u, 0u}, "function LDS: lane-strided slots must be addressed without the lane term"); + } + + const auto [mixed, mixedRequirements] = build([](IrProgram&, IrBlock&, auto& emit, auto& constant, IrValue& active) { + auto& lane = emit(IrOpcode::LaneId, IrType::U32, {}); + auto& address = emit(IrOpcode::ShiftLeftLogical32, IrType::U32, {&lane, &constant(2u)}); + emit(IrOpcode::WriteSharedU32, IrType::Void, {&address, &lane, &active}, 0u); + emit(IrOpcode::LoadSharedU32, IrType::U32, {&constant(8u), &active}, 0u); + }); + require(mixed == 64u && mixedRequirements.functionLdsAddresses.empty(), "function LDS: a constant address beside lane-strided ones must keep the lane term"); + + const auto [strides, stridesRequirements] = build([](IrProgram&, IrBlock&, auto& emit, auto& constant, IrValue& active) { + auto& lane = emit(IrOpcode::LaneId, IrType::U32, {}); + auto& four = emit(IrOpcode::ShiftLeftLogical32, IrType::U32, {&lane, &constant(2u)}); + auto& eight = emit(IrOpcode::IMul32, IrType::U32, {&lane, &constant(8u)}); + emit(IrOpcode::WriteSharedU32, IrType::Void, {&four, &lane, &active}, 0u); + emit(IrOpcode::LoadSharedU32, IrType::U32, {&eight, &active}, 0u); + }); + require(strides == 128u && stridesRequirements.functionLdsAddresses.empty(), "function LDS: two lane strides must keep the lane term"); + + const auto [halfWord, halfWordRequirements] = build([](IrProgram&, IrBlock&, auto& emit, auto& constant, IrValue& active) { + auto& lane = emit(IrOpcode::LaneId, IrType::U32, {}); + auto& address = emit(IrOpcode::ShiftLeftLogical32, IrType::U32, {&lane, &constant(1u)}); + emit(IrOpcode::WriteSharedU32, IrType::Void, {&address, &lane, &active}, 0u); + }); + require(halfWord == 64u && halfWordRequirements.functionLdsAddresses.empty(), "function LDS: a lane stride that is not whole dwords must keep the lane term"); + + const auto [summed, summedRequirements] = build([](IrProgram&, IrBlock&, auto& emit, auto& constant, IrValue& active) { + auto& lane = emit(IrOpcode::LaneId, IrType::U32, {}); + auto& scaled = emit(IrOpcode::IMul32, IrType::U32, {&constant(16u), &lane}); + auto& low = emit(IrOpcode::IAdd32, IrType::U32, {&scaled, &constant(4u)}); + auto& other = emit(IrOpcode::LaneId, IrType::U32, {}); + auto& shifted = emit(IrOpcode::ShiftLeftLogical32, IrType::U32, {&other, &constant(4u)}); + auto& high = emit(IrOpcode::IAdd32, IrType::U32, {&constant(1024u), &shifted}); + emit(IrOpcode::WriteSharedU32x2, IrType::Void, {&low, &lane, &lane, &active}, 0u); + emit(IrOpcode::LoadSharedU32x4, IrType::U32x4, {&high, &active}, 1u); + }); + { + std::vector rebased; + for (const auto& [inst, address] : summedRequirements.functionLdsAddresses) rebased.push_back(address); + std::sort(rebased.begin(), rebased.end()); + require(rebased == std::vector{4u, 1024u}, "function LDS: lane ids of separate instructions with one stride must share the rebase"); + require(summedRequirements.functionLdsDwords == 384u && summed == 576u, "function LDS: the rebased array must hold the rebased addresses"); + } const auto [wide, wideRequirements] = build([](IrProgram&, IrBlock&, auto& emit, auto& constant, IrValue& active) { auto& lane = emit(IrOpcode::LaneId, IrType::U32, {}); @@ -947,6 +995,7 @@ void verifyFunctionLdsBound() { }); require(wide == 128u, "function LDS: a 4-dword access at byte 0x100 must take 68 dwords, rounded to 128"); require(wideRequirements.functionLdsDwords == 128u, "function LDS: the requirements do not carry the wide access's size"); + require(wideRequirements.functionLdsAddresses.empty(), "function LDS: a selected address must keep the lane term"); const auto [unbounded, unboundedRequirements] = build([](IrProgram&, IrBlock&, auto& emit, auto& constant, IrValue& active) { auto& user = emit(IrOpcode::GetUserData, IrType::U32, {&constant(0u)}); diff --git a/core/libs/prx/libSceAgcDriver/tests/execution/PixelLdsLaneSlots.cpp b/core/libs/prx/libSceAgcDriver/tests/execution/PixelLdsLaneSlots.cpp new file mode 100644 index 000000000..a1aa594a1 --- /dev/null +++ b/core/libs/prx/libSceAgcDriver/tests/execution/PixelLdsLaneSlots.cpp @@ -0,0 +1,165 @@ +#include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp" +#include "prx/libSceAgcDriver/Graphics/include/Draw.hpp" +#include "Recompiler.hpp" +#include "VulkanTestDevice.hpp" +#include +#include +#include +#include +#include +#include + +namespace { + +using AgcDriver::Graphics::Require; +using ShaderRecompiler::ShaderStage; + +constexpr std::uint32_t Sentinel = 0xdeadbeefu; +constexpr std::uint32_t Width = 64; +constexpr std::uint32_t Height = 32; +alignas(256) std::array Pixels{}; +alignas(256) std::array Results{}; + +alignas(256) constexpr std::array VertexCode{ + 0xe0382000, 0x80000005, 0xbf8c3f70, 0xf80008cf, 0x03020100, 0xbf810000, +}; + +std::vector PixelCode() { + std::vector code{0x7e040f00, 0x7e060f01, 0xd5430006, 0x0409ff03, Width}; + const auto vop2 = [&](std::uint32_t opcode, std::uint32_t destination, std::uint32_t source, std::uint32_t vector) { + code.push_back(opcode | (destination << 17u) | (vector << 9u) | source); + }; + vop2(0x34000000, 8, 0x82, 6); + code.push_back(0xbefc0380); + for (std::uint32_t bank = 0; bank < 3; ++bank) { + vop2(0x4a000000, 5, 0x80 + bank, 8); + code.push_back(0xdac00000 | (bank * 256u)); + code.push_back(0x00000500); + } + code.push_back(0xbf8cc07f); + for (std::uint32_t bank = 0; bank < 3; ++bank) { + code.push_back(0xdac40000 | (((bank + 1u) % 3u) * 256u)); + code.push_back((10u + bank) << 24u); + } + code.push_back(0xbf8cc07f); + for (std::uint32_t bank = 0; bank < 3; ++bank) { + code.push_back(0xe0702000 | (bank * 4u)); + code.push_back(0x80010006 | ((10u + bank) << 8u)); + } + code.insert(code.end(), {0x7e0e02f2, 0xf800180f, 0x07070707, 0xbf810000}); + return code; +} + +std::array BufferDescriptor(const void* data, std::uint32_t stride, std::uint32_t count) { + const auto address = reinterpret_cast(data); + return {static_cast(address), static_cast((address >> 32u) & 0xffffu) | (stride << 16u), count, 0x01016facu}; +} + +constexpr std::array, 3> Triangles{{ + {-1.0f, -1.0f, 0.5f, 1.0f}, {3.0f, -1.0f, 0.5f, 1.0f}, {-1.0f, 3.0f, 0.5f, 1.0f} +}}; + +void Draw(AgcDriver::VulkanDevice& device, std::uint32_t waveSize) { + const auto& triangles = Triangles; + const auto pixelCode = PixelCode(); + Pixels.fill(std::byte{0}); + Results.fill(Sentinel); + const auto target = device.Target(); + + std::vector vertexUserData(4, 0u); + const auto vertexBuffer = BufferDescriptor(triangles.data(), 16u, static_cast(triangles.size())); + std::copy(vertexBuffer.begin(), vertexBuffer.end(), vertexUserData.begin()); + const std::array vertexMemory{{{reinterpret_cast(VertexCode.data()), std::as_bytes(std::span(VertexCode))}}}; + ShaderRecompiler::RecompileRequest vertex{ + {ShaderStage::Vertex, reinterpret_cast(VertexCode.data()), VertexCode, 0, {}}, + {waveSize, 0, vertexUserData, std::nullopt, std::nullopt, ShaderRecompiler::ShaderVertexStageInfo{}, vertexMemory}, + target, + {0, 0, 0, 64} + }; + vertex.useCache = false; + const auto vertexResult = ShaderRecompiler::Recompile(vertex); + const auto vertexPush = static_cast(vertexResult.pushConstants.size()); + + ShaderRecompiler::ShaderPixelStageInfo pixel{}; + pixel.wave32 = waveSize == 32u; + pixel.inputAddr = ShaderRecompiler::PixelInputBit(ShaderRecompiler::PixelInput::PositionX) | ShaderRecompiler::PixelInputBit(ShaderRecompiler::PixelInput::PositionY); + pixel.posX = true; + pixel.posY = true; + pixel.targetOutputMode[0] = 9; + pixel.targetExportMapping.fill(0xe4u); + std::vector pixelUserData(8, 0u); + const auto results = BufferDescriptor(Results.data(), 16u, Width * Height); + std::copy(results.begin(), results.end(), pixelUserData.begin() + 4); + const std::array pixelMemory{{{reinterpret_cast(pixelCode.data()), std::as_bytes(std::span(pixelCode))}}}; + ShaderRecompiler::RecompileRequest fragment{ + {ShaderStage::Fragment, reinterpret_cast(pixelCode.data()), pixelCode, 0, {}}, + {waveSize, 0, pixelUserData, std::nullopt, pixel, std::nullopt, pixelMemory}, + target, + {0, 0, vertexPush, 128 - vertexPush} + }; + fragment.useCache = false; + const auto pixelResult = ShaderRecompiler::Recompile(fragment); + const std::array shaders{{ + {ShaderStage::Vertex, &vertexResult, 0}, + {ShaderStage::Fragment, &pixelResult, vertexPush} + }}; + + AgcDriver::Graphics::State state{}; + state.stages = {AgcDriver::Graphics::ShaderPath::Vertex, 0u, waveSize, waveSize, std::nullopt, std::nullopt}; + state.color = {reinterpret_cast(Pixels.data()), {Width, Height}, VK_FORMAT_R8G8B8A8_UNORM, Pixels.size(), 0xe4u}; + state.colors = {state.color}; + state.hasColorTarget = true; + state.renderExtent = {Width, Height}; + state.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST; + state.viewport = {0, static_cast(Height), static_cast(Width), -static_cast(Height), 0, 1}; + state.negativeOneToOne = false; + state.scissor = {{0, 0}, {Width, Height}}; + state.cullMode = VK_CULL_MODE_NONE; + state.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE; + state.blend.colorWriteMask = 15; + state.blends = {state.blend}; + state.blendConstants = {}; + const AgcDriver::Pm4::DrawParameters draw{0, static_cast(triangles.size()), 0, 1, 0, false}; + device.Draw(state, draw, shaders); + device.WaitIdle(); +} + +void Check(std::uint32_t waveSize) { + const auto what = "pixel LDS lane slots wave" + std::to_string(waveSize); + for (std::uint32_t pixel = 0; pixel < Width * Height; ++pixel) { + for (std::uint32_t channel = 0; channel < 4; ++channel) { + Require(Pixels[pixel * 4u + channel] == std::byte{255}, what + ": missing color at pixel " + std::to_string(pixel)); + } + for (std::uint32_t bank = 0; bank < 3; ++bank) { + const auto expected = pixel * 4u + (bank + 1u) % 3u; + const auto actual = Results[pixel * 4u + bank]; + Require(actual == expected, what + ": pixel " + std::to_string(pixel) + " bank " + std::to_string(bank) + " expected " + std::to_string(expected) + " got " + std::to_string(actual)); + } + Require(Results[pixel * 4u + 3u] == Sentinel, what + ": output stride guard overwritten"); + } +} + +} + +int main() { + try { + const auto device = OpenVulkanTestDevice(); + if (!device) return VulkanTestSkipped; + if (device->Target().subgroupSize < 32u) { + std::printf("skipped, subgroup size %u cannot hold a wave32\n", device->Target().subgroupSize); + return VulkanTestSkipped; + } + for (const auto waveSize : {64u, 32u}) { + std::fprintf(stderr, "pixel LDS lane slots wave%u\n", waveSize); + for (std::uint32_t pass = 0; pass < 3; ++pass) { + Draw(*device, waveSize); + Check(waveSize); + } + } + std::puts("pixel LDS lane slot tests passed: 6 draws, 12288 pixels, 36864 LDS results"); + return 0; + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/core/shader/recompiler/SpirvBackend/include/SpirvBackend/SpirvAnalysis.hpp b/core/shader/recompiler/SpirvBackend/include/SpirvBackend/SpirvAnalysis.hpp index 0492acf15..0b433bfeb 100644 --- a/core/shader/recompiler/SpirvBackend/include/SpirvBackend/SpirvAnalysis.hpp +++ b/core/shader/recompiler/SpirvBackend/include/SpirvBackend/SpirvAnalysis.hpp @@ -4,6 +4,7 @@ #include "IntermediateRepresentation/IrProgram.hpp" #include #include +#include #include #include @@ -18,6 +19,7 @@ struct SpirvRequirements { bool functionLds = false; bool ldsLock = false; std::uint32_t functionLdsDwords = 0; + std::unordered_map functionLdsAddresses; bool functionScratch = false; bool pixelValidMask = false; bool bufferInt64Atomics = false; @@ -31,6 +33,7 @@ struct SpirvRequirements { [[nodiscard]] SpirvRequirements AnalyzeProgramRequirements(const IrProgram& program); inline constexpr std::uint32_t FunctionLdsDwordLimit = 8192u; [[nodiscard]] std::uint32_t FunctionLdsDwords(const IrProgram& program); +[[nodiscard]] std::unordered_map FunctionLdsLaneAddresses(const IrProgram& program); [[nodiscard]] std::unordered_set WaveUniformValues(const IrProgram& program); [[nodiscard]] bool IsWaveMaskBranch(BranchCondition condition); diff --git a/core/shader/recompiler/SpirvBackend/src/SpirvAnalysis.cpp b/core/shader/recompiler/SpirvBackend/src/SpirvAnalysis.cpp index d85b226ab..82fdd645b 100644 --- a/core/shader/recompiler/SpirvBackend/src/SpirvAnalysis.cpp +++ b/core/shader/recompiler/SpirvBackend/src/SpirvAnalysis.cpp @@ -238,6 +238,67 @@ std::uint32_t FunctionLdsDwords(const IrProgram& program) { return static_cast(std::min(std::max((needed + 63u) & ~63ull, 64u), FunctionLdsDwordLimit)); } +namespace { + +struct LaneAffine { + std::uint64_t stride = 0; + std::uint64_t constant = 0; +}; + +std::optional LaneAffineOf(const IrValue* value, std::uint32_t depth = 0) { + if (value == nullptr || depth > 64u) return std::nullopt; + value = value->Resolve(); + if (value == nullptr) return std::nullopt; + if (value->HasImmediate()) return LaneAffine{0u, value->ImmediateU32()}; + const auto argument = [&](std::size_t index) { return index < value->ArgumentCount() ? LaneAffineOf(value->Argument(index), depth + 1u) : std::nullopt; }; + switch (value->Opcode()) { + case IrOpcode::LaneId: + return LaneAffine{1u, 0u}; + case IrOpcode::IAdd32: { + const auto a = argument(0), b = argument(1); + if (!a || !b) return std::nullopt; + return LaneAffine{a->stride + b->stride, a->constant + b->constant}; + } + case IrOpcode::IMul32: { + const auto a = argument(0), b = argument(1); + if (!a || !b) return std::nullopt; + if (a->stride == 0u) return LaneAffine{b->stride * a->constant, b->constant * a->constant}; + if (b->stride == 0u) return LaneAffine{a->stride * b->constant, a->constant * b->constant}; + return std::nullopt; + } + case IrOpcode::ShiftLeftLogical32: + case IrOpcode::IShiftLeft32: { + const auto a = argument(0), b = argument(1); + if (!a || !b || b->stride != 0u || b->constant >= 32u) return std::nullopt; + return LaneAffine{a->stride << b->constant, a->constant << b->constant}; + } + default: + return std::nullopt; + } +} + +} + +std::unordered_map FunctionLdsLaneAddresses(const IrProgram& program) { + std::unordered_map addresses; + if (FunctionLdsDwords(program) >= FunctionLdsDwordLimit) return {}; + std::optional stride; + for (const IrBlock* block : program.BlockOrder()) { + for (const IrValue* inst : block->Instructions()) { + if (SharedAccessOf(inst->Opcode()) == SharedAccess::None) continue; + const auto& memory = program.Resources().memoryInfo[inst->Flags().index]; + if (memory.kind != ResourceKind::Lds) continue; + const auto affine = LaneAffineOf(inst->Argument(0)); + if (!affine || affine->stride % 4u != 0u || (stride && *stride != affine->stride)) return {}; + stride = affine->stride; + if (affine->constant + memory.offset > 0xffffffffull) return {}; + addresses.emplace(inst, static_cast(affine->constant)); + } + } + if (!stride || *stride == 0u) return {}; + return addresses; +} + SpirvRequirements AnalyzeProgramRequirements(const IrProgram& program) { SpirvRequirements requirements {}; for (const IrBlock* block : program.BlockOrder()) { @@ -372,7 +433,18 @@ SpirvRequirements AnalyzeProgramRequirements(const IrProgram& program) { requirements.subgroupBallot = true; } } - if (requirements.functionLds) requirements.functionLdsDwords = FunctionLdsDwords(program); + if (requirements.functionLds) { + requirements.functionLdsDwords = FunctionLdsDwords(program); + requirements.functionLdsAddresses = FunctionLdsLaneAddresses(program); + if (!requirements.functionLdsAddresses.empty()) { + std::uint64_t needed = 0; + for (const auto& [inst, address] : requirements.functionLdsAddresses) { + const auto offset = program.Resources().memoryInfo[inst->Flags().index].offset; + needed = std::max(needed, ((std::uint64_t{address} + offset) >> 2u) + SharedAccessDwords(inst->Opcode())); + } + requirements.functionLdsDwords = static_cast(std::max((needed + 63u) & ~63ull, 64u)); + } + } return requirements; } diff --git a/core/shader/recompiler/SpirvBackend/src/SpirvMemory/SpirvMemoryInstructions.cpp b/core/shader/recompiler/SpirvBackend/src/SpirvMemory/SpirvMemoryInstructions.cpp index 5df19bf32..60d891f99 100644 --- a/core/shader/recompiler/SpirvBackend/src/SpirvMemory/SpirvMemoryInstructions.cpp +++ b/core/shader/recompiler/SpirvBackend/src/SpirvMemory/SpirvMemoryInstructions.cpp @@ -88,6 +88,11 @@ std::uint32_t ByteAddress(SpirvValueEmitContext& ctx, const IrValue& inst, const return BufferByteAddress(ctx, inst, mem, ctx.Arg(inst, 1), ctx.Arg(inst, 2), ctx.Arg(inst, 3)); case ResourceKind::Lds: case ResourceKind::Gds: + if (mem.kind == ResourceKind::Lds) { + if (const auto found = state.requirements.functionLdsAddresses.find(&inst); found != state.requirements.functionLdsAddresses.end()) { + return ConstantU32(state, found->second + mem.offset); + } + } if (mem.offset == 0u) { return ctx.Arg(inst, 0); }