perf(recompiler): emit wave-uniform values once for both lanes of a two-lane invocation

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

@@ -5,13 +5,15 @@
 #include "Optimization/ResourceProgram.hpp"
 #include "Optimization/ShaderStageInputInfo.hpp"
 #include "Optimization/SrtWalker/SrtFlatSlotClasses.hpp"
+#include "SpirvBackend/SpirvAnalysis.hpp"
 #if ANYPS5_ENABLE_SPIRV_TOOLS
 #include "SpirvBackend/SpirvOptimizer.hpp"
 #endif
 #include "CacheKey.hpp"
 #include <spirv/unified1/spirv.hpp>
 #include <algorithm>
 #include <array>
+#include <cstring>
 #include <initializer_list>
 #include <iostream>
 #include <map>
@@ -663,6 +665,94 @@ void verifyComputedTexelOffsets() {
     }
 }
 
+void verifyWaveUniformValues() {
+    using namespace ShaderRecompiler;
+    IrProgram program;
+    program.Resources().stage = IrShaderStage::Compute;
+    MemoryInfo scalar;
+    scalar.kind = ResourceKind::ScalarAddress;
+    MemoryInfo global;
+    global.kind = ResourceKind::Global;
+    program.Resources().memoryInfo = {scalar, global};
+    auto& block = program.CreateBlock();
+    program.SetEntryBlock(block);
+    program.BlockOrder().push_back(&block);
+    const auto emit = [&](IrOpcode opcode, IrType type, std::initializer_list<IrValue*> arguments, std::uint64_t flags = 0) -> IrValue& {
+        auto& value = program.CreateValue(opcode, type, flags);
+        for (auto* argument : arguments) value.AddArgument(argument);
+        block.AppendInstruction(&value);
+        return value;
+    };
+    const auto memory = [](std::uint32_t index) {
+        MemoryFlags flags{index, 0u};
+        std::uint64_t bits = 0;
+        std::memcpy(&bits, &flags, sizeof(flags));
+        return bits;
+    };
+    auto& zero = program.CreateValue(IrOpcode::Void, IrType::U32);
+    zero.SetImmediateU32(0u);
+    auto& active = program.CreateValue(IrOpcode::Void, IrType::U1);
+    active.SetImmediateBool(true);
+    auto& userData = emit(IrOpcode::GetUserData, IrType::U32, {&zero});
+    auto& lane = emit(IrOpcode::LaneId, IrType::U32, {});
+    auto& uniformSum = emit(IrOpcode::IAdd32, IrType::U32, {&userData, &userData});
+    auto& laneSum = emit(IrOpcode::IAdd32, IrType::U32, {&userData, &lane});
+    auto& address = emit(IrOpcode::GetAddressResource, IrType::AddressResource, {&userData, &userData});
+    auto& scalarLoad = emit(IrOpcode::LoadAddressU32, IrType::U32, {&address, &uniformSum, &zero, &active}, memory(0u));
+    auto& laneOffsetLoad = emit(IrOpcode::LoadAddressU32, IrType::U32, {&address, &laneSum, &zero, &active}, memory(0u));
+    auto& globalLoad = emit(IrOpcode::LoadAddressU32, IrType::U32, {&address, &uniformSum, &zero, &active}, memory(1u));
+    auto& fromScalarLoad = emit(IrOpcode::IAdd32, IrType::U32, {&scalarLoad, &uniformSum});
+    auto& fromGlobalLoad = emit(IrOpcode::IAdd32, IrType::U32, {&globalLoad, &uniformSum});
+    auto& compare = emit(IrOpcode::ULessThan32, IrType::U1, {&laneSum, &userData});
+    auto& ballot = emit(IrOpcode::Ballot, IrType::U32x4, {&compare});
+    const auto uniform = WaveUniformValues(program);
+    for (const auto* value : {&userData, &uniformSum, &address, &scalarLoad, &fromScalarLoad, &ballot}) {
+        require(uniform.contains(value), "wave-uniform values: a value every lane of the wave computes alike was not found uniform");
+    }
+    for (const auto* value : {&lane, &laneSum, &laneOffsetLoad, &globalLoad, &fromGlobalLoad, &compare}) {
+        require(!uniform.contains(value), "wave-uniform values: a value that may differ between lanes was found uniform");
+    }
+}
+
+void verifyTwoLaneUniformValues() {
+    using namespace ShaderRecompiler;
+    static std::array<std::uint32_t, 64> output{};
+    const auto outputBase = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(output.data()));
+    static std::array<std::uint32_t, 8> srt{};
+    srt = {6u, 7u, 0u, 0u, static_cast<std::uint32_t>(outputBase), static_cast<std::uint32_t>((outputBase >> 32u) & 0xffffu), 64u, 0xfacu};
+    const auto srtAddress = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(srt.data()));
+    const std::array<std::uint32_t, 2> userData{static_cast<std::uint32_t>(srtAddress), static_cast<std::uint32_t>(srtAddress >> 32u)};
+    const std::array<std::uint32_t, 10> code{0xf4040080u, 0xfa000000u, 0xf4080200u, 0xfa000010u, 0xbf8cc07fu, 0x93040302u, 0x4a020004u, 0xe0700000u, 0x80020100u, 0xbf810000u};
+    const auto multiplies = [&](std::uint32_t subgroupSize) {
+        RecompileRequest request{};
+        request.shader = {ShaderStage::Compute, 0x40000u, code, 0, {}};
+        request.context.waveSize = 64;
+        request.context.userDataBaseRegister = 0;
+        request.context.userData = userData;
+        request.context.compute = ShaderComputeStageInfo{{64u, 1u, 1u}, 0u, {false, false, false}, false, 1u};
+        request.target.vulkanVersion = 0x00401000u;
+        request.target.spirvVersion = 0x00010300u;
+        request.target.subgroupSize = subgroupSize;
+        request.target.fragmentShaderBarycentricEnabled = false;
+        request.layout.pushConstantSizeBytes = 128;
+        AgcDriver::ShaderMemory memory({});
+        const auto capture = memory.Capture(request);
+        request.context.memory = memory.Regions();
+        const auto words = Recompile(request, *capture)->spirv;
+        std::size_t count = 0;
+        for (std::size_t cursor = 5; cursor < words.size();) {
+            const auto length = words[cursor] >> 16u;
+            require(length != 0 && length <= words.size() - cursor, "two-lane uniform values: truncated SPIR-V instruction");
+            if ((words[cursor] & 0xffffu) == spv::OpIMul) ++count;
+            cursor += length;
+        }
+        return count;
+    };
+    const auto oneLane = multiplies(64u);
+    require(oneLane != 0u, "two-lane uniform values: the scalar multiply is missing from the module");
+    require(multiplies(32u) == oneLane, "two-lane uniform values: a two-lane invocation computes a scalar value once per lane");
+}
+
 int main() {
     try {
         using namespace ShaderRecompiler;
@@ -674,6 +764,8 @@ int main() {
         verifyPixelInputs();
         verifyPixelParameterSlots();
         verifyComputedTexelOffsets();
+        verifyWaveUniformValues();
+        verifyTwoLaneUniformValues();
 #if ANYPS5_ENABLE_SPIRV_TOOLS
         const std::vector<std::uint32_t> minimalSpirv{
             0x07230203u, 0x00010000u, 0u, 5u, 0u,

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

@@ -0,0 +1,88 @@
+#include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/Draw.hpp"
+#include "Recompiler.hpp"
+#include "VulkanTestDevice.hpp"
+#include <array>
+#include <cstdio>
+#include <iostream>
+#include <string>
+#include <vector>
+
+namespace {
+
+using AgcDriver::Graphics::Require;
+using ShaderRecompiler::ShaderStage;
+
+constexpr std::uint32_t Threads = 64;
+alignas(256) std::array<std::uint32_t, 4> Input{3u, 5u, 0u, 0u};
+alignas(256) std::array<std::uint32_t, Threads> Output{};
+
+alignas(256) constexpr std::array<std::uint32_t, 22> WaveCode{
+    0xf4240200, 0xfa000000, 0xbf8cc07f, 0xbe8a0380, 0xbe8b0380, 0x800a080a, 0x930c0b0a, 0x800a0c0a,
+    0x800b810b, 0xbf0a090b, 0xbf85fffa, 0x7d880094, 0xbe8d106a, 0x4a060087, 0x7e1c0503, 0xd5430001,
+    0x00281b00, 0x4a02020e, 0x34040082, 0xe0701000, 0x80010102, 0xbf810000,
+};
+
+std::array<std::uint32_t, 4> BufferDescriptor(const void* data, 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), count * 4u, 0x01016facu};
+}
+
+std::uint32_t Accumulated() {
+    std::uint32_t sum = 0;
+    for (std::uint32_t count = 0; count < Input[1]; ++count) {
+        sum += Input[0];
+        sum += sum * count;
+    }
+    return sum;
+}
+
+void Run(AgcDriver::VulkanDevice& device) {
+    Output.fill(0xdeadbeefu);
+    std::vector<std::uint32_t> userData(8, 0u);
+    const auto input = BufferDescriptor(Input.data(), static_cast<std::uint32_t>(Input.size()));
+    const auto output = BufferDescriptor(Output.data(), static_cast<std::uint32_t>(Output.size()));
+    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(WaveCode);
+    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, {false, false, false}, false, 1};
+    ShaderRecompiler::RecompileRequest request{
+        {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}
+    };
+    request.useCache = false;
+    const auto result = ShaderRecompiler::Recompile(request);
+    device.Dispatch(result, 1, 1, 1, {}, reinterpret_cast<std::uintptr_t>(code.data()));
+    device.WaitIdle();
+}
+
+void Check() {
+    const auto uniform = Accumulated() + 7u;
+    for (std::uint32_t tid = 0; tid < Threads; ++tid) {
+        const auto expected = tid * 20u + uniform;
+        Require(Output[tid] == expected, "wave uniform: thread " + std::to_string(tid) + " stored " + std::to_string(Output[tid]) + ", expected " + std::to_string(expected));
+    }
+}
+
+}
+
+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 wave64 in two lanes\n", device->Target().subgroupSize);
+            return VulkanTestSkipped;
+        }
+        Run(*device);
+        Check();
+        std::puts("wave uniform tests passed");
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}

core/shader/recompiler/SpirvBackend/include/SpirvBackend/SpirvAnalysis.hpp

@@ -4,6 +4,7 @@
 #include "IntermediateRepresentation/IrProgram.hpp"
 #include <cstdint>
 #include <string>
+#include <unordered_set>
 #include <vector>
 
 namespace ShaderRecompiler {
@@ -24,6 +25,7 @@ struct SpirvRequirements {
 };
 
 [[nodiscard]] SpirvRequirements AnalyzeProgramRequirements(const IrProgram& program);
+[[nodiscard]] std::unordered_set<const IrValue*> WaveUniformValues(const IrProgram& program);
 
 }
 

core/shader/recompiler/SpirvBackend/include/SpirvBackend/SpirvEmitterState.hpp

@@ -10,6 +10,7 @@
 #include <span>
 #include <string_view>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -81,6 +82,7 @@ struct SpirvEmitterState {
 
     SpirvRequirements requirements;
     std::uint32_t laneCount = 1;
+    std::unordered_set<const IrValue*> sharedLaneValues;
     std::uint32_t laneHalf = 0;
     // The target's SPIR-V version and what the device accepts, for capabilities an emitter adds
     // only when needed (bindless image tables: see TableImageIndex).

core/shader/recompiler/SpirvBackend/src/SpirvAnalysis.cpp

@@ -1,8 +1,66 @@
 #include "SpirvBackend/SpirvAnalysis.hpp"
 #include <stdexcept>
+#include <unordered_map>
 
 namespace ShaderRecompiler {
 
+namespace {
+
+bool UniformSource(const IrValue& value) {
+    switch (value.Opcode()) {
+    case IrOpcode::Ballot:
+    case IrOpcode::ReadFirstLane:
+    case IrOpcode::ReadLane:
+    case IrOpcode::GetUserData:
+    case IrOpcode::GetShaderBase:
+    case IrOpcode::ReadConst:
+    case IrOpcode::ReadConstBuffer:
+    case IrOpcode::MeshDrawParameter:
+    case IrOpcode::GetSrtResource:
+    case IrOpcode::GetBufferResource:
+    case IrOpcode::GetAddressResource:
+    case IrOpcode::GetScratchResource:
+    case IrOpcode::GetImageResource:
+    case IrOpcode::GetSamplerResource:
+        return true;
+    default:
+        return false;
+    }
+}
+
+bool LaneSource(const IrProgram& program, const IrValue& value) {
+    switch (value.Opcode()) {
+    case IrOpcode::LaneId:
+    case IrOpcode::GetAttribute:
+    case IrOpcode::GetInterpolationParameter:
+    case IrOpcode::GetTessellationAttribute:
+    case IrOpcode::DppMoveU32:
+    case IrOpcode::DppUpdateU32:
+    case IrOpcode::Permlane16U32:
+    case IrOpcode::BpermuteU32:
+    case IrOpcode::SwizzleU32:
+    case IrOpcode::WriteLane:
+    case IrOpcode::DataAppend:
+    case IrOpcode::DataConsume:
+        return true;
+    case IrOpcode::GetBuiltin: {
+        const auto* kind = value.Argument(0)->Resolve();
+        return kind == nullptr || !kind->HasImmediate() || static_cast<StageInputKind>(kind->ImmediateU32()) != StageInputKind::WorkgroupId;
+    }
+    default:
+        break;
+    }
+    const auto address = AddressOpcodeInfoOf(value.Opcode()).access;
+    if (address == AddressAccess::Read) {
+        const auto index = value.Flags<MemoryFlags>().index;
+        const auto& memory = program.Resources().memoryInfo;
+        if (index < memory.size() && memory[index].kind == ResourceKind::ScalarAddress) return false;
+    }
+    return address != AddressAccess::None || BufferAccessOf(value.Opcode()) != BufferAccess::None || SharedAccessOf(value.Opcode()) != SharedAccess::None || ImageOpcodeInfoOf(value.Opcode()).access != ImageAccess::None;
+}
+
+}
+
 SpirvRequirements AnalyzeProgramRequirements(const IrProgram& program) {
     SpirvRequirements requirements {};
     for (const IrBlock* block : program.BlockOrder()) {
@@ -128,4 +186,36 @@ SpirvRequirements AnalyzeProgramRequirements(const IrProgram& program) {
     return requirements;
 }
 
+std::unordered_set<const IrValue*> WaveUniformValues(const IrProgram& program) {
+    std::unordered_map<const IrValue*, bool> varying;
+    const auto isVarying = [&](const IrValue* value) {
+        const auto* resolved = value->Resolve();
+        if (resolved == nullptr || resolved->HasImmediate()) return false;
+        const auto found = varying.find(resolved);
+        return found != varying.end() && found->second;
+    };
+    bool changed = true;
+    while (changed) {
+        changed = false;
+        for (const IrBlock* block : program.BlockOrder()) {
+            for (const IrValue* inst : block->Instructions()) {
+                bool lanes = !UniformSource(*inst) && LaneSource(program, *inst);
+                for (std::size_t index = 0; !lanes && !UniformSource(*inst) && index < inst->ArgumentCount(); index++) {
+                    if (inst->Argument(index) != nullptr) lanes = isVarying(inst->Argument(index));
+                }
+                auto& entry = varying[inst];
+                if (lanes && !entry) {
+                    entry = true;
+                    changed = true;
+                }
+            }
+        }
+    }
+    std::unordered_set<const IrValue*> uniform;
+    for (const auto& [value, lanes] : varying) {
+        if (!lanes) uniform.insert(value);
+    }
+    return uniform;
+}
+
 }

core/shader/recompiler/SpirvBackend/src/SpirvEmitter.cpp

@@ -249,6 +249,7 @@ std::vector<std::uint32_t> SpirvEmitter::Emit(const IrProgram& program, const Sh
     state.nonConstantImageOffsets = target.nonConstantImageOffsets;
     const auto* workgroup = ShaderWorkgroupInputFor(state);
     state.laneCount = workgroup != nullptr && program.WaveSize() == 64u && workgroup->hostSubgroupSize == 32u ? 2u : 1u;
+    if (state.laneCount == 2u) state.sharedLaneValues = WaveUniformValues(program);
     if (program.Resources().stage == IrShaderStage::Compute && workgroup != nullptr) {
         // The key comes from a subgroup ballot (ReadFirstLane), so the slot is uniform over the
         // workgroup only when the workgroup is one wave held by one host subgroup; a wave64 program

core/shader/recompiler/SpirvBackend/src/SpirvFlowEmitter.cpp

@@ -571,10 +571,17 @@ void EmitStructuredBlock(SpirvValueEmitContext& ctx, StructuredFunctionState& fu
                 ldsRead |= access != SharedAccess::Write;
             }
         }
+        const bool shared = state.laneCount == 2u && inst->Type() != IrType::Void && !IrOpcodeHasSideEffects(inst->Opcode()) && state.sharedLaneValues.contains(inst);
         for (std::uint32_t half = 0; half < state.laneCount; half++) {
             if (half != 0 && ctx.otherHalf == nullptr) {
                 ctx.Fail(*inst, "requires a second lane context");
             }
+            if (half != 0 && shared) {
+                if (const auto found = ctx.definitions.find(inst); found != ctx.definitions.end()) {
+                    ctx.otherHalf->Define(*inst, found->second);
+                    continue;
+                }
+            }
             SpirvValueEmitContext& lane = half == 0 ? ctx : *ctx.otherHalf;
             state.laneHalf = half;
             if (half == 0 || (inst->Opcode() != IrOpcode::Barrier && inst->Opcode() != IrOpcode::MeshAllocate)) {