fix(recompiler): give wave64 vertex programs on 32-lane hosts a 32-lane s3

core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp

@@ -1210,6 +1210,101 @@ void verifyHalfWaveReduction() {
     expectFailure([&] { static_cast<void>(pixel(umin, 8u, true)); }, "v_permlanex16_b32 from lanes 16-31 is outside the 8-lane host subgroup", "half-wave reduction: the scan built on an 8-lane subgroup");
 }
 
+void verifyNarrowVertexWave() {
+    using namespace ShaderRecompiler;
+    static constexpr std::array<std::uint32_t, 4> capabilities{spv::CapabilityGroupNonUniform, spv::CapabilityGroupNonUniformBallot, spv::CapabilityGroupNonUniformShuffle, spv::CapabilityGroupNonUniformArithmetic};
+    static constexpr std::array<std::uint32_t, 4> userData{0u, 0u, 0u, 0x2020u};
+    const auto vertex = [](std::span<const std::uint32_t> code, std::uint32_t waveSize, std::uint32_t subgroupSize, std::uint32_t userDataBase) {
+        RecompileRequest request{};
+        request.shader = {ShaderStage::Vertex, 0x50000u, code, 0, {}};
+        request.context.waveSize = waveSize;
+        request.context.userDataBaseRegister = userDataBase;
+        if (userDataBase == 0u) request.context.userData = userData;
+        request.context.vertex = ShaderVertexStageInfo{};
+        request.target.vulkanVersion = 0x00401000u;
+        request.target.spirvVersion = 0x00010300u;
+        request.target.subgroupSize = subgroupSize;
+        request.target.supportedCapabilities = capabilities;
+        request.target.fragmentShaderBarycentricEnabled = false;
+        request.layout.pushConstantSizeBytes = 128;
+        return request;
+    };
+    const auto prepared = [&vertex](std::span<const std::uint32_t> code, std::uint32_t subgroupSize, std::uint32_t userDataBase = 8u) {
+        return GetPreparedArtifact(*PrepareShader(vertex(code, 64u, subgroupSize, userDataBase))).spirv;
+    };
+    const auto prepares = [&prepared](std::span<const std::uint32_t> code, std::uint32_t subgroupSize, std::uint32_t userDataBase = 8u) {
+        return !prepared(code, subgroupSize, userDataBase).empty();
+    };
+    const auto reductions = [](const std::vector<std::uint32_t>& spirv) {
+        std::uint32_t count = 0;
+        for (std::size_t cursor = 5; cursor < spirv.size() && (spirv[cursor] >> 16u) != 0u; cursor += spirv[cursor] >> 16u) {
+            if ((spirv[cursor] & 0xffffu) == spv::OpGroupNonUniformBitwiseOr && (spirv[cursor] >> 16u) == 6u && spirv[cursor + 4u] == spv::GroupOperationReduce) ++count;
+        }
+        return count;
+    };
+    const std::vector<std::uint32_t> culling{
+        0xd7650015u, 0x000100c1u, 0xd7660015u, 0x00022ac1u, 0x7e0e0280u, 0x7da206f9u, 0x81060015u, 0x360e0a81u,
+        0xbefe04c1u, 0xd4c50002u, 0x00020e80u, 0xd763000fu, 0x00022a81u, 0xd5010000u, 0x000a1e80u,
+        0x380000fau, 0xff091100u, 0x380000fau, 0xff091200u, 0x380000fau, 0xff091400u, 0x380000fau, 0xff091800u,
+        0xd778000fu, 0x03058300u, 0x38001f00u, 0xd7600002u, 0x00013f00u, 0xd760006bu, 0x00017f00u, 0x88026b02u,
+        0x7e020202u, 0xf80008cfu, 0x01010101u, 0xbf810000u,
+    };
+    const auto patched = [&culling](std::initializer_list<std::pair<std::size_t, std::uint32_t>> words) {
+        auto code = culling;
+        for (const auto& [index, word] : words) code[index] = word;
+        return code;
+    };
+    const char* lane63 = "v_readlane_b32 of lane 63 is outside the 32-lane host subgroup";
+    require(reductions(prepared(culling, 32u)) == 1u, "narrow vertex wave: the NGG culling scan, keyed on lanes below the primitive count in s3, did not read lane 31 as a subgroup OR on 32 lanes");
+    const auto wide64 = prepared(culling, 64u);
+    require(!wide64.empty() && reductions(wide64) == 0u, "narrow vertex wave: the NGG culling scan, keyed on lanes below 64 on 64 lanes, did not keep its DPP scan");
+    require(prepares(patched({{6, 0x80060015u}}), 32u), "narrow vertex wave: the scan keyed on lanes below the vertex count in s3 did not prepare on 32 lanes");
+    static constexpr std::array<std::uint32_t, 3> noArithmetic{spv::CapabilityGroupNonUniform, spv::CapabilityGroupNonUniformBallot, spv::CapabilityGroupNonUniformShuffle};
+    const auto otherKeys = patched({{7, 0x360e0a82u}});
+    auto withoutArithmetic = vertex(otherKeys, 64u, 32u, 8u);
+    withoutArithmetic.target.supportedCapabilities = noArithmetic;
+    const auto shuffled = GetPreparedArtifact(*PrepareShader(withoutArithmetic)).spirv;
+    require(!shuffled.empty() && reductions(shuffled) == 0u, "narrow vertex wave: without subgroup arithmetic, the scan did not keep lane 31 as a shuffle with lane 63 read as 0");
+    require(prepares(patched({{5, 0x7da82aa0u}, {6, 0xbf800000u}}), 32u), "narrow vertex wave: the scan keyed on lanes below 32 did not prepare on 32 lanes");
+    expectFailure([&] { static_cast<void>(prepares(patched({{5, 0x7da82aa1u}, {6, 0xbf800000u}}), 32u)); }, lane63, "narrow vertex wave: a scan keyed on lanes below 33 read lane 63");
+    expectFailure([&] { static_cast<void>(prepares(culling, 32u, 0u)); }, lane63, "narrow vertex wave: a scan keyed on a count from user data read lane 63");
+    expectFailure([&] { static_cast<void>(prepares(patched({{6, 0x81060005u}}), 32u)); }, lane63, "narrow vertex wave: a scan keyed on the vertex index instead of the lane read lane 63");
+    auto otherRegister = patched({{21, 0x380200fau}, {24, 0x03058301u}, {25, 0x38001f01u}});
+    otherRegister.insert(otherRegister.begin(), 0x7e020305u);
+    require(reductions(prepared(otherRegister, 32u)) == 1u, "narrow vertex wave: a scan whose bound_ctrl:1 step is written to another register did not read lane 31 as a subgroup OR on 32 lanes");
+    otherRegister[23] = 0xff011800u;
+    expectFailure([&] { static_cast<void>(prepares(otherRegister, 32u)); }, lane63, "narrow vertex wave: a scan step written to another register without bound_ctrl:1, which keeps that register's lanes outside the row, read lane 63");
+    auto threeSteps = culling;
+    threeSteps.erase(threeSteps.begin() + 21, threeSteps.begin() + 23);
+    const auto partialScan = prepared(threeSteps, 32u);
+    require(!partialScan.empty() && reductions(partialScan) == 0u, "narrow vertex wave: lane 63 of row_shr 1/2/4 steps over keys of 0 in lanes 32-63 did not read 0, or lane 31 became a subgroup OR");
+    threeSteps[5] = 0x7da82aa1u;
+    threeSteps[6] = 0xbf800000u;
+    expectFailure([&] { static_cast<void>(prepares(threeSteps, 32u)); }, lane63, "narrow vertex wave: lane 63 of row_shr 1/2/4 steps over keys of lanes below 33 read 0");
+
+    std::array<std::vector<std::uint32_t>, 2> narrow;
+    std::array<std::vector<std::uint32_t>, 2> wide;
+    for (const bool narrowFirst : {false, true}) {
+        const std::vector<std::uint32_t> code{0xbf800001u + static_cast<std::uint32_t>(narrowFirst), 0x7e020203u, 0xf80008cfu, 0x01010101u, 0xbf810000u};
+        const auto narrowRequest = vertex(code, 64u, 32u, 8u);
+        const auto wideRequest = vertex(code, 64u, 64u, 8u);
+        ShaderPreparationContext preparation;
+        const auto first = PrepareShader(narrowFirst ? narrowRequest : wideRequest, &preparation);
+        const auto second = PrepareShader(narrowFirst ? wideRequest : narrowRequest, &preparation);
+        const auto& narrowHandle = narrowFirst ? *first : *second;
+        const auto& wideHandle = narrowFirst ? *second : *first;
+        require(first->source != second->source, "narrow vertex wave: the program prepared for 32 and for 64 lanes shares one translated source");
+        require(MatchesPreparedShader(narrowRequest, narrowHandle) && MatchesPreparedShader(wideRequest, wideHandle), "narrow vertex wave: a prepared program did not match its own request");
+        require(!MatchesPreparedShader(narrowRequest, wideHandle) && !MatchesPreparedShader(wideRequest, narrowHandle), "narrow vertex wave: a program prepared for one subgroup width matched a request for the other");
+        narrow[narrowFirst ? 1u : 0u] = GetPreparedArtifact(narrowHandle).spirv;
+        wide[narrowFirst ? 1u : 0u] = GetPreparedArtifact(wideHandle).spirv;
+    }
+    require(narrow[0] == narrow[1] && wide[0] == wide[1], "narrow vertex wave: preparing for one subgroup width after the other gave another program than preparing it first");
+    require(narrow[0] != wide[0], "narrow vertex wave: s3 of a wave64 vertex program was the same on 32 and 64 lanes");
+    const std::vector<std::uint32_t> wave32{0x7e020203u, 0xf80008cfu, 0x01010101u, 0xbf810000u};
+    require(GetPreparedArtifact(*PrepareShader(vertex(wave32, 32u, 32u, 8u))).spirv == GetPreparedArtifact(*PrepareShader(vertex(wave32, 32u, 64u, 8u))).spirv, "narrow vertex wave: s3 of a wave32 vertex program depended on the subgroup width");
+}
+
 void verifyMeshConfiguration() {
     using namespace ShaderRecompiler;
     ShaderMeshInputInfo list;
@@ -2583,6 +2678,7 @@ int main(int argc, char** argv) {
         verifyProgramCounterRelativeData();
         verifyLanesOutsideHostSubgroup();
         verifyHalfWaveReduction();
+        verifyNarrowVertexWave();
         verifyMeshConfiguration();
         verifyPixelInputs();
         verifyPixelRequestSerialization();

core/libs/prx/libSceAgcDriver/tests/execution/NggVertexWaveReduction.cpp

@@ -0,0 +1,149 @@
+#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 <initializer_list>
+#include <iostream>
+#include <span>
+#include <string>
+#include <vector>
+
+namespace {
+
+using AgcDriver::Graphics::Require;
+using ShaderRecompiler::ShaderStage;
+
+constexpr std::uint32_t Width = 64;
+constexpr std::uint32_t Height = 32;
+constexpr std::uint32_t Columns = 7;
+constexpr std::uint32_t Rows = 5;
+alignas(256) std::array<std::byte, Width * Height * 4> Pixels{};
+
+alignas(256) constexpr std::array<std::uint32_t, 44> VertexCode{
+    0xd7650015, 0x000100c1, 0xd7660015, 0x00022ac1, 0x7e0e0280, 0x7da206f9, 0x81060015, 0x360e0a81,
+    0xbefe04c1, 0xd4c50002, 0x00020e80, 0xd763000f, 0x00022a81, 0xd5010000, 0x000a1e80, 0x380000fa,
+    0xff091100, 0x380000fa, 0xff091200, 0x380000fa, 0xff091400, 0x380000fa, 0xff091800, 0xd778000f,
+    0x03058300, 0x38001f00, 0x88040302, 0xd7600002, 0x00013f00, 0xd760006b, 0x00017f00, 0x88026b02,
+    0xbf060402, 0x850580f2, 0xe0382000, 0x80020005, 0xbf8c3f70, 0x10000005, 0x10020205, 0x10040405,
+    0x10060605, 0xf80008cf, 0x03020100, 0xbf810000,
+};
+
+alignas(256) constexpr std::array<std::uint32_t, 4> PixelCode{
+    0x7e0e02f2, 0xf800180f, 0x07070707, 0xbf810000,
+};
+
+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};
+}
+
+std::vector<std::array<float, 4>> Triangles() {
+    std::vector<std::array<float, 4>> vertices;
+    for (std::uint32_t cell = 0; cell < Columns * Rows; ++cell) {
+        const float left = -1.0f + 2.0f * static_cast<float>(cell % Columns) / Columns;
+        const float right = -1.0f + 2.0f * static_cast<float>(cell % Columns + 1u) / Columns;
+        const float bottom = -1.0f + 2.0f * static_cast<float>(cell / Columns) / Rows;
+        const float top = -1.0f + 2.0f * static_cast<float>(cell / Columns + 1u) / Rows;
+        for (const auto& corner : {std::array{left, bottom}, std::array{right, bottom}, std::array{right, top}, std::array{left, bottom}, std::array{right, top}, std::array{left, top}}) {
+            vertices.push_back({corner[0], corner[1], 0.5f, 1.0f});
+        }
+    }
+    return vertices;
+}
+
+void Draw(AgcDriver::VulkanDevice& device) {
+    static const auto triangles = Triangles();
+    Pixels.fill(std::byte{0});
+    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, {}},
+        {64u, 8, 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.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);
+    const std::vector<std::uint32_t> pixelUserData(8, 0u);
+    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, {}},
+        {64u, 0, pixelUserData, std::nullopt, pixel, std::nullopt, pixelMemory},
+        target,
+        PixelPushLayout(vertexPush, target)
+    };
+    fragment.useCache = false;
+    const auto pixelResult = ShaderRecompiler::Recompile(fragment);
+    const std::array<AgcDriver::Graphics::CompiledShader, 2> shaders{{
+        {ShaderStage::Vertex, &vertexResult, 0},
+        {ShaderStage::Fragment, &pixelResult, PixelPushOffset(vertexPush, target)}
+    }};
+
+    AgcDriver::Graphics::State state{};
+    state.stages = {AgcDriver::Graphics::ShaderPath::Vertex, 0u, 64u, 64u, 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() {
+    for (std::uint32_t pixel = 0; pixel < Width * Height; ++pixel) {
+        const auto actual = std::to_integer<unsigned>(Pixels[pixel * 4u]);
+        Require(actual == 255u, "NGG vertex wave reduction: pixel " + std::to_string(pixel) + " is " + std::to_string(actual) + ", expected 255 (a vertex whose wave OR of its lanes' bits, read from lanes 31 and 63, differs from its wave's ballot collapses its triangle)");
+    }
+}
+
+}
+
+int main() {
+    try {
+        const auto device = OpenVulkanTestDevice();
+        if (!device) return VulkanTestSkipped;
+        const auto subgroupSize = device->Target().subgroupSize;
+        if (subgroupSize != 32u) {
+            std::printf("skipped, s3 describes a 32-lane wave only on 32-lane subgroups, and the device has %u lanes\n", subgroupSize);
+            return VulkanTestSkipped;
+        }
+        if ((device->SubgroupStages() & VK_SHADER_STAGE_VERTEX_BIT) == 0u || !TargetHasCapability(device->Target(), spv::CapabilityGroupNonUniformArithmetic)) {
+            std::puts("skipped, the device has no subgroup arithmetic in vertex shaders");
+            return VulkanTestSkipped;
+        }
+        for (std::uint32_t pass = 0; pass < 3; ++pass) {
+            Draw(*device);
+            Check();
+        }
+        std::puts("NGG vertex wave reduction tests passed");
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}

core/shader/recompiler/Optimization/include/Optimization/ShaderStageInputInfo.hpp

@@ -129,6 +129,7 @@ struct ShaderVertexInputInfo {
     int fetchBufferReg = 0;
     int buffersNum = 0;
     std::uint32_t scratchSizeDwords = 0;
+    std::uint32_t hostSubgroupSize = 64;
     std::uint32_t paClVsOutCntl = 0;
     ShaderClipSpaceTransform clipSpace;
     ShaderMeshInputInfo mesh;