perf(shader): choose the fma NaN result only when the result is NaN

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -922,6 +922,14 @@ if(BUILD_TESTING)
         set_tests_properties(vector_min_max_med16 PROPERTIES ENVIRONMENT_MODIFICATION
             "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
     endif()
+    add_test_executable(vector_fma_nan_result_tests ${LIBS_INCLUDE_DIR}/tests/VectorFmaNanResult.cpp)
+    target_link_libraries(vector_fma_nan_result_tests PRIVATE shader_recompiler_agc_driver libc)
+    configure_windows_unwind(vector_fma_nan_result_tests)
+    add_test(NAME vector_fma_nan_result COMMAND vector_fma_nan_result_tests)
+    if(WIN32)
+        set_tests_properties(vector_fma_nan_result PROPERTIES ENVIRONMENT_MODIFICATION
+            "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
+    endif()
     add_test_executable(vector_mad_legacy_tests ${LIBS_INCLUDE_DIR}/tests/VectorMadLegacy.cpp)
     target_link_libraries(vector_mad_legacy_tests PRIVATE shader_recompiler_agc_driver libc)
     configure_windows_unwind(vector_mad_legacy_tests)

core/libs/tests/VectorFmaNanResult.cpp

@@ -0,0 +1,75 @@
+#include "Translation/TranslationContext.hpp"
+#include "RdnaDecoder/RdnaVectorOpDecoder.hpp"
+#include "Optimization/SrtWalker/SrtEvaluator.hpp"
+#include <cstdint>
+#include <cstdio>
+#include <iostream>
+#include <span>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+using namespace ShaderRecompiler;
+
+namespace {
+
+void Require(bool value, const std::string& message) {
+    if (!value) throw std::runtime_error("fma nan result: " + message);
+}
+
+std::vector<const IrValue*> NanResults(const std::string& name, const std::vector<std::uint32_t>& words, bool ieee) {
+    const RdnaInstruction instruction = DecodeRdnaVectorOp(std::span<const std::uint32_t>(words), 0u);
+    IrProgram program;
+    auto& block = program.CreateBlock();
+    program.SetEntryBlock(block);
+    TranslationContext context(program, block, 256);
+    context.SetFloatMode(ShaderFloatMode{0xc0u, true, ieee, false});
+    context.TranslateInstruction(instruction);
+    std::vector<const IrValue*> results;
+    for (const auto* value : block.Instructions()) {
+        if (value->Opcode() != IrOpcode::FPNanResultFma32) continue;
+        Require(value->ArgumentCount() == 4u, name + " has the wrong operands");
+        Require(value->Flags<std::uint64_t>() == (ieee ? 1u : 0u), name + " carries the wrong IEEE flag");
+        results.push_back(value);
+    }
+    return results;
+}
+
+std::uint32_t Evaluate(std::uint32_t result, std::uint32_t lhs, std::uint32_t rhs, std::uint32_t addend, bool ieee) {
+    IrProgram program;
+    const auto constant = [&](std::uint32_t bits) -> IrValue& {
+        auto& value = program.CreateValue(IrOpcode::Void, IrType::F32);
+        value.SetImmediateU32(bits);
+        return value;
+    };
+    auto& nan = program.CreateValue(IrOpcode::FPNanResultFma32, IrType::F32, ieee ? 1u : 0u);
+    for (auto* argument : {&constant(result), &constant(lhs), &constant(rhs), &constant(addend)}) nan.AddArgument(argument);
+    const SrtRuntime runtime{};
+    Detail::Evaluator evaluator(program.Resources(), runtime);
+    std::uint32_t value = 0u;
+    Require(evaluator.Evaluate(&nan, value), "the SRT evaluator refuses FPNanResultFma32");
+    return value;
+}
+
+}
+
+int main() {
+    try {
+        for (const bool ieee : {false, true}) {
+            Require(NanResults("v_fma_f32 v5, v6, v7, v8", {0xd54b0005u, 0x04220f06u}, ieee).size() == 1u, "v_fma_f32 emits no single NaN result");
+            Require(NanResults("v_fmac_f32 v5, v6, v7", {0x560a0f06u}, ieee).size() == 1u, "v_fmac_f32 emits no single NaN result");
+            Require(NanResults("v_add_f32 v5, v6, v7", {0x060a0f06u}, ieee).empty(), "v_add_f32 emits an fma NaN result");
+        }
+        Require(Evaluate(0x3f800000u, 0x3f800000u, 0x3f800000u, 0x00000000u, true) == 0x3f800000u, "a number result is not kept");
+        Require(Evaluate(0x7fffffffu, 0x3f800000u, 0x7f812345u, 0x7fc00003u, false) == 0x7f812345u, "the first NaN source is not chosen");
+        Require(Evaluate(0x7fffffffu, 0x3f800000u, 0x7f812345u, 0x7fc00003u, true) == 0x7fc12345u, "the NaN source is not quieted in IEEE mode");
+        Require(Evaluate(0x7fffffffu, 0xffc12345u, 0x7f812345u, 0x00000000u, false) == 0xffc12345u, "the first source does not win");
+        Require(Evaluate(0x7fffffffu, 0x7f800000u, 0x00000000u, 0x7fc00003u, false) == 0xffc00000u, "a NaN addend of an invalid product is not ignored");
+        Require(Evaluate(0x7fffffffu, 0x7f800000u, 0xff800000u, 0x7fc00003u, false) == 0x7fc00003u, "the NaN addend of a valid product is not chosen");
+        Require(Evaluate(0x7fffffffu, 0x7f800000u, 0x00000000u, 0x3f800000u, false) == 0xffc00000u, "an invalid product does not give the default NaN");
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}

core/shader/recompiler/IntermediateRepresentation/include/IntermediateRepresentation/IrOpcode.hpp

@@ -224,6 +224,7 @@ enum class IrOpcode : std::uint16_t {
     FPFma32,
     FPMad32,
     FPMul32,
+    FPNanResultFma32,
     FPMin32,
     FPMax32,
     FPMinTri32,

core/shader/recompiler/IntermediateRepresentation/src/IrOpcode.cpp

@@ -265,6 +265,7 @@ constexpr std::array<OpcodeMeta, static_cast<std::size_t>(IrOpcode::Count)> Meta
     makeMeta("FPFma32", F32, F32, F32, F32),
     makeMeta("FPMad32", F32, F32, F32, F32),
     makeMeta("FPMul32", F32, F32, F32),
+    makeMeta("FPNanResultFma32", F32, F32, F32, F32, F32),
     makeMeta("FPMin32", F32, F32, F32),
     makeMeta("FPMax32", F32, F32, F32),
     makeMeta("FPMinTri32", F32, F32, F32, F32),

core/shader/recompiler/Optimization/src/MaskedSelectEliminator.cpp

@@ -136,6 +136,7 @@ bool isLaneLocal(IrOpcode opcode) {
         case IrOpcode::FPFma32:
         case IrOpcode::FPMad32:
         case IrOpcode::FPMul32:
+        case IrOpcode::FPNanResultFma32:
         case IrOpcode::FPMin32:
         case IrOpcode::FPMax32:
         case IrOpcode::FPMinTri32:

core/shader/recompiler/Optimization/src/SrtWalker/SrtEvaluator.cpp

@@ -279,6 +279,26 @@ bool Evaluator::EvaluateInst(IrValue& inst, std::uint64_t& result) {
                 return true;
             }
             return false;
+        case IrOpcode::FPNanResultFma32: {
+            std::uint64_t values[4] = {};
+            const std::size_t count = inst.ArgumentCount();
+            for (std::size_t index = 0; index < count; ++index) {
+                if (!Arg(inst, index, values[index])) return false;
+            }
+            const auto bits = [&](std::size_t index) { return static_cast<std::uint32_t>(values[index]); };
+            const auto isNan = [&](std::size_t index) { return (bits(index) & 0x7fffffffu) > 0x7f800000u; };
+            result = bits(0);
+            if (!isNan(0)) return true;
+            const auto infZero = [&](std::size_t inf, std::size_t zero) { return (bits(inf) & 0x7fffffffu) == 0x7f800000u && (bits(zero) & 0x7fffffffu) == 0u; };
+            const bool invalidProduct = count == 4u && (infZero(1, 2) || infZero(2, 1));
+            const std::uint32_t quiet = (inst.Flags<std::uint64_t>() & 1u) != 0u ? 0x00400000u : 0u;
+            std::uint32_t chosen = 0xffc00000u;
+            for (std::size_t index = count; index-- > 1u;) {
+                if (isNan(index) && !(index == 3u && invalidProduct)) chosen = bits(index) | quiet;
+            }
+            result = chosen;
+            return true;
+        }
         case IrOpcode::FPTrunc32:
             if (Arg(inst, 0, a)) {
                 result = Float32Bits(std::trunc(Float32(a)));

core/shader/recompiler/Optimization/src/SrtWalker/SrtInstructionPredicates.cpp

@@ -76,6 +76,7 @@ bool IsRuntimeUniformOp(IrOpcode opcode) {
         case IrOpcode::FPOrdGreaterThanEqual32:
         case IrOpcode::FPIsNan32:
         case IrOpcode::FPMul32:
+        case IrOpcode::FPNanResultFma32:
         case IrOpcode::FPTrunc32: return true;
         default: return false;
     }

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

@@ -172,6 +172,7 @@ std::uint32_t EmitPackHalf2x16(SpirvEmitterState& state, std::uint32_t arg0);
 std::uint32_t EmitF32ToF16BitsRte(SpirvEmitterState& state, std::uint32_t value);
 std::uint32_t EmitFPFma32(SpirvEmitterState& state, std::uint32_t arg0, std::uint32_t arg1, std::uint32_t arg2);
 std::uint32_t EmitFPMad32(SpirvEmitterState& state, std::uint32_t arg0, std::uint32_t arg1, std::uint32_t arg2);
+std::uint32_t EmitFPNanResultFma32(SpirvEmitterState& state, const IrValue& inst, std::uint32_t result, std::uint32_t lhs, std::uint32_t rhs, std::uint32_t addend);
 std::uint32_t EmitFPRoundEven32(SpirvEmitterState& state, std::uint32_t arg0);
 std::uint32_t EmitFPFloor32(SpirvEmitterState& state, std::uint32_t arg0);
 std::uint32_t EmitFPCeil32(SpirvEmitterState& state, std::uint32_t arg0);

core/shader/recompiler/SpirvBackend/src/SpirvAlu/SpirvAluEmitterMath.cpp

@@ -1179,6 +1179,25 @@ std::uint32_t EmitFPMad32(SpirvEmitterState& state, std::uint32_t arg0, std::uin
     return EmitFPAdd32(state, EmitFPMul32(state, arg0, arg1), arg2);
 }
 
+std::uint32_t EmitFPNanResultFma32(SpirvEmitterState& state, const IrValue& inst, std::uint32_t result, std::uint32_t lhs, std::uint32_t rhs, std::uint32_t addend) {
+    const auto u32 = TypeU32(state);
+    const auto boolean = TypeBool(state);
+    const auto bits = [&](std::uint32_t value) { return Unary(state, spv::OpBitcast, u32, value); };
+    const auto magnitude = [&](std::uint32_t value) { return Binary(state, spv::OpBitwiseAnd, u32, bits(value), ConstantU32(state, 0x7fffffffu)); };
+    const auto isNan = [&](std::uint32_t value) { return Binary(state, spv::OpUGreaterThan, boolean, magnitude(value), ConstantU32(state, 0x7f800000u)); };
+    const bool quiet = (inst.Flags<std::uint64_t>() & 1u) != 0u;
+    return EmitValueOrDefaultIfCondition(state, isNan(result), TypeF32(state), result, [&] {
+        const auto equals = [&](std::uint32_t value, std::uint32_t constant) { return Binary(state, spv::OpIEqual, boolean, value, ConstantU32(state, constant)); };
+        const auto infZero = [&](std::uint32_t inf, std::uint32_t zero) { return Binary(state, spv::OpLogicalAnd, boolean, equals(magnitude(inf), 0x7f800000u), equals(magnitude(zero), 0u)); };
+        const auto invalidProduct = Binary(state, spv::OpLogicalOr, boolean, infZero(lhs, rhs), infZero(rhs, lhs));
+        const auto quieted = [&](std::uint32_t value) { return quiet ? Binary(state, spv::OpBitwiseOr, u32, bits(value), ConstantU32(state, 0x00400000u)) : bits(value); };
+        auto chosen = Select(state, u32, Binary(state, spv::OpLogicalAnd, boolean, isNan(addend), Unary(state, spv::OpLogicalNot, boolean, invalidProduct)), quieted(addend), ConstantU32(state, 0xffc00000u));
+        chosen = Select(state, u32, isNan(rhs), quieted(rhs), chosen);
+        chosen = Select(state, u32, isNan(lhs), quieted(lhs), chosen);
+        return Unary(state, spv::OpBitcast, TypeF32(state), chosen);
+    });
+}
+
 std::uint32_t EmitFPRoundEven32(SpirvEmitterState& state, std::uint32_t arg0) {
     return EmitGlsl<GLSLstd450RoundEven, IrType::F32>(state, arg0);
 }

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

@@ -437,6 +437,7 @@ void EmitDirectInstruction(SpirvValueEmitContext& ctx, const IrValue& inst) {
         case IrOpcode::FPSub32: return Invoke(EmitFPSub32, ctx, inst);
         case IrOpcode::FPFma32: return Invoke(EmitFPFma32, ctx, inst);
         case IrOpcode::FPMad32: return Invoke(EmitFPMad32, ctx, inst);
+        case IrOpcode::FPNanResultFma32: return Invoke(EmitFPNanResultFma32, ctx, inst);
         case IrOpcode::FPMul32: return Invoke(EmitFPMul32, ctx, inst);
         case IrOpcode::FPMin32: return Invoke(EmitFPMin32, ctx, inst);
         case IrOpcode::FPMax32: return Invoke(EmitFPMax32, ctx, inst);

core/shader/recompiler/Translation/include/Translation/TranslationContext.hpp

@@ -62,8 +62,7 @@ class TranslationContext {
     IrF32 flushTinyProduct(IrValue* lhs, IrValue* rhs, IrValue* product, IrValue* addend = nullptr);
     IrU32 quietNan32(IrU32 bits);
     IrU32 quietNan16(IrU32 bits);
-    IrValue* nanResultF32(std::initializer_list<IrValue*> sources, IrValue* result, IrValue* invalidProduct = nullptr);
-    IrValue& invalidProductF32(IrValue* lhs, IrValue* rhs);
+    IrValue* nanResultF32(std::initializer_list<IrValue*> sources, IrValue* result);
     std::array<IrU32, 2> readU32Pair(const RdnaOperand& operand);
     IrU64 readU64(const RdnaOperand& operand);
     std::array<IrU32, 2> readF64Bits(const RdnaOperand& operand);

core/shader/recompiler/Translation/src/FloatInstructions.cpp

@@ -529,7 +529,7 @@ bool TranslationContext::floatTernary(const RdnaInstruction& inst, IrOpcode opco
     IrValue* result = &ir.Emit(opcode, IrOpcodeType(opcode), {args[0], args[1], args[2]});
     if (opcode == IrOpcode::FPFma32) result = &flushTinyProduct(args[0], args[1], result, args[2]).Value();
     const bool f32Sources = !mix || !(sourceAt(inst, 0u).opSelHi || sourceAt(inst, 1u).opSelHi || (!accumulator && sourceAt(inst, 2u).opSelHi));
-    if (opcode == IrOpcode::FPFma32 && f32Sources) result = nanResultF32({args[0], args[1], args[2]}, result, &invalidProductF32(args[0], args[1]));
+    if (opcode == IrOpcode::FPFma32 && f32Sources) result = nanResultF32({args[0], args[1], args[2]}, result);
     writeOperand(inst.destination, result);
     return true;
 }
@@ -574,7 +574,7 @@ bool TranslationContext::vFmaLegacyF32(const RdnaInstruction& inst) {
     const auto factor = [&](IrValue* value) { return &ir.Emit(IrOpcode::SelectF32, IrType::F32, {&zero.Value(), &ir.ConstantF32(0.0f), value}); };
     IrValue* lhsFactor = factor(lhs);
     IrValue* rhsFactor = factor(rhs);
-    writeOperand(inst.destination, nanResultF32({lhsFactor, rhsFactor, addend}, &ir.Emit(IrOpcode::FPFma32, IrType::F32, {lhsFactor, rhsFactor, addend}), &invalidProductF32(lhsFactor, rhsFactor)));
+    writeOperand(inst.destination, nanResultF32({lhsFactor, rhsFactor, addend}, &ir.Emit(IrOpcode::FPFma32, IrType::F32, {lhsFactor, rhsFactor, addend})));
     return true;
 }
 

core/shader/recompiler/Translation/src/TranslationContext/OperandAccess.cpp

@@ -286,27 +286,23 @@ IrU32 TranslationContext::quietNan32(IrU32 bits) {
     return ieeeMode ? IrU32(ir.BitwiseOr(bits.Value(), ir.Constant(0x00400000u))) : bits;
 }
 
-IrValue* TranslationContext::nanResultF32(std::initializer_list<IrValue*> sources, IrValue* result, IrValue* invalidProduct) {
+IrValue* TranslationContext::nanResultF32(std::initializer_list<IrValue*> sources, IrValue* result) {
+    if (sources.size() == 3u) {
+        const auto operands = sources.begin();
+        return &ir.Emit(IrOpcode::FPNanResultFma32, IrType::F32, {result, operands[0], operands[1], operands[2]}, ieeeMode ? std::uint64_t{1u} : std::uint64_t{0u});
+    }
     const auto isNan = [&](IrValue& bits) -> IrValue& { return ir.UGreaterThan(ir.BitwiseAnd(bits, ir.Constant(0x7fffffffu)), ir.Constant(0x7f800000u)); };
     IrValue* bits = &ir.BitCastU32(*result);
     bits = &ir.Select(isNan(*bits), ir.Constant(0xffc00000u), *bits);
     for (auto source = sources.end(); source != sources.begin();) {
         --source;
         IrValue& sourceBits = ir.BitCastU32(**source);
         IrValue* selected = &isNan(sourceBits);
-        if (invalidProduct != nullptr && source + 1 == sources.end() && sources.size() == 3u) selected = &ir.LogicalAnd(*selected, ir.LogicalNot(*invalidProduct));
         bits = &ir.Select(*selected, quietNan32(IrU32(sourceBits)).Value(), *bits);
     }
     return &ir.BitCastF32(*bits);
 }
 
-IrValue& TranslationContext::invalidProductF32(IrValue* lhs, IrValue* rhs) {
-    IrValue& lhsMagnitude = ir.BitwiseAnd(ir.BitCastU32(*lhs), ir.Constant(0x7fffffffu));
-    IrValue& rhsMagnitude = ir.BitwiseAnd(ir.BitCastU32(*rhs), ir.Constant(0x7fffffffu));
-    const auto infZero = [&](IrValue& inf, IrValue& zero) -> IrValue& { return ir.LogicalAnd(ir.IEqual(inf, ir.Constant(0x7f800000u)), ir.IEqual(zero, ir.Constant(0u))); };
-    return ir.LogicalOr(infZero(lhsMagnitude, rhsMagnitude), infZero(rhsMagnitude, lhsMagnitude));
-}
-
 IrU32 TranslationContext::quietNan16(IrU32 bits) {
     return ieeeMode ? IrU32(ir.BitwiseOr(bits.Value(), ir.Constant(0x0200u))) : bits;
 }