fix(shader): mask writes after barrier-safe BDA faults

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

@@ -8,6 +8,9 @@
 #define NOMINMAX
 #endif
 #include <windows.h>
+#include <io.h>
+#else
+#include <unistd.h>
 #endif
 #include <algorithm>
 #include <array>
@@ -34,9 +37,9 @@ constexpr std::uint32_t OutputBase = 0x8000;
 
 alignas(256) std::array<std::array<std::uint32_t, 64>, 4> Programs{};
 
-std::span<const std::uint32_t> Code(bool coherent, bool barrier) {
+std::span<const std::uint32_t> Code(bool coherent, bool barrier, std::uint32_t loadBase = LoadBase) {
     auto& program = Programs[(coherent ? 2u : 0u) + (barrier ? 1u : 0u)];
-    std::vector<std::uint32_t> code{0x34040084, 0x4a0404ff, LoadBase, 0x34080081, 0x36080882, 0x4a040504, 0x34060083, 0x4a0606ff, OutputBase, 0xdc388ff8 | (coherent ? 0x10000u : 0u), 0x04000002, 0xbf8c3f70};
+    std::vector<std::uint32_t> code{0x34040084, 0x4a0404ff, loadBase, 0x34080081, 0x36080882, 0x4a040504, 0x34060083, 0x4a0606ff, OutputBase, 0xdc388ff8 | (coherent ? 0x10000u : 0u), 0x04000002, 0xbf8c3f70};
     if (barrier) code.push_back(0xbf8a0000);
     code.insert(code.end(), {0xdc708000, 0x00000503, 0xdc708004, 0x00000703, 0xbf810000});
     std::copy(code.begin(), code.end(), program.begin());
@@ -118,6 +121,86 @@ ShaderRecompiler::RecompileResult Recompile(std::span<const std::uint32_t> code,
     return ShaderRecompiler::Recompile(request);
 }
 
+class StderrCapture {
+public:
+    StderrCapture() {
+        std::fflush(stderr);
+        file = std::tmpfile();
+        Require(file != nullptr, "BDA span reads: cannot capture the fault report");
+#ifdef _WIN32
+        saved = _dup(_fileno(stderr));
+        _dup2(_fileno(file), _fileno(stderr));
+#else
+        saved = dup(fileno(stderr));
+        dup2(fileno(file), fileno(stderr));
+#endif
+    }
+    std::string Finish() {
+        std::fflush(stderr);
+#ifdef _WIN32
+        _dup2(saved, _fileno(stderr));
+        _close(saved);
+#else
+        dup2(saved, fileno(stderr));
+        close(saved);
+#endif
+        std::rewind(file);
+        std::string text;
+        std::array<char, 4096> buffer{};
+        for (std::size_t read = 0; (read = std::fread(buffer.data(), 1, buffer.size(), file)) != 0;) text.append(buffer.data(), read);
+        std::fclose(file);
+        std::fwrite(text.data(), 1, text.size(), stderr);
+        return text;
+    }
+private:
+    std::FILE* file = nullptr;
+    int saved = -1;
+};
+
+bool RunFault(AgcDriver::VulkanDevice& device, GuestBlock& guest, bool coherent, bool barrier, std::uint32_t waveSize, const ShaderRecompiler::SpirvTarget& target, const std::string& name, bool mixed = false) {
+    const auto loadBase = static_cast<std::uint32_t>(mixed ? BlockBytes - 512u : BlockBytes + LoadBase);
+    const auto code = Code(coherent, barrier, loadBase);
+    auto memory = Initial();
+    if (mixed) {
+        for (std::size_t dword = memory.size() - 128u; dword < memory.size(); ++dword) memory[dword] = 0x51000000u + static_cast<std::uint32_t>(dword) * 0x00010203u;
+    }
+    std::memcpy(guest.Data(), memory.data(), BlockBytes);
+    auto expected = memory;
+    if (mixed) {
+        for (std::uint32_t tid = 0u; tid < Threads; ++tid) {
+            const auto load = loadBase + tid * 16u + (tid & 1u) * 2u - 8u;
+            if (load + 16u > BlockBytes) continue;
+            expected[OutputBase / 4u + tid * 2u] = Bytes(memory, load + 4u);
+            expected[OutputBase / 4u + tid * 2u + 1u] = Bytes(memory, load + 12u);
+        }
+        Require(expected != memory, "BDA span reads: the mixed control has no distinguishable valid lane writes");
+    }
+    const auto address = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(guest.Data()));
+    Require(!GuestAllocations::Mutation().Overlaps(reinterpret_cast<const void*>(address + BlockBytes), Threads * 16u + LoadBase + 16u), "BDA span reads: the fault input overlaps a registered allocation");
+    const std::array<std::uint32_t, 2> userData{static_cast<std::uint32_t>(address), static_cast<std::uint32_t>(address >> 32u)};
+    const auto result = Recompile(code, waveSize, userData, target);
+    Require(std::ranges::any_of(result.bindings, [](const auto& binding) { return binding.role == ShaderRecompiler::DescriptorRole::BdaPagetable; }), "BDA span reads: the fault control has no BDA table");
+    std::string failure;
+    StderrCapture log;
+    try {
+        device.Dispatch(result, 1, 1, 1, {}, reinterpret_cast<std::uintptr_t>(code.data()));
+        device.WaitIdle();
+    } catch (const std::exception& error) {
+        failure = error.what();
+    }
+    const auto deferred = log.Finish();
+    if (failure.empty()) failure = deferred;
+    Require(failure.find("BDA access failed:") != std::string::npos && failure.find("reason=1") != std::string::npos, "BDA span reads: no unmapped fault report: " + failure);
+    std::uint32_t changed = 0u;
+    for (std::size_t dword = 0; dword < memory.size(); ++dword) {
+        std::uint32_t actual = 0u;
+        std::memcpy(&actual, guest.Data() + dword * 4u, sizeof(actual));
+        if (actual != expected[dword]) ++changed;
+    }
+    std::printf("BDA fault span %s%s coherent=%u barrier=%u changed=%u refused=1\n", name.c_str(), mixed ? "-mixed" : "", coherent ? 1u : 0u, barrier ? 1u : 0u, changed);
+    return changed != 0u;
+}
+
 void Run(AgcDriver::VulkanDevice& device, GuestBlock& guest, bool coherent, bool barrier, std::uint32_t waveSize, const ShaderRecompiler::SpirvTarget& target, const std::string& name) {
     const auto code = Code(coherent, barrier);
     auto memory = Initial();
@@ -143,13 +226,24 @@ int main() {
         GuestBlock guest;
         const auto native = device->Target();
         const auto split = device->ComputeTarget(32);
+        std::uint32_t faultingWrites = 0u;
         for (const bool coherent : {false, true}) {
             for (const bool barrier : {false, true}) {
                 Run(*device, guest, coherent, barrier, 32, native, "wave32");
                 Run(*device, guest, coherent, barrier, 64, native, "wave64");
                 if (split.subgroupSize != native.subgroupSize) Run(*device, guest, coherent, barrier, 64, split, "wave64 split");
+                faultingWrites += RunFault(*device, guest, coherent, barrier, 32, native, "wave32") ? 1u : 0u;
+                faultingWrites += RunFault(*device, guest, coherent, barrier, 64, native, "wave64") ? 1u : 0u;
+                if (split.subgroupSize != native.subgroupSize) faultingWrites += RunFault(*device, guest, coherent, barrier, 64, split, "wave64-split") ? 1u : 0u;
+                if (barrier && native.subgroupSize == 32u) {
+                    faultingWrites += RunFault(*device, guest, coherent, barrier, 32, native, "wave32", true) ? 1u : 0u;
+                    faultingWrites += RunFault(*device, guest, coherent, barrier, 64, native, "wave64", true) ? 1u : 0u;
+                    if (split.subgroupSize != native.subgroupSize) faultingWrites += RunFault(*device, guest, coherent, barrier, 64, split, "wave64-split", true) ? 1u : 0u;
+                }
             }
         }
+        if (native.subgroupSize != 32u) std::printf("skipped, mixed-lane barrier BDA fault controls need subgroup size 32, device reports %u\n", native.subgroupSize);
+        Require(faultingWrites == 0u, "BDA span reads: " + std::to_string(faultingWrites) + " faulting programs still wrote guest memory");
         std::puts("BDA span read tests passed");
         return 0;
     } catch (const std::exception& error) {

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/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

@@ -136,6 +136,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;
@@ -199,6 +200,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,7 +107,34 @@ 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.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;
+    }
     if (!state.bdaStopsInvocations) return;
     const auto failed = state.module.AllocateId();
     const auto next = state.module.AllocateId();