Merge pull request #612 from oneandonlydean/perf/recompiler-lds-constant-slots

perf(recompiler): index lane-strided per-invocation LDS by constant slots
This commit is contained in:
Dean 2026-10-05 21:24:42 +13:00 committed by GitHub
commit d18fa1367c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 295 additions and 1 deletions

View File

@ -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<std::uint32_t> rebased;
for (const auto& [inst, address] : laneRequirements.functionLdsAddresses) rebased.push_back(address);
std::sort(rebased.begin(), rebased.end());
require(rebased == std::vector<std::uint32_t>{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<std::uint32_t> rebased;
for (const auto& [inst, address] : summedRequirements.functionLdsAddresses) rebased.push_back(address);
std::sort(rebased.begin(), rebased.end());
require(rebased == std::vector<std::uint32_t>{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)});

View File

@ -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 <algorithm>
#include <array>
#include <cstdio>
#include <iostream>
#include <string>
#include <vector>
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<std::byte, Width * Height * 4> Pixels{};
alignas(256) std::array<std::uint32_t, Width * Height * 4> Results{};
alignas(256) constexpr std::array<std::uint32_t, 6> VertexCode{
0xe0382000, 0x80000005, 0xbf8c3f70, 0xf80008cf, 0x03020100, 0xbf810000,
};
std::vector<std::uint32_t> PixelCode() {
std::vector<std::uint32_t> 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<std::uint32_t, 4> BufferDescriptor(const void* data, std::uint32_t stride, std::uint32_t count) {
const auto address = reinterpret_cast<std::uintptr_t>(data);
return {static_cast<std::uint32_t>(address), static_cast<std::uint32_t>((address >> 32u) & 0xffffu) | (stride << 16u), count, 0x01016facu};
}
constexpr std::array<std::array<float, 4>, 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<std::uint32_t> vertexUserData(4, 0u);
const auto vertexBuffer = BufferDescriptor(triangles.data(), 16u, static_cast<std::uint32_t>(triangles.size()));
std::copy(vertexBuffer.begin(), vertexBuffer.end(), vertexUserData.begin());
const std::array<ShaderRecompiler::MemoryRegion, 1> vertexMemory{{{reinterpret_cast<std::uintptr_t>(VertexCode.data()), std::as_bytes(std::span(VertexCode))}}};
ShaderRecompiler::RecompileRequest vertex{
{ShaderStage::Vertex, reinterpret_cast<std::uintptr_t>(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<std::uint32_t>(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<std::uint32_t> pixelUserData(8, 0u);
const auto results = BufferDescriptor(Results.data(), 16u, Width * Height);
std::copy(results.begin(), results.end(), pixelUserData.begin() + 4);
const std::array<ShaderRecompiler::MemoryRegion, 1> pixelMemory{{{reinterpret_cast<std::uintptr_t>(pixelCode.data()), std::as_bytes(std::span(pixelCode))}}};
ShaderRecompiler::RecompileRequest fragment{
{ShaderStage::Fragment, reinterpret_cast<std::uintptr_t>(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<AgcDriver::Graphics::CompiledShader, 2> 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<std::uintptr_t>(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<float>(Height), static_cast<float>(Width), -static_cast<float>(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<std::uint32_t>(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;
}
}

View File

@ -4,6 +4,7 @@
#include "IntermediateRepresentation/IrProgram.hpp"
#include <cstdint>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
@ -18,6 +19,7 @@ struct SpirvRequirements {
bool functionLds = false;
bool ldsLock = false;
std::uint32_t functionLdsDwords = 0;
std::unordered_map<const IrValue*, std::uint32_t> 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<const IrValue*, std::uint32_t> FunctionLdsLaneAddresses(const IrProgram& program);
[[nodiscard]] std::unordered_set<const IrValue*> WaveUniformValues(const IrProgram& program);
[[nodiscard]] bool IsWaveMaskBranch(BranchCondition condition);

View File

@ -238,6 +238,67 @@ std::uint32_t FunctionLdsDwords(const IrProgram& program) {
return static_cast<std::uint32_t>(std::min<std::uint64_t>(std::max<std::uint64_t>((needed + 63u) & ~63ull, 64u), FunctionLdsDwordLimit));
}
namespace {
struct LaneAffine {
std::uint64_t stride = 0;
std::uint64_t constant = 0;
};
std::optional<LaneAffine> 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<const IrValue*, std::uint32_t> FunctionLdsLaneAddresses(const IrProgram& program) {
std::unordered_map<const IrValue*, std::uint32_t> addresses;
if (FunctionLdsDwords(program) >= FunctionLdsDwordLimit) return {};
std::optional<std::uint64_t> 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<MemoryFlags>().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<std::uint32_t>(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<MemoryFlags>().index].offset;
needed = std::max<std::uint64_t>(needed, ((std::uint64_t{address} + offset) >> 2u) + SharedAccessDwords(inst->Opcode()));
}
requirements.functionLdsDwords = static_cast<std::uint32_t>(std::max<std::uint64_t>((needed + 63u) & ~63ull, 64u));
}
}
return requirements;
}

View File

@ -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);
}