fix(shader): keep faulting invocations in their wave and mask their later writes

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

@@ -9817,6 +9817,7 @@ void verifyGuardedNullPointers() {
         std::string function;
         std::size_t calls = 0;
         std::size_t stores = 0;
+        std::size_t latches = 0;
         for (std::size_t cursor = 5; cursor < words.size();) {
             const auto count = words[cursor] >> 16u;
             const auto opcode = words[cursor] & 0xffffu;
@@ -9827,9 +9828,13 @@ void verifyGuardedNullPointers() {
             if (opcode == spv::OpAccessChain && count == 6u && names[words[cursor + 3]] == "fault_buffer" && constants.contains(words[cursor + 5]) && constants[words[cursor + 5]] == BdaAbi::SrtGuardWord) guardWords.push_back(words[cursor + 2]);
             if (function == "main" && opcode == spv::OpFunctionCall && names[words[cursor + 3]] == "record_bda_fault") ++calls;
             if (function == "main" && opcode == spv::OpAtomicStore && std::find(guardWords.begin(), guardWords.end(), words[cursor + 1]) != guardWords.end()) ++stores;
+            if (function == "main" && opcode == spv::OpAtomicCompareExchange && std::find(guardWords.begin(), guardWords.end(), words[cursor + 3]) != guardWords.end()) {
+                require(count == 9u && constants.contains(words[cursor + 8]) && constants[words[cursor + 8]] == 0u, "guarded pointer: the first-fault latch does not compare against zero");
+                ++latches;
+            }
             cursor += count;
         }
-        return std::pair{calls, stores};
+        return std::array{calls, stores, latches};
     };
 
     NestedRequest branch(branchCode, root.data());
@@ -9842,8 +9847,8 @@ void verifyGuardedNullPointers() {
     branch.request.context.memory = validRegions;
     const auto validResult = Recompile(branch.request, *valid);
     require(validResult->poisonedSrtReads == 0u && faultBinding(*validResult) && validResult->bdaAbiVersion == BdaAbi::Version, "guarded pointer: a guarded program has no fault buffer");
-    const auto [branchCalls, branchStores] = guardSites(*validResult);
-    require(branchCalls == 0u && branchStores == 1u, ("guarded pointer: a guarded read calls record_bda_fault " + std::to_string(branchCalls) + " times in main and stores its record to the fault buffer's guard word " + std::to_string(branchStores) + " times, expected one store and no call").c_str());
+    const auto [branchCalls, branchStores, branchLatches] = guardSites(*validResult);
+    require(branchCalls == 0u && branchStores == 0u && branchLatches == 1u, "guarded pointer: a continuing guarded read must latch its first fault without overwriting the guard word or calling record_bda_fault");
 
     for (const auto pointer : {std::uint64_t{0}, unmapped}) {
         for (const auto condition : {1u, 0u}) {
@@ -9929,8 +9934,8 @@ void verifyGuardedNullPointers() {
     chain.request.context.memory = mappedChainRegions;
     const auto mappedChainResult = Recompile(chain.request, *mappedChain);
     require(chainResult->poisonedSrtReads == 3u && nullChainResult->poisonedSrtReads == 1u && mappedChainResult->poisonedSrtReads == 0u, "guarded pointer: a chain result does not count its poisoned reads");
-    const auto [chainCalls, chainStores] = guardSites(*mappedChainResult);
-    require(chainCalls == 0u && chainStores == mappedChain->plan->guardedSrtSlots.size(), ("guarded pointer: a chain's guarded reads call record_bda_fault " + std::to_string(chainCalls) + " times in main and store to the guard word " + std::to_string(chainStores) + " times, expected one store per guard and no call").c_str());
+    const auto [chainCalls, chainStores, chainLatches] = guardSites(*mappedChainResult);
+    require(chainCalls == 0u && chainStores == 0u && chainLatches == mappedChain->plan->guardedSrtSlots.size(), "guarded pointer: a continuing chain must latch once per guard without overwriting the guard word or calling record_bda_fault");
     require(samePipeline(*chainResult, *mappedChainResult) && samePipeline(*nullChainResult, *mappedChainResult), "guarded pointer: a chain's poison changed the variant, the specialization or the module");
 
     const auto poisonedWords = [](const ResourceCapture& capture, std::uint32_t pc, std::uint64_t base) {

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

@@ -0,0 +1,123 @@
+#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 <algorithm>
+#include <array>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <iostream>
+#include <span>
+#include <string>
+#include <vector>
+
+namespace {
+
+using AgcDriver::Graphics::Require;
+using ShaderRecompiler::ShaderStage;
+
+constexpr std::uint32_t Threads = 64;
+constexpr std::uint32_t FaultingLane = 5;
+constexpr std::uint32_t PassLimit = 100;
+constexpr std::uint32_t CountOffset = 0x200;
+constexpr std::uint32_t ValueOffset = 0x300;
+constexpr std::uint32_t Untouched = 0xdeadbeefu;
+constexpr std::size_t BlockBytes = 65536;
+
+alignas(256) constexpr std::array<std::uint32_t, 22> Code{
+    0x34020082u, 0x7d8a0085u, 0x020202ffu, 0x00010000u, 0xdc308000u, 0x02000001u, 0xbf8c3f70u, 0xbe8a0380u,
+    0x7da2040au, 0xbf880004u, 0x800a810au, 0xbf0aff0au, PassLimit, 0xbf85fffau, 0xbefe04c1u, 0x7e06020au,
+    0x34080082u, 0xdc708200u, 0x00000304u, 0xdc708300u, 0x00000204u, 0xbf810000u,
+};
+
+class GuestBlock {
+public:
+    GuestBlock() {
+#ifdef _WIN32
+        block = static_cast<std::uint32_t*>(VirtualAlloc(nullptr, BlockBytes, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE));
+#else
+        block = static_cast<std::uint32_t*>(std::aligned_alloc(BlockBytes, BlockBytes));
+#endif
+        Require(block != nullptr, "faulting wave invocation: cannot allocate the guest block");
+        GuestAllocations::Mutation().Add(block, BlockBytes, 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;
+
+    std::uint32_t& At(std::uint32_t offset) { return block[offset / 4u]; }
+    std::uint64_t Address() const { return reinterpret_cast<std::uintptr_t>(block); }
+
+private:
+    std::uint32_t* block = nullptr;
+};
+
+std::uint32_t Passes(std::uint32_t lane) {
+    return lane % 4u + 1u;
+}
+
+}
+
+int main() {
+    try {
+        const auto device = OpenVulkanTestDevice();
+        if (!device) return VulkanTestSkipped;
+        GuestBlock guest;
+        for (std::uint32_t lane = 0; lane < Threads; ++lane) {
+            guest.At(lane * 4u) = Passes(lane);
+            guest.At(CountOffset + lane * 4u) = Untouched;
+            guest.At(ValueOffset + lane * 4u) = Untouched;
+        }
+        const std::vector<std::uint32_t> userData{static_cast<std::uint32_t>(guest.Address()), static_cast<std::uint32_t>(guest.Address() >> 32u)};
+        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, {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();
+        std::uint32_t passes = 0;
+        for (std::uint32_t lane = 0; lane < Threads; ++lane) {
+            if (lane != FaultingLane) passes = std::max(passes, Passes(lane));
+        }
+        for (std::uint32_t lane = 0; lane < Threads; ++lane) {
+            const auto what = "faulting wave invocation: lane " + std::to_string(lane);
+            const auto count = guest.At(CountOffset + lane * 4u);
+            const auto value = guest.At(ValueOffset + lane * 4u);
+            if (lane == FaultingLane) {
+                Require(count == Untouched && value == Untouched, what + " stored after its read faulted");
+                continue;
+            }
+            Require(value == Passes(lane), what + " stored " + std::to_string(value) + ", expected its loaded " + std::to_string(Passes(lane)));
+            Require(count == passes, what + " left the EXEC loop after " + std::to_string(count) + " passes, expected " + std::to_string(passes) + " (the loop stops at " + std::to_string(PassLimit) + ")");
+        }
+        std::puts("faulting wave invocation tests passed");
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}

core/shader/recompiler/ShaderDiskCache.hpp

@@ -13,7 +13,7 @@
 
 namespace ShaderRecompiler::ShaderDiskCache {
 
-inline constexpr std::uint32_t FormatVersion = 51;
+inline constexpr std::uint32_t FormatVersion = 52;
 
 enum class LoadStatus {
     Loaded,

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

@@ -25,6 +25,7 @@ void ValidateBdaTarget(const IrProgram& program, const SpirvTargetOptions& targe
 // every invocation running, so their faulting reads return zero instead.
 bool BdaInvocationsMayStop(const IrProgram& program);
 void StopBdaInvocationIf(SpirvEmitterState& state, std::uint32_t condition);
+std::uint32_t BdaInvocationFaulted(SpirvEmitterState& state);
 std::uint32_t BdaInstructionPc(SpirvEmitterState& state, const IrValue& inst);
 std::uint32_t EmitBdaRead(SpirvValueEmitContext& ctx, const IrValue& inst, std::uint32_t address, std::uint32_t bits);
 std::uint32_t EmitBdaReadBytes(SpirvValueEmitContext& ctx, const IrValue& inst, std::uint32_t address, std::uint32_t bytes);

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

@@ -134,6 +134,7 @@ struct SpirvEmitterState {
     bool ldsLockFences = false;
     std::unordered_set<const IrValue*> sharedLaneValues;
     std::uint32_t laneHalf = 0;
+    bool laneValueShared = false;
     std::uint32_t spirvVersion = 0x00010300u;
     std::span<const std::uint32_t> supportedCapabilities;
     std::span<const std::string_view> supportedExtensions;
@@ -196,6 +197,8 @@ struct SpirvEmitterState {
     std::uint32_t bdaPcOverride = 0;
     // False for programs with workgroup barriers: faulting BDA accesses then continue (see BdaInvocationsMayStop).
     bool bdaStopsInvocations = true;
+    bool bdaFaultMasksWrites = false;
+    std::array<std::uint32_t, 2> bdaFaultedVariables {};
     bool continueTarget = false;
     // Execution scope of the barriers that keep one guest wave's LDS accesses in program order across
     // host invocations (see WaveLdsScope); 0 when none are emitted.

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

@@ -107,8 +107,35 @@ void ReturnBdaFailureIf(SpirvEmitterState& state, std::uint32_t condition, std::
     EmitLabel(state, next);
 }
 
+namespace {
+
+std::uint32_t BdaFaultedVariable(SpirvEmitterState& state, std::uint32_t half) {
+    auto& variable = state.bdaFaultedVariables.at(half);
+    if (variable == 0) variable = state.module.DefineInitializedGlobalVariable(TypePointer(state, spv::StorageClassPrivate, TypeBool(state)), spv::StorageClassPrivate, ConstantBool(state, false));
+    return variable;
+}
+
+std::uint32_t LoadBdaFaulted(SpirvEmitterState& state, std::uint32_t half) {
+    const auto faulted = state.module.AllocateId();
+    state.module.AddFunction(spv::OpLoad, TypeBool(state), faulted, BdaFaultedVariable(state, half));
+    return faulted;
+}
+
+}
+
+std::uint32_t BdaInvocationFaulted(SpirvEmitterState& state) {
+    return LoadBdaFaulted(state, state.laneHalf);
+}
+
 void StopBdaInvocationIf(SpirvEmitterState& state, std::uint32_t condition) {
     if (!state.bdaStopsInvocations) return;
+    if (state.bdaFaultMasksWrites && state.bdaStopValue == 0) {
+        for (std::uint32_t half = 0; half < state.laneCount; half++) {
+            if (half != state.laneHalf && !state.laneValueShared) continue;
+            state.module.AddFunction(spv::OpStore, BdaFaultedVariable(state, half), Binary(state, spv::OpLogicalOr, TypeBool(state), LoadBdaFaulted(state, half), condition));
+        }
+        return;
+    }
     const auto failed = state.module.AllocateId();
     const auto next = state.module.AllocateId();
     state.module.AddFunction(spv::OpSelectionMerge, next, spv::SelectionControlMaskNone);

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

@@ -41,8 +41,8 @@ std::uint32_t EmitBdaByte(SpirvEmitterState& state, std::uint32_t guest, std::ui
     const auto bytePointer = TypePointer(state, spv::StorageClassPhysicalStorageBuffer, byteType);
     const auto physical = state.module.AllocateId();
     state.module.AddFunction(spv::OpFunctionCall, TypeScalarU64(state), physical, state.bdaPointerFunction, guest, ConstantU32(state, 1u), instruction);
-    if (state.bdaStopsInvocations) {
-        StopBdaInvocationIf(state, Binary(state, spv::OpIEqual, TypeBool(state), physical, BdaConstant(state, 0u)));
+    if (state.bdaStopsInvocations) StopBdaInvocationIf(state, Binary(state, spv::OpIEqual, TypeBool(state), physical, BdaConstant(state, 0u)));
+    if (state.bdaStopsInvocations && (!state.bdaFaultMasksWrites || state.bdaStopValue != 0)) {
         const auto pointer = state.module.AllocateId();
         state.module.AddFunction(spv::OpConvertUToPtr, bytePointer, pointer, physical);
         const auto loaded = state.module.AllocateId();

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

@@ -328,8 +328,11 @@ std::vector<std::uint32_t> SpirvEmitter::Emit(const IrProgram& program, const Sh
     state.waveLdsScope = WaveLdsScope(program, workgroup, state.laneCount);
     if (const char* guard = std::getenv("APS5_LOOP_GUARD")) state.loopGuardLimit = static_cast<std::uint32_t>(std::strtoul(guard, nullptr, 0));
     // Stopped invocations would leave the wave LDS barriers incomplete.
-    state.bdaStopsInvocations = state.waveLdsScope == 0 && !state.spannedWave && BdaInvocationsMayStop(program) && program.Resources().stage != IrShaderStage::Mesh;
+    state.bdaStopsInvocations = state.waveLdsScope == 0 && BdaInvocationsMayStop(program) && program.Resources().stage != IrShaderStage::Mesh;
     state.meshWindowBuffer = target.meshWindowBuffer && program.Resources().stage == IrShaderStage::Mesh && program.Metadata().vectorWindowSize != 0u && state.laneCount * MeshWindowInvocations(state) * program.Metadata().vectorWindowSize <= MeshWindowSlotDwords;
+    const bool waveOperations = state.laneCount != 1u || state.requirements.subgroupBallot || state.requirements.subgroupShuffle;
+    const bool faults = program.Info().usesDma || program.Info().usesFaultBuffer;
+    state.bdaFaultMasksWrites = state.bdaStopsInvocations && faults && (waveOperations || state.spannedWave) && program.Resources().stage != IrShaderStage::Pixel;
     EmitModuleHeader(state, bindings);
     EmitProgram(state);
     state.module.EmitEntryPoint(ExecutionModelForStage(state.program.Resources().stage), state.mainFunc, "main", state.interfaceVariables);

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

@@ -735,6 +735,41 @@ void EmitStructuredInstruction(SpirvValueEmitContext& ctx, StructuredFunctionSta
     EmitDirectInstruction(ctx, inst);
 }
 
+bool WritesGuestState(IrOpcode opcode) {
+    if (opcode == IrOpcode::SetAttribute || opcode == IrOpcode::SetTessellationAttribute || opcode == IrOpcode::EmitVertex) return true;
+    const auto shared = SharedAccessOf(opcode);
+    const auto buffer = BufferAccessOf(opcode);
+    const auto address = AddressOpcodeInfoOf(opcode).access;
+    const auto image = ImageOpcodeInfoOf(opcode).access;
+    return (shared != SharedAccess::None && shared != SharedAccess::Read) || buffer == BufferAccess::Write || buffer == BufferAccess::Atomic || address == AddressAccess::Write || address == AddressAccess::Atomic || image == ImageAccess::Write || image == ImageAccess::Atomic;
+}
+
+void EmitUnlessInvocationFaulted(SpirvValueEmitContext& ctx, StructuredFunctionState& functionState, const IrValue& inst) {
+    auto& state = ctx.state;
+    const bool returnsValue = inst.Type() != IrType::Void;
+    const auto result = returnsValue ? ctx.Result(inst) : 0u;
+    ctx.definitions.erase(&inst);
+    const auto faulted = BdaInvocationFaulted(state);
+    const auto body = state.module.AllocateId();
+    const auto skip = state.module.AllocateId();
+    const auto merge = state.module.AllocateId();
+    state.module.AddFunction(spv::OpSelectionMerge, merge, spv::SelectionControlMaskNone);
+    state.module.AddFunction(spv::OpBranchConditional, faulted, skip, body);
+    EmitLabel(state, body);
+    EmitStructuredInstruction(ctx, functionState, inst);
+    const auto value = returnsValue ? ctx.Def(&inst) : 0u;
+    const auto exit = state.currentLabel;
+    state.module.AddFunction(spv::OpBranch, merge);
+    EmitLabel(state, skip);
+    state.module.AddFunction(spv::OpBranch, merge);
+    EmitLabel(state, merge);
+    if (!returnsValue) return;
+    ctx.definitions.erase(&inst);
+    const auto type = TypeId(state, inst.Type());
+    state.module.AddFunction(spv::OpPhi, type, result, value, exit, state.module.Constant(spv::OpConstantNull, type), skip);
+    ctx.definitions.emplace(&inst, result);
+}
+
 void EmitStructuredBlock(SpirvValueEmitContext& ctx, StructuredFunctionState& functionState, const IrBlock* block) {
     auto& state = ctx.state;
     state.currentBlock = block;
@@ -783,11 +818,14 @@ void EmitStructuredBlock(SpirvValueEmitContext& ctx, StructuredFunctionState& fu
             }
             SpirvValueEmitContext& lane = half == 0 ? ctx : *ctx.otherHalf;
             state.laneHalf = half;
+            state.laneValueShared = shared;
             if (half == 0 || (inst->Opcode() != IrOpcode::Barrier && inst->Opcode() != IrOpcode::MeshAllocate)) {
-                EmitStructuredInstruction(lane, functionState, *inst);
+                if (state.bdaFaultMasksWrites && WritesGuestState(inst->Opcode())) EmitUnlessInvocationFaulted(lane, functionState, *inst);
+                else EmitStructuredInstruction(lane, functionState, *inst);
             }
         }
         state.laneHalf = 0;
+        state.laneValueShared = false;
     }
 }
 

core/shader/recompiler/SpirvBackend/src/SpirvMemory/SpirvMemoryInstructions.cpp

@@ -1584,7 +1584,7 @@ std::uint32_t EmitReadConst(SpirvValueEmitContext& ctx, const IrValue& inst) {
                 const auto scope = ConstantU32(state, spv::ScopeDevice);
                 const auto relaxed = ConstantU32(state, spv::MemorySemanticsMaskNone);
                 const auto value = Binary(state, spv::OpBitwiseOr, u32, record, stage);
-                if (state.bdaStopsInvocations) {
+                if (state.bdaStopsInvocations && !state.bdaFaultMasksWrites) {
                     state.module.AddFunction(spv::OpAtomicStore, word, scope, relaxed, value);
                 } else {
                     state.module.AddFunction(spv::OpAtomicCompareExchange, u32, state.module.AllocateId(), word, scope, relaxed, relaxed, value, ConstantU32(state, 0u));