perf(shader): keep the device's f16 NaN results with IEEE_MODE clear when a startup probe matches RDNA

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -411,6 +411,7 @@ set(agcDriverSources
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/FrameCapture.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/FrameRecorder.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/VulkanDevice.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Float16NanProbe.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/FloatFlushProbe.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/FloatProbeTolerance.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Float16ConversionProbe.cpp

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

@@ -0,0 +1,21 @@
+#ifndef CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_FLOAT16NANPROBE_HPP
+#define CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_FLOAT16NANPROBE_HPP
+
+#include "Recompiler.hpp"
+
+#include <cstdint>
+
+namespace AgcDriver {
+
+class VulkanDevice;
+
+struct Float16NanProbeResult {
+    std::uint32_t mismatches = 0;
+    std::uint32_t compared = 0;
+};
+
+Float16NanProbeResult ProbeFloat16HostNans(VulkanDevice& device, const ShaderRecompiler::SpirvTarget& target);
+
+}
+
+#endif

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

@@ -0,0 +1,159 @@
+#include "prx/libSceAgcDriver/Execution/include/Float16NanProbe.hpp"
+#include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp"
+
+#include <algorithm>
+#include <array>
+#include <optional>
+#include <span>
+#include <vector>
+
+namespace AgcDriver {
+namespace {
+
+using ShaderRecompiler::ShaderStage;
+
+constexpr std::uint32_t Threads = 256;
+constexpr std::uint32_t Inputs = 4;
+constexpr std::uint32_t Results = 16;
+constexpr std::uint32_t Combinations = 4096;
+alignas(256) std::array<std::uint32_t, Threads * Inputs> Input{};
+alignas(256) std::array<std::uint32_t, Threads * Results> Output{};
+
+alignas(256) constexpr std::array<std::uint32_t, 26> PackedCode{
+    0x34020084u, 0x34060085u, 0xe0381000u, 0x80000401u, 0xbf8c3f70u,
+    0xcc0f400au, 0x18020b04u,
+    0xcc10400bu, 0x18020b04u,
+    0xcc0e400cu, 0x1c1a0b04u,
+    0xcc0f4a0du, 0x30020b04u,
+    0xcc10c00eu, 0x18020d04u,
+    0xcc0e330fu, 0x8c1a0b04u,
+    0x7e200306u, 0x78200b04u,
+    0xcc0ec011u, 0x3c1a0f04u,
+    0xe0781000u, 0x80010a03u, 0xe0781010u, 0x80010e03u, 0xbf810000u,
+};
+constexpr std::uint32_t PackedResults = 8;
+
+alignas(256) constexpr std::array<std::uint32_t, 41> ScalarCode{
+    0x34020084u, 0x34060086u, 0xe0381000u, 0x80000401u, 0xbf8c3f70u,
+    0x64140b04u,
+    0x66160b04u,
+    0x68180b04u,
+    0x6a1a0b04u,
+    0xd74b000eu, 0x041a0b04u,
+    0xd74b980fu, 0x041a0b04u,
+    0xd74b3110u, 0xc4160d04u,
+    0x7e220306u, 0x6c220b04u,
+    0x6e240b04u, 0x00003d55u,
+    0x70260b04u, 0x00000001u,
+    0xd74ba214u, 0x241a0b04u,
+    0x662a08f2u,
+    0xd5328116u, 0x20020b04u,
+    0xd5350217u, 0x40020d04u,
+    0xd5338018u, 0x00020b04u,
+    0xd5340019u, 0x20020b04u,
+    0xe0781000u, 0x80010a03u, 0xe0781010u, 0x80010e03u, 0xe0781020u, 0x80011203u, 0xe0781030u, 0x80011603u, 0xbf810000u,
+};
+constexpr std::uint32_t ScalarFmaFirst = 4;
+constexpr std::uint32_t ScalarFmaLast = 10;
+
+constexpr std::array<std::uint32_t, 16> Halves{
+    0x7e00u, 0xfe00u, 0xff2au, 0x7d55u, 0xfc01u, 0x7c01u, 0x7c00u, 0xfc00u,
+    0x0000u, 0x8000u, 0x0001u, 0x0400u, 0x3c00u, 0xbc00u, 0x3555u, 0x7bffu,
+};
+
+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};
+}
+
+bool IsNan16(std::uint32_t half) {
+    return (half & 0x7fffu) > 0x7c00u;
+}
+
+std::uint32_t Half(std::uint32_t combination, std::uint32_t operand) {
+    return Halves[(combination >> (operand * 4u)) & 15u];
+}
+
+void FillRows(std::uint32_t first) {
+    for (std::uint32_t tid = 0; tid < Threads; ++tid) {
+        const std::uint32_t low = (first + tid) % Combinations;
+        const std::uint32_t high = Combinations - 1u - low;
+        for (std::uint32_t operand = 0; operand < Inputs; ++operand) {
+            const std::uint32_t lowOperand = operand < 3u ? operand : 1u;
+            const std::uint32_t highOperand = operand < 3u ? operand : 2u;
+            Input[tid * Inputs + operand] = Half(low, lowOperand) | (Half(high, highOperand) << 16u);
+        }
+    }
+}
+
+std::vector<std::uint32_t> Run(VulkanDevice& device, const ShaderRecompiler::RecompileResult& result, std::span<const std::uint32_t> code, std::uint32_t first) {
+    FillRows(first);
+    Output.fill(0xdeadbeefu);
+    device.Dispatch(result, 1, 1, 1, {}, reinterpret_cast<std::uintptr_t>(code.data()));
+    device.WaitIdle();
+    return {Output.begin(), Output.end()};
+}
+
+ShaderRecompiler::RecompileResult Compile(ShaderRecompiler::SpirvTarget target, bool hostNans, const std::optional<ShaderRecompiler::ShaderFloatMode>& floatMode, std::span<const std::uint32_t> code) {
+    std::vector<std::uint32_t> userData(8, 0u);
+    const auto input = BufferDescriptor(Input.data(), static_cast<std::uint32_t>(Input.size() * 4u));
+    const auto output = BufferDescriptor(Output.data(), static_cast<std::uint32_t>(Output.size() * 4u));
+    std::copy(input.begin(), input.end(), userData.begin());
+    std::copy(output.begin(), output.end(), userData.begin() + 4);
+    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}, 0u, {false, false, false}, false, 1};
+    target.float16HostNans = hostNans;
+    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;
+    request.context.floatMode = floatMode;
+    return ShaderRecompiler::Recompile(request);
+}
+
+}
+
+Float16NanProbeResult ProbeFloat16HostNans(VulkanDevice& device, const ShaderRecompiler::SpirvTarget& target) {
+    Float16NanProbeResult result;
+    const std::array<std::optional<ShaderRecompiler::ShaderFloatMode>, 3> modes{std::nullopt, ShaderRecompiler::ShaderFloatMode{0xc0u, true, false, true}, ShaderRecompiler::ShaderFloatMode{0xc0u, false, false, true}};
+    for (const auto& mode : modes) {
+        const auto packedSelected = Compile(target, false, mode, PackedCode);
+        const auto packedHost = Compile(target, true, mode, PackedCode);
+        for (std::uint32_t first = 0; first < Combinations; first += Threads) {
+            const auto selected = Run(device, packedSelected, PackedCode, first);
+            const auto host = Run(device, packedHost, PackedCode, first);
+            for (std::uint32_t tid = 0; tid < Threads; ++tid) {
+                for (std::uint32_t index = 0; index < PackedResults; ++index) {
+                    ++result.compared;
+                    if (selected[tid * PackedResults + index] != host[tid * PackedResults + index]) ++result.mismatches;
+                }
+            }
+            if (result.mismatches != 0u) return result;
+        }
+        const auto scalarSelected = Compile(target, false, mode, ScalarCode);
+        const auto scalarHost = Compile(target, true, mode, ScalarCode);
+        for (std::uint32_t first = 0; first < Combinations; first += Threads) {
+            const auto scalarExact = Run(device, scalarSelected, ScalarCode, first);
+            const auto scalar = Run(device, scalarHost, ScalarCode, first);
+            for (std::uint32_t tid = 0; tid < Threads; ++tid) {
+                for (std::uint32_t index = ScalarFmaFirst; index <= ScalarFmaLast; ++index) {
+                    const std::uint32_t expected = scalarExact[tid * Results + index];
+                    const std::uint32_t actual = scalar[tid * Results + index];
+                    for (const std::uint32_t shift : {0u, 16u}) {
+                        const std::uint32_t expectedHalf = (expected >> shift) & 0xffffu;
+                        const std::uint32_t actualHalf = (actual >> shift) & 0xffffu;
+                        ++result.compared;
+                        if (expectedHalf != actualHalf && !(IsNan16(expectedHalf) && IsNan16(actualHalf) && expectedHalf != 0xfe00u)) ++result.mismatches;
+                    }
+                }
+            }
+            if (result.mismatches != 0u) return result;
+        }
+    }
+    return result;
+}
+
+}

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

@@ -13,6 +13,7 @@
 #include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/MemoryAccounting.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Float16ConversionProbe.hpp"
+#include "prx/libSceAgcDriver/Execution/include/Float16NanProbe.hpp"
 #include "prx/libSceAgcDriver/Execution/include/FloatFlushProbe.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Shaders/ShaderRegistry.hpp"
@@ -576,6 +577,8 @@ struct VulkanDevice::State {
     bool float16WidenNans = false;
     bool float16NarrowNans = false;
     bool float32DenormPreserve = false;
+    bool float16HostNanCandidate = false;
+    bool float16HostNans = false;
     bool probingFloatControls = false;
     bool demoteToHelperInvocation = false;
     bool earlyFragmentSampleMaskTestBeforeSampleCounting = false;
@@ -1770,6 +1773,7 @@ VulkanDevice::VulkanDevice(const PresentationWindow* window, bool compileOnly) :
         floatControls.shaderDenormPreserveFloat16 == VK_TRUE && floatControls.shaderSignedZeroInfNanPreserveFloat16 == VK_TRUE;
     float16Features = {VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES_KHR};
     float16Features.shaderFloat16 = VK_TRUE;
+    state->float16HostNanCandidate = float16Rte && std::getenv("APS5_NO_HOST_F16_NANS") == nullptr;
     if (float16Rte) {
         deviceExtensions.push_back(VK_KHR_SHADER_FLOAT16_INT8_EXTENSION_NAME);
         state->capabilities.push_back(spv::CapabilityFloat16);
@@ -2722,6 +2726,14 @@ VulkanDevice::VulkanDevice(const PresentationWindow* window, bool compileOnly) :
         };
         std::fprintf(stderr, "[float-controls] f16 to f32 NaN results: %s; f32 to f16 NaN results: %s (%.0f ms)\n", describe(probe.widen, state->float16WidenNans).c_str(), describe(probe.narrow, state->float16NarrowNans).c_str(), std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - probeStart).count());
     }
+    if (state->float16HostNanCandidate) {
+        const auto probeStart = std::chrono::steady_clock::now();
+        state->probingFloatControls = true;
+        const auto probe = ProbeFloat16HostNans(*this, Target());
+        state->probingFloatControls = false;
+        state->float16HostNans = probe.mismatches == 0u;
+        std::fprintf(stderr, "[float-controls] f16 NaN results with IEEE_MODE clear: %s (%u of %u probe results differ from the selected NaN, %.0f ms)\n", state->float16HostNans ? "host" : "selected", probe.mismatches, probe.compared, std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - probeStart).count());
+    }
     // A device without a window is usually replaced by the window's device before the title draws
     // (Driver::AttachWindow): its replay starts at its first draw or dispatch instead, so the
     // journal is replayed once, on the device the title draws with.
@@ -4556,6 +4568,7 @@ ShaderRecompiler::SpirvTarget VulkanDevice::Target() const {
     require(state->shaderProfile != nullptr, "shader device profile is unavailable");
     auto target = state->shaderProfile->Target();
     target.float32FlushToZero = state->float32FlushToZero;
+    target.float16HostNans = state->float16HostNans;
     target.fusedExactArithmetic = state->fusedExactArithmetic;
     target.hostInterpolationFirst = state->hostInterpolationFirst;
     target.float32DenormPreserve = state->float32DenormPreserve;

core/libs/prx/libSceAgcDriver/tests/cmake/agc_driver_float16_nan_probe_tests.cmake

@@ -0,0 +1,7 @@
+if(BUILD_TESTING)
+    add_test(NAME agc_driver_float16_nan_probe_spirv COMMAND agc_driver_float16_nan_probe_tests --verify-spirv)
+    set_tests_properties(agc_driver_float16_nan_probe_spirv PROPERTIES ENVIRONMENT "ANYPS5_NO_SHADER_CACHE=1" SKIP_RETURN_CODE 77 WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}")
+    if(WIN32)
+        set_tests_properties(agc_driver_float16_nan_probe_spirv PROPERTIES ENVIRONMENT_MODIFICATION "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
+    endif()
+endif()

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

@@ -0,0 +1,145 @@
+#include "prx/libSceAgcDriver/Execution/include/Float16NanProbe.hpp"
+#include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/Draw.hpp"
+#include "Recompiler.hpp"
+#include "VulkanTestDevice.hpp"
+#include <spirv/unified1/spirv.hpp>
+#include <algorithm>
+#include <array>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <iostream>
+#include <optional>
+#include <span>
+#include <string>
+#include <string_view>
+#include <vector>
+
+namespace {
+
+using AgcDriver::Graphics::Require;
+using ShaderRecompiler::ShaderStage;
+
+constexpr std::uint32_t Threads = 64;
+
+alignas(256) constexpr std::array<std::uint32_t, 26> PackedCode{
+    0x34020084u, 0x34060085u, 0xe0381000u, 0x80000401u, 0xbf8c3f70u,
+    0xcc0f400au, 0x18020b04u,
+    0xcc10400bu, 0x18020b04u,
+    0xcc0e400cu, 0x1c1a0b04u,
+    0xcc0f4a0du, 0x30020b04u,
+    0xcc10c00eu, 0x18020d04u,
+    0xcc0e330fu, 0x8c1a0b04u,
+    0x7e200306u, 0x78200b04u,
+    0xcc0ec011u, 0x3c1a0f04u,
+    0xe0781000u, 0x80010a03u, 0xe0781010u, 0x80010e03u, 0xbf810000u,
+};
+
+alignas(256) constexpr std::array<std::uint32_t, 20> ScalarFmaCode{
+    0x34020084u, 0x34060085u, 0xe0381000u, 0x80000401u, 0xbf8c3f70u,
+    0xd74b000au, 0x041a0b04u,
+    0xd74b980bu, 0x041a0b04u,
+    0x7e180306u, 0x6c180b04u,
+    0x6e1a0b04u, 0x00003d55u,
+    0x701c0b04u, 0x00000001u,
+    0xe0781000u, 0x80010a03u, 0xe0701010u, 0x80010e03u, 0xbf810000u,
+};
+
+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 Selects(std::span<const std::uint32_t> words) {
+    std::uint32_t selects = 0;
+    for (std::size_t cursor = 5; cursor < words.size();) {
+        const auto count = words[cursor] >> 16u;
+        Require(count != 0 && cursor + count <= words.size(), "f16 NaN probe: malformed SPIR-V");
+        if ((words[cursor] & 0xffffu) == spv::OpSelect) ++selects;
+        cursor += count;
+    }
+    return selects;
+}
+
+void VerifySpirv() {
+    constexpr std::array<std::uint32_t, 12> capabilities{
+        spv::CapabilityShader, spv::CapabilityGroupNonUniform, spv::CapabilityGroupNonUniformBallot,
+        spv::CapabilityInt64, spv::CapabilityPhysicalStorageBufferAddresses, spv::CapabilityStorageBuffer8BitAccess,
+        spv::CapabilityDenormPreserve, spv::CapabilitySignedZeroInfNanPreserve, spv::CapabilityFloat16,
+        spv::CapabilityRoundingModeRTE, spv::CapabilityInt16, spv::CapabilityFMAKHR,
+    };
+    constexpr std::array<std::string_view, 6> extensions{
+        "SPV_KHR_physical_storage_buffer", "SPV_KHR_8bit_storage", "SPV_KHR_float_controls",
+        "SPV_KHR_storage_buffer_storage_class", "SPV_KHR_16bit_storage", "SPV_KHR_fma",
+    };
+    std::vector<std::uint32_t> input(Threads * 4u), output(Threads * 8u);
+    std::vector<std::uint32_t> userData(8, 0u);
+    const auto inputDescriptor = BufferDescriptor(input.data(), static_cast<std::uint32_t>(input.size() * 4u));
+    const auto outputDescriptor = BufferDescriptor(output.data(), static_cast<std::uint32_t>(output.size() * 4u));
+    std::copy(inputDescriptor.begin(), inputDescriptor.end(), userData.begin());
+    std::copy(outputDescriptor.begin(), outputDescriptor.end(), userData.begin() + 4);
+    ShaderRecompiler::SpirvTarget target{};
+    target.vulkanVersion = 0x00403000u;
+    target.spirvVersion = 0x00010600u;
+    target.subgroupSize = 32;
+    target.bdaAbiVersion = 1;
+    target.supportedCapabilities = capabilities;
+    target.supportedExtensions = extensions;
+    target.maxWorkgroupSize = {1024, 1024, 64};
+    target.maxWorkgroupInvocations = 1024;
+    target.maxWorkgroupSharedMemoryBytes = 49152;
+    target.float16DenormPreserve = true;
+    const auto compile = [&](std::span<const std::uint32_t> code, bool hostNans, bool ieee) {
+        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}, 0u, {false, false, false}, false, 1};
+        target.float16HostNans = hostNans;
+        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;
+        request.context.floatMode = ShaderRecompiler::ShaderFloatMode{0xc0u, false, ieee, true};
+        return ShaderRecompiler::Recompile(request).spirv.Copy();
+    };
+    for (const auto& [code, name] : {std::pair{std::span<const std::uint32_t>(PackedCode), "packed add, mul, fma and fmac"}, std::pair{std::span<const std::uint32_t>(ScalarFmaCode), "scalar fma forms"}}) {
+        const auto selected = compile(code, false, false);
+        const auto host = compile(code, true, false);
+        Require(Selects(host) < Selects(selected), std::string("f16 NaN probe: the ") + name + " with float16HostNans and IEEE_MODE clear have " + std::to_string(Selects(host)) + " selects, without " + std::to_string(Selects(selected)));
+        Require(compile(code, true, true) == compile(code, false, true), std::string("f16 NaN probe: float16HostNans changed the ") + name + " with IEEE_MODE set");
+    }
+}
+
+}
+
+int main(int argc, char** argv) {
+    try {
+        if (argc == 2 && std::string_view(argv[1]) == "--verify-spirv") {
+            VerifySpirv();
+            std::puts("f16 NaN results drop their selection only when the target says the device returns RDNA's NaN with IEEE_MODE clear");
+            return 0;
+        }
+        const auto device = OpenVulkanTestDevice();
+        if (!device) return VulkanTestSkipped;
+        const auto target = device->Target();
+        const auto supports = [&](spv::Capability capability) {
+            return std::ranges::find(target.supportedCapabilities, static_cast<std::uint32_t>(capability)) != target.supportedCapabilities.end();
+        };
+        if (!supports(spv::CapabilityFloat16) || !supports(spv::CapabilityRoundingModeRTE) || !supports(spv::CapabilityDenormPreserve)) {
+            std::puts("f16 NaN probe: skipped, the device has no native f16 with RTE rounding and kept denormals");
+            return VulkanTestSkipped;
+        }
+        const auto probe = AgcDriver::ProbeFloat16HostNans(*device, target);
+        std::printf("f16 NaN probe: %u of %u results differ from the selected NaN; the device %s its own NaN results\n", probe.mismatches, probe.compared, target.float16HostNans ? "keeps" : "does not keep");
+        Require(probe.compared != 0u, "f16 NaN probe: nothing was compared");
+        Require(!target.float16HostNans || probe.mismatches == 0u, "f16 NaN probe: the device keeps its NaN results although they differ from RDNA's");
+        Require(target.float16HostNans || probe.mismatches != 0u || std::getenv("APS5_NO_HOST_F16_NANS") != nullptr, "f16 NaN probe: the device matches RDNA but keeps the selection");
+        std::puts("f16 NaN probe tests passed");
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}