fix(shader): widen LDS atomic fences to device scope on AMD's proprietary driver

core/libs/prx/libSceAgcDriver/Execution/include/SubgroupBarrier.hpp

@@ -9,6 +9,10 @@ inline bool SubgroupControlBarrierUsable(VkDriverId driver) {
     return driver != VK_DRIVER_ID_AMD_PROPRIETARY;
 }
 
+inline bool WorkgroupScopeFenceUsable(VkDriverId driver) {
+    return driver != VK_DRIVER_ID_AMD_PROPRIETARY;
+}
+
 }
 
 #endif

core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp

@@ -388,6 +388,7 @@ struct VulkanDevice::State {
     bool earlyFragmentSampleMaskTestBeforeSampleCounting = false;
     bool robustBufferAccess2 = false;
     bool subgroupControlBarriers = true;
+    bool workgroupScopeFences = true;
     // VK_EXT_descriptor_indexing with non-uniform image array indexing (bindless image tables in
     // graphics stages, and compute workgroups wider than a wave).
     bool descriptorIndexing = false;
@@ -1274,6 +1275,7 @@ VulkanDevice::VulkanDevice(const PresentationWindow* window, bool compileOnly) :
         state->individuallyFreeDescriptorSets = driver.driverID == VK_DRIVER_ID_MESA_LLVMPIPE;
         state->fusedExactArithmetic = driver.driverID == VK_DRIVER_ID_MOLTENVK;
         state->subgroupControlBarriers = SubgroupControlBarrierUsable(driver.driverID);
+        state->workgroupScopeFences = WorkgroupScopeFenceUsable(driver.driverID);
     }
     if (driverId == VK_DRIVER_ID_NVIDIA_PROPRIETARY && hasExtension(VK_NV_SHADER_SM_BUILTINS_EXTENSION_NAME)) {
         VkPhysicalDeviceShaderSMBuiltinsPropertiesNV smBuiltins{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_SM_BUILTINS_PROPERTIES_NV};
@@ -3833,6 +3835,7 @@ ShaderRecompiler::SpirvTarget VulkanDevice::buildTarget() const {
     target.waveLdsMemoryBarriers = !state->subgroupControlBarriers;
     target.largePixelFunctionLdsBuffer = state->driverId == VK_DRIVER_ID_NVIDIA_PROPRIETARY;
     target.pixelFunctionStorageDwords = state->pixelFunctionStorageDwords;
+    target.deviceScopeLdsMemoryBarriers = !state->workgroupScopeFences;
     target.srgbDecodeFormats = state->srgbDecodeFormats & ~state->srgbCopyFormats;
     target.vertexSubgroups = (state->subgroup.supportedStages & VK_SHADER_STAGE_VERTEX_BIT) != 0;
     if (state->meshShader) {

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

@@ -750,6 +750,7 @@ void verifyKeySensitivity() {
     changes("spanned waves", [](SampleRequest& sample) { sample.request.target.spannedWaves = true; });
     changes("robust buffer access 2", [](SampleRequest& sample) { sample.request.target.robustBufferAccess2 = true; });
     changes("wave LDS memory barriers", [](SampleRequest& sample) { sample.request.target.waveLdsMemoryBarriers = true; });
+    changes("device-scope LDS memory barriers", [](SampleRequest& sample) { sample.request.target.deviceScopeLdsMemoryBarriers = true; });
     changes("the workgroup size limit", [](SampleRequest& sample) { sample.request.target.maxWorkgroupSize[2] = 128; });
     changes("the invocation limit", [](SampleRequest& sample) { sample.request.target.maxWorkgroupInvocations = 512; });
     changes("the shared memory limit", [](SampleRequest& sample) { sample.request.target.maxWorkgroupSharedMemoryBytes = 32768; });

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

@@ -5,10 +5,12 @@
 #include <cstdint>
 #include <cstdio>
 #include <exception>
+#include <span>
 #include <stdexcept>
 #include <string>
 #include <string_view>
 #include <unordered_map>
+#include <utility>
 
 namespace {
 
@@ -34,6 +36,10 @@ constexpr std::array<std::uint32_t, 19> BranchedLdsStep{
     0xbf8cc07fu, 0xbf8a0000u, 0xbf810000u,
 };
 
+constexpr std::array<std::uint32_t, 5> LdsHalfStore{0x34040082u, 0xd87c0400u, 0x00000102u, 0xbf8cc07fu, 0xbf810000u};
+constexpr std::array<std::uint32_t, 5> LdsAtomicAdd{0x34040082u, 0xd8000400u, 0x00000102u, 0xbf8cc07fu, 0xbf810000u};
+constexpr std::array<std::uint32_t, 5> LdsAtomicMax{0x34040082u, 0xd84c0400u, 0x00000102u, 0xbf8cc07fu, 0xbf810000u};
+
 struct Barriers {
     std::uint32_t subgroupControl = 0;
     std::uint32_t subgroupMemory = 0;
@@ -43,14 +49,19 @@ void require(bool condition, const std::string& message) {
     if (!condition) throw std::runtime_error(message);
 }
 
-Barriers recompile(std::uint32_t threads, bool waveLdsMemoryBarriers) {
+struct MemoryBarriers {
+    std::uint32_t workgroup = 0;
+    std::uint32_t device = 0;
+};
+
+ShaderRecompiler::RecompileResult recompileProgram(std::span<const std::uint32_t> code, std::uint32_t threads, bool waveLdsMemoryBarriers, bool deviceScopeLdsMemoryBarriers) {
     using namespace ShaderRecompiler;
     const std::array<std::uint32_t, 2> capabilities{spv::CapabilityShader, spv::CapabilityGroupNonUniform};
     const std::array<std::string_view, 1> extensions{"SPV_KHR_storage_buffer_storage_class"};
     RecompileRequest request{};
-    request.shader = {ShaderStage::Compute, 0x20000u, BranchedLdsStep, 0, {}};
+    request.shader = {ShaderStage::Compute, 0x20000u, code, 0, {}};
     request.context.waveSize = 32;
-    request.context.compute = ShaderComputeStageInfo{{threads, 1u, 1u}, 256u, {false, false, false}, false, 1u};
+    request.context.compute = ShaderComputeStageInfo{{threads, 1u, 1u}, 512u, {false, false, false}, false, 1u};
     request.target.vulkanVersion = 0x00403000u;
     request.target.spirvVersion = 0x00010600u;
     request.target.subgroupSize = 32;
@@ -61,15 +72,42 @@ Barriers recompile(std::uint32_t threads, bool waveLdsMemoryBarriers) {
     request.target.maxWorkgroupInvocations = 1024;
     request.target.maxWorkgroupSharedMemoryBytes = 49152;
     request.target.waveLdsMemoryBarriers = waveLdsMemoryBarriers;
+    request.target.deviceScopeLdsMemoryBarriers = deviceScopeLdsMemoryBarriers;
     request.layout = {0, 0, 0, 128};
     request.useCache = false;
-    const auto result = Recompile(request);
-    const auto& words = result.spirv.Words();
+    return Recompile(request);
+}
+
+std::unordered_map<std::uint32_t, std::uint32_t> constantsOf(std::span<const std::uint32_t> words) {
     std::unordered_map<std::uint32_t, std::uint32_t> constants;
+    for (std::size_t at = 5; at < words.size() && (words[at] >> 16u) != 0; at += words[at] >> 16u) {
+        if (static_cast<spv::Op>(words[at] & 0xffffu) == spv::OpConstant && (words[at] >> 16u) == 4u) constants[words[at + 2]] = words[at + 3];
+    }
+    return constants;
+}
+
+MemoryBarriers memoryBarriers(std::span<const std::uint32_t> code, bool deviceScopeLdsMemoryBarriers) {
+    const auto result = recompileProgram(code, 128, false, deviceScopeLdsMemoryBarriers);
+    const auto& words = result.spirv.Words();
+    const auto constants = constantsOf(words);
+    MemoryBarriers barriers;
+    for (std::size_t at = 5; at < words.size() && (words[at] >> 16u) != 0; at += words[at] >> 16u) {
+        if (static_cast<spv::Op>(words[at] & 0xffffu) != spv::OpMemoryBarrier) continue;
+        const auto scope = constants.find(words[at + 1]);
+        require(scope != constants.end(), "a memory barrier scope is not a 32-bit constant");
+        if (scope->second == spv::ScopeWorkgroup) ++barriers.workgroup;
+        if (scope->second == spv::ScopeDevice) ++barriers.device;
+    }
+    return barriers;
+}
+
+Barriers recompile(std::uint32_t threads, bool waveLdsMemoryBarriers) {
+    const auto result = recompileProgram(BranchedLdsStep, threads, waveLdsMemoryBarriers, false);
+    const auto& words = result.spirv.Words();
+    const auto constants = constantsOf(words);
     Barriers barriers;
     for (std::size_t at = 5; at < words.size() && (words[at] >> 16u) != 0; at += words[at] >> 16u) {
         const auto op = static_cast<spv::Op>(words[at] & 0xffffu);
-        if (op == spv::OpConstant && (words[at] >> 16u) == 4u) constants[words[at + 2]] = words[at + 3];
         if (op != spv::OpControlBarrier && op != spv::OpMemoryBarrier) continue;
         const auto scope = constants.find(words[at + 1]);
         require(scope != constants.end(), "a barrier scope is not a 32-bit constant");
@@ -83,6 +121,17 @@ Barriers recompile(std::uint32_t threads, bool waveLdsMemoryBarriers) {
 void verifyDriverQuirk() {
     for (const auto& item : Cases) {
         require(AgcDriver::SubgroupControlBarrierUsable(item.driver) == item.usable, "driver " + std::to_string(static_cast<int>(item.driver)) + ": expected subgroup control barriers to be " + (item.usable ? "usable" : "unusable"));
+        require(AgcDriver::WorkgroupScopeFenceUsable(item.driver) == item.usable, "driver " + std::to_string(static_cast<int>(item.driver)) + ": expected workgroup-scope fences to be " + (item.usable ? "usable" : "unusable"));
+    }
+}
+
+void verifyLdsMemoryBarrierScope() {
+    const std::array<std::pair<const char*, std::span<const std::uint32_t>>, 3> programs{{{"ds_write_b16", LdsHalfStore}, {"ds_add_u32", LdsAtomicAdd}, {"ds_max_f32", LdsAtomicMax}}};
+    for (const auto& [name, code] : programs) {
+        const auto workgroup = memoryBarriers(code, false);
+        require(workgroup.workgroup == 1 && workgroup.device == 0, std::string(name) + " did not order its LDS update with one workgroup-scope OpMemoryBarrier");
+        const auto device = memoryBarriers(code, true);
+        require(device.workgroup == 0 && device.device == 1, std::string(name) + " with device-scope LDS memory barriers did not order its LDS update with one device-scope OpMemoryBarrier");
     }
 }
 
@@ -101,6 +150,7 @@ int main() {
     try {
         verifyDriverQuirk();
         verifyEmittedBarriers();
+        verifyLdsMemoryBarrierScope();
         std::puts("subgroup barrier tests passed");
         return 0;
     } catch (const std::exception& error) {

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

@@ -0,0 +1,92 @@
+#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 <cstdint>
+#include <cstdio>
+#include <iostream>
+#include <span>
+#include <string>
+#include <vector>
+
+namespace {
+
+using AgcDriver::Graphics::Require;
+using ShaderRecompiler::ShaderStage;
+
+constexpr std::uint32_t Threads = 128;
+constexpr std::uint32_t BranchThreads = 64;
+constexpr std::uint32_t Iterations = 0x1000;
+constexpr std::uint32_t Rounds = 16;
+constexpr std::uint32_t LdsDwords = 512;
+alignas(256) std::array<std::uint32_t, 2u * Threads> Output{};
+
+alignas(256) constexpr std::array<std::uint32_t, 36> BranchCode{
+    0x34040082, 0xe0301000, 0x80000102, 0xbf8c3f70, 0xbe86037e, 0x7da800c0, 0xbf880002, 0xd87c0400,
+    0x00000102, 0xbefe0306, 0x7e0e0500, 0xbf0ac007, 0x850580ff, 0x00001000, 0x7e0c0300, 0xbf068005,
+    0xbf850005, 0xd5430006, 0x02050706, 0x80858105, 0xbf078005, 0xbf85fffb, 0xe0701200, 0x80000602,
+    0xd8340000, 0x00000102, 0xbf8cc07f, 0xbf8a0000, 0x3a0604ff, 0x00000100, 0xd8d80000, 0x04000003,
+    0xbf8cc07f, 0xe0701000, 0x80000402, 0xbf810000,
+};
+
+std::array<std::uint32_t, 4> BufferDescriptor(const void* data, std::uint32_t bytes) {
+    const auto address = reinterpret_cast<std::uintptr_t>(data);
+    return {static_cast<std::uint32_t>(address), static_cast<std::uint32_t>((address >> 32u) & 0xffffu), bytes, 0x31016facu};
+}
+
+std::uint32_t BranchResult(std::uint32_t thread) {
+    std::uint32_t value = thread;
+    for (std::uint32_t iteration = 0; iteration < Iterations; ++iteration) value = (value & 0xffffffu) * 3u + 1u;
+    return value;
+}
+
+}
+
+int main() {
+    try {
+        const auto device = OpenVulkanTestDevice();
+        if (!device) return VulkanTestSkipped;
+        const auto target = device->ComputeTarget(32);
+        if (target.subgroupSize != 32u) {
+            std::printf("skipped, wave32 programs do not run on 32-lane subgroups (%u lanes)\n", target.subgroupSize);
+            return VulkanTestSkipped;
+        }
+        std::vector<std::uint32_t> userData(4, 0u);
+        const auto output = BufferDescriptor(Output.data(), static_cast<std::uint32_t>(Output.size() * 4u));
+        std::copy(output.begin(), output.end(), userData.begin());
+        const std::span<const std::uint32_t> code(BranchCode);
+        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}, LdsDwords, {false, false, false}, false, 1};
+        ShaderRecompiler::RecompileRequest request{
+            {ShaderStage::Compute, reinterpret_cast<std::uintptr_t>(code.data()), code, 0, {}},
+            {32, 0, userData, compute, std::nullopt, std::nullopt, memory},
+            target,
+            {0, 0, 0, 128}
+        };
+        request.useCache = false;
+        const auto result = ShaderRecompiler::Recompile(request);
+        std::array<std::uint32_t, Threads> input{};
+        for (std::uint32_t round = 0; round < Rounds; ++round) {
+            for (std::uint32_t thread = 0; thread < Threads; ++thread) input[thread] = (round + 1u) * 0x10000u + thread;
+            Output.fill(0xdeadbeefu);
+            std::copy(input.begin(), input.end(), Output.begin());
+            device->Dispatch(result, 1, 1, 1, {}, reinterpret_cast<std::uintptr_t>(code.data()));
+            device->WaitIdle();
+            for (std::uint32_t thread = 0; thread < Threads; ++thread) {
+                const auto expected = input[thread ^ BranchThreads];
+                Require(Output[thread] == expected, "LDS exchange after a 16-bit LDS store in a branch that waves 0-1 take, round " + std::to_string(round) + ": thread " + std::to_string(thread) + " (wave " + std::to_string(thread / 32u) + ") read " + std::to_string(Output[thread]) + ", expected " + std::to_string(expected));
+            }
+            for (std::uint32_t thread = 0; thread < Threads; ++thread) {
+                const auto expected = thread < BranchThreads ? BranchResult(thread) : thread;
+                Require(Output[Threads + thread] == expected, "loop result of thread " + std::to_string(thread) + " in round " + std::to_string(round));
+            }
+        }
+        std::puts("LDS subword store branch tests passed");
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}

core/libs/prx/libSceAgcDriver/tools/ShaderReplay.cpp

@@ -49,6 +49,7 @@ bool g_graph = false;
 bool g_maintenance8 = false;
 bool g_robust2 = false;
 bool g_waveLdsMemoryBarriers = false;
+bool g_deviceScopeLdsMemoryBarriers = false;
 bool g_code = false;
 
 bool Replay(const char* path) {
@@ -149,6 +150,7 @@ bool Replay(const char* path) {
         if (g_maintenance8) request.request.target.nonConstantImageOffsets = true;
         if (g_robust2) request.request.target.robustBufferAccess2 = true;
         if (g_waveLdsMemoryBarriers) request.request.target.waveLdsMemoryBarriers = true;
+        if (g_deviceScopeLdsMemoryBarriers) request.request.target.deviceScopeLdsMemoryBarriers = true;
         const auto result = ShaderRecompiler::Recompile(request.request);
         std::printf("  recompiled: %zu SPIR-V words\n", result.spirv.size());
         if (g_spirv) {
@@ -186,7 +188,7 @@ bool Replay(const char* path) {
 
 int main(int argc, char** argv) {
     if (argc < 2) {
-        std::fprintf(stderr, "usage: agc_shader_replay [--dis] [--asm] [--cfg] [--mem] [--spv] [--prepare] [--materialize] [--code] [--maintenance8] [--robust2] [--wave-lds-memory-barriers] <shader.req>...\n  the driver writes shader_<address>.req files when APS5_DUMP_SHADERS is set; shader_<address>.prepare.req holds a request whose preparation failed before any resource capture, so replay it with --prepare\n");
+        std::fprintf(stderr, "usage: agc_shader_replay [--dis] [--asm] [--cfg] [--mem] [--spv] [--prepare] [--materialize] [--code] [--maintenance8] [--robust2] [--wave-lds-memory-barriers] [--device-scope-lds-memory-barriers] <shader.req>...\n  the driver writes shader_<address>.req files when APS5_DUMP_SHADERS is set; shader_<address>.prepare.req holds a request whose preparation failed before any resource capture, so replay it with --prepare\n");
         return 2;
     }
     int failures = 0;
@@ -215,6 +217,10 @@ int main(int argc, char** argv) {
             g_waveLdsMemoryBarriers = true;
             continue;
         }
+        if (std::string(argv[i]) == "--device-scope-lds-memory-barriers") {
+            g_deviceScopeLdsMemoryBarriers = true;
+            continue;
+        }
         if (std::string(argv[i]) == "--mem") {
             g_memory = true;
             continue;

core/shader/recompiler/CacheKey.hpp

@@ -300,6 +300,7 @@ class RecompileCacheKey {
         append(key, value.fusedExactArithmetic);
         append(key, value.largePixelFunctionLdsBuffer);
         append(key, value.pixelFunctionStorageDwords);
+        append(key, value.deviceScopeLdsMemoryBarriers);
     }
 };
 

core/shader/recompiler/Recompiler.cpp

@@ -642,6 +642,7 @@ CompiledVariant emitVariant(const RecompileRequest& request, IrProgram& program)
         if (budget.bytes > limits.maxOutputMemoryBytes) throw std::runtime_error("barycentric emulation on the mesh path needs " + std::to_string(budget.bytes) + " bytes of mesh output memory, the device allows " + std::to_string(limits.maxOutputMemoryBytes));
     }
     targetOptions.waveLdsMemoryBarriers = request.target.waveLdsMemoryBarriers;
+    targetOptions.deviceScopeLdsMemoryBarriers = request.target.deviceScopeLdsMemoryBarriers;
 
     constexpr SpirvEmitter spirvEmitter;
     CompiledShaderArtifact result;

core/shader/recompiler/Recompiler.hpp

@@ -301,6 +301,7 @@ struct SpirvTarget {
     bool fusedExactArithmetic = false;
     bool largePixelFunctionLdsBuffer = false;
     std::uint32_t pixelFunctionStorageDwords = 4096u;
+    bool deviceScopeLdsMemoryBarriers = false;
 };
 
 struct BindingLayout {

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

@@ -41,6 +41,7 @@ struct SpirvTargetOptions {
     bool waveLdsMemoryBarriers = false;
     bool largePixelFunctionLdsBuffer = false;
     std::uint32_t pixelFunctionStorageDwords = 4096u;
+    bool deviceScopeLdsMemoryBarriers = false;
 };
 
 struct SpirvEmitInfo {

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

@@ -112,6 +112,7 @@ void EmitMemoryOffsets(SpirvEmitterState& state);
 std::uint32_t LdsDwordCount(const SpirvEmitterState& state);
 bool LdsInDeviceMemory(const SpirvEmitterState& state);
 std::uint32_t LdsMemorySemantics(const SpirvEmitterState& state);
+std::uint32_t LdsMemoryBarrierScope(const SpirvEmitterState& state);
 void EmitLdsBufferBase(SpirvEmitterState& state);
 std::uint32_t EmitLdsLockPointer(SpirvEmitterState& state);
 std::uint32_t EmitLdsBufferU64Pointer(SpirvEmitterState& state, std::uint32_t index);
@@ -276,7 +277,8 @@ std::uint32_t AtomicUpdateTyped(SpirvEmitterState& state, std::uint32_t pointer,
     EmitLabel(state, cont);
     state.module.AddFunction(spv::OpBranch, header);
     EmitLabel(state, merge);
-    if (kind != ResourceKind::Lds || !state.insideLdsLock) state.module.AddFunction(spv::OpMemoryBarrier, ConstantU32(state, scope), ConstantU32(state, spv::MemorySemanticsAcquireReleaseMask | memory));
+    const auto fenceScope = kind == ResourceKind::Lds ? LdsMemoryBarrierScope(state) : scope;
+    if (kind != ResourceKind::Lds || !state.insideLdsLock) state.module.AddFunction(spv::OpMemoryBarrier, ConstantU32(state, fenceScope), ConstantU32(state, spv::MemorySemanticsAcquireReleaseMask | memory));
     return observed;
 }
 

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

@@ -129,6 +129,7 @@ struct SpirvEmitterState {
     std::uint32_t pixelFlatPositions = 0;
     std::map<std::pair<std::uint32_t, std::uint32_t>, std::uint32_t> solvedBarycentrics;
     bool waveLdsMemoryBarriers = false;
+    bool deviceScopeLdsMemoryBarriers = false;
     std::unordered_set<const IrValue*> sharedLaneValues;
     std::uint32_t laneHalf = 0;
     std::uint32_t spirvVersion = 0x00010300u;

core/shader/recompiler/SpirvBackend/include/SpirvBackend/SpirvMemory/SpirvBufferAccess.hpp

@@ -13,6 +13,7 @@ void EmitMemoryOffsets(SpirvEmitterState& state);
 std::uint32_t LdsDwordCount(const SpirvEmitterState& state);
 bool LdsInDeviceMemory(const SpirvEmitterState& state);
 std::uint32_t LdsMemorySemantics(const SpirvEmitterState& state);
+std::uint32_t LdsMemoryBarrierScope(const SpirvEmitterState& state);
 void EmitLdsBufferBase(SpirvEmitterState& state);
 std::uint32_t EmitLdsLockPointer(SpirvEmitterState& state);
 std::uint32_t EmitLdsBufferU64Pointer(SpirvEmitterState& state, std::uint32_t index);