fix(shader): isolate terminating paths from loop continuations (open PR boykopovar#2311)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: 2 people

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -953,6 +953,7 @@ target_include_directories(agc_driver_prepared_shaders_tests PRIVATE
         $<TARGET_PROPERTY:shader_recompiler_agc_driver,INTERFACE_INCLUDE_DIRECTORIES>
 )
 
+target_include_directories(agc_driver_runtime_images_tests PRIVATE ${PROJECT_SOURCE_DIR}/3rdparty/SPIRV-Headers/include)
 if(BUILD_TESTING)
     set_target_properties(agc_shader_mode_register_tests PROPERTIES EXCLUDE_FROM_ALL OFF)
     add_test(NAME agc_shader_mode_register COMMAND agc_shader_mode_register_tests)

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

@@ -220,6 +220,45 @@ void testStructExtractsAreNotShuffled() {
             if ((specialized[cursor] & 0xffffu) == spv::OpVectorShuffle) shuffled = true;
         }
         check(shuffled == vectorSource, vectorSource ? "extracts of one vector were not folded into a shuffle" : "extracts of a struct were folded into a vector shuffle");
+void testLoopContinuePatchingAndEmissionRanges() {
+    SpirvModule module;
+    check(!module.FunctionHasTerminationSince(0u), "an empty emission range contains a termination");
+    const auto merge = module.AllocateId();
+    const auto oldContinue = module.AllocateId();
+    const auto newContinue = module.AllocateId();
+    module.AddFunction(spv::OpNop);
+    const auto loop = module.AddLoopMerge(merge, oldContinue);
+    const auto body = module.FunctionInstructionOffset();
+    module.AddFunction(spv::OpBranch, oldContinue);
+    check(!module.FunctionHasTerminationSince(body), "an ordinary branch was classified as a termination");
+    module.AddFunction(spv::OpUnreachable);
+    check(module.FunctionHasTerminationSince(body), "the emitted unreachable path was not detected");
+    check(!module.FunctionHasTerminationSince(module.FunctionInstructionOffset()), "an empty tail includes an earlier termination");
+    for (std::uint32_t i = 0u; i < 40000u; ++i) module.AddFunction(spv::OpNop);
+    module.Type(spv::OpTypeVoid);
+    module.AddName(newContinue, "patched_continue");
+    module.PatchLoopContinue(loop, newContinue);
+    const auto words = module.Finalize();
+    for (std::size_t offset = 5u; offset < words.size(); offset += words[offset] >> spv::WordCountShift) {
+        if ((words[offset] & spv::OpCodeMask) == spv::OpLoopMerge) {
+            check(words[offset + 1u] == merge && words[offset + 2u] == newContinue, "loop patch moved during function vector growth or declaration insertion");
+        }
+    }
+    const auto refuses = [&](auto&& action, const char* what) {
+        try { action(); check(false, what); } catch (const std::runtime_error&) {}
+    };
+    refuses([&] { module.PatchLoopContinue(0u, newContinue); }, "a non-loop instruction was patched as a loop merge");
+    refuses([&] { module.PatchLoopContinue(loop, 0u); }, "a zero continue id was accepted");
+    refuses([&] { module.PatchLoopContinue(module.FunctionInstructionOffset(), newContinue); }, "a patch beyond the instruction range was accepted");
+    refuses([&] { static_cast<void>(module.FunctionHasTerminationSince(loop + 1u)); }, "a range inside an instruction was accepted");
+    refuses([&] { static_cast<void>(module.FunctionHasTerminationSince(module.FunctionInstructionOffset() + 1u)); }, "a range past the function end was accepted");
+    SpirvModule truncated;
+    truncated.EmitFunctionInstruction({spv::OpLoopMerge});
+    refuses([&] { truncated.PatchLoopContinue(0u, newContinue); }, "a truncated loop merge was patched");
+    for (const auto opcode : {spv::OpReturn, spv::OpReturnValue, spv::OpKill, spv::OpTerminateInvocation}) {
+        SpirvModule terminated;
+        terminated.AddFunction(opcode);
+        check(terminated.FunctionHasTerminationSince(0u), "a terminating opcode was missed");
     }
 }
 
@@ -235,6 +274,7 @@ int main() {
     testModulesAreIndependent();
     testIdsStayUniqueBeyondCacheCapacity();
     testEmittedWordsAreStable();
+    testLoopContinuePatchingAndEmissionRanges();
     if (failures != 0) {
         std::fprintf(stderr, "%d check(s) failed\n", failures);
         return 1;

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

@@ -210,6 +210,19 @@ void requireOriginalDominanceSets(ControlFlowGraph graph, const char* name) {
 
 int main() {
     try {
+        for (const auto& edges : {std::vector<std::vector<std::uint32_t>>{{1}, {2, 3}, {1}, {}},
+                std::vector<std::vector<std::uint32_t>>{{1}, {2}, {1, 3}, {}}}) {
+            auto graph = makeGraph(edges);
+            const auto original = graph;
+            Structurizer{}.Structurize(graph);
+            for (const auto& source : original.blocks) {
+                const auto copies = std::count_if(graph.blocks.begin(), graph.blocks.end(), [&](const BasicBlock& block) {
+                    return block.instructionBegin == source.instructionBegin && block.instructionEnd == source.instructionEnd;
+                });
+                if (copies != 1) throw std::runtime_error("loop normalization changed a guest instruction range");
+            }
+            requireStructuredBranches(graph, "dedicated loop continuation");
+        }
         auto nested = makeGraph({{1}, {2}, {5, 3}, {5, 4}, {7}, {6, 7}, {}, {1}});
         Structurizer{}.Structurize(nested);
         requireExactPostDominators(nested, "nested selections");

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

@@ -1,10 +1,15 @@
 #include "prx/libSceAgcDriver/Graphics/include/Draw.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/Texture.hpp"
+#include "prx/libc/include/GuestAllocations.hpp"
 #include "Recompiler.hpp"
 #include "VulkanTestDevice.hpp"
 #include <algorithm>
 #include <array>
 #include <bit>
 #include <iostream>
+#include <memory>
+#include <set>
+#include <spirv/unified1/spirv.hpp>
 
 namespace {
 
@@ -14,6 +19,61 @@ alignas(256) std::array<std::array<std::uint32_t, 64>, 2> textures{};
 alignas(256) std::array<std::uint32_t, 64> output{};
 constexpr std::array<std::uint32_t, 10> code{0x7e3c02ffu, 0x3fa00000u, 0x7e3e02ffu, 0x3f000000u, 0xf09c0108u, 0x00610a1eu, 0x34060082u, 0xe0701000u, 0x80000a03u, 0xbf810000u};
 
+void RunLoop(AgcDriver::VulkanDevice& device, std::uint32_t waveSize) {
+    alignas(256) static std::array<std::uint32_t, 256> texels{};
+    constexpr std::array<std::uint32_t, 12> loopCode{
+        0xbe8c0381u, 0x7e040280u, 0x7e020280u, 0xbf0a840cu, 0xbf840006u,
+        0x4a04040cu, 0xf0201108u, 0x00010200u, 0x4a020281u, 0x800c810cu, 0xbf82fff8u, 0xbf810000u,
+    };
+    GuestAllocations::Mutation().Add(texels.data(), sizeof(texels), true, true, true);
+    const auto deregister = [](std::uint32_t* pointer) { GuestAllocations::Mutation().Remove(pointer); };
+    const std::unique_ptr<std::uint32_t, decltype(deregister)> registration(texels.data(), deregister);
+    texels.fill(0xdeadbeefu);
+    const auto imageAddress = reinterpret_cast<std::uintptr_t>(texels.data());
+    std::array<std::uint32_t, 16> userData{};
+    const std::array<std::uint32_t, 8> image{
+        static_cast<std::uint32_t>(imageAddress >> 8u), static_cast<std::uint32_t>(imageAddress >> 40u) | (20u << 20u) | (3u << 30u),
+        15u | (3u << 14u), 0x90000facu, 0u, 0u, 0u, 0u,
+    };
+    std::copy(image.begin(), image.end(), userData.begin() + 4u);
+    const ShaderRecompiler::ShaderComputeStageInfo compute{{64u, 1u, 1u}, 0u, {false, false, false}, false, 1u};
+    ShaderRecompiler::RecompileRequest request{
+        {ShaderRecompiler::ShaderStage::Compute, reinterpret_cast<std::uintptr_t>(loopCode.data()), loopCode, 0u, {}},
+        {waveSize, 0u, userData, compute, std::nullopt, std::nullopt, {}}, device.Target(), {0u, 0u, 0u, 128u},
+    };
+    request.useCache = false;
+    const auto shader = ShaderRecompiler::Recompile(request);
+    std::set<std::uint32_t> continuations;
+    const auto& words = shader.spirv;
+    for (std::size_t offset = 5u; offset < words.size(); offset += words[offset] >> 16u) {
+        if ((words[offset] & 0xffffu) == spv::OpLoopMerge) continuations.insert(words[offset + 2u]);
+    }
+    Require(!continuations.empty(), "runtime image loop lost its loop construct");
+    std::uint32_t label = 0u;
+    for (std::size_t offset = 5u; offset < words.size(); offset += words[offset] >> 16u) {
+        const auto opcode = words[offset] & 0xffffu;
+        if (opcode == spv::OpLabel) label = words[offset + 1u];
+        if (continuations.contains(label)) {
+            Require(opcode != spv::OpUnreachable && opcode != spv::OpReturn && opcode != spv::OpReturnValue &&
+                opcode != spv::OpKill && opcode != spv::OpTerminateInvocation,
+                "runtime image loop continuation includes a terminating path: label=" +
+                std::to_string(label) + " opcode=" + std::to_string(opcode));
+        }
+    }
+    device.Dispatch(shader, 1u, 1u, 1u);
+    device.WaitIdle();
+    AgcDriver::Graphics::StorageTexture::FlushPending(imageAddress, sizeof(texels), nullptr, "test");
+    device.WaitIdle();
+    constexpr std::array expected{1u, 3u, 6u, 0xdeadbeefu};
+    for (std::uint32_t row = 0u; row < expected.size(); ++row) {
+        for (std::uint32_t lane = 0u; lane < 64u; ++lane) {
+            Require(texels[row * 64u + lane] == expected[row], "runtime image loop did not preserve all three stores and exit: wave=" +
+                std::to_string(waveSize) + " row=" + std::to_string(row) + " lane=" + std::to_string(lane) + " actual=" +
+                std::to_string(texels[row * 64u + lane]) + " expected=" + std::to_string(expected[row]));
+        }
+    }
+}
+
 void Run(AgcDriver::VulkanDevice& device) {
     ShaderRecompiler::CompiledShaderArtifact first;
     const std::array formats{20u, 21u, 22u, 34u};
@@ -50,6 +110,8 @@ int main() {
     try {
         const auto device = OpenVulkanTestDevice();
         if (!device) return VulkanTestSkipped;
+        RunLoop(*device, 32u);
+        RunLoop(*device, 64u);
         Run(*device);
         std::cout << "runtime image and sampler tests passed\n";
         return 0;

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

@@ -205,6 +205,7 @@ struct SpirvDeferredPhiPatch {
 
 struct StructuredFunctionState {
     std::unordered_map<const IrBlock*, std::uint32_t> blockExitLabels;
+    std::unordered_map<const IrBlock*, std::vector<std::size_t>> loopMergeOffsets;
     std::vector<SpirvDeferredPhiPatch> deferredPhis;
 };
 

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

@@ -57,6 +57,10 @@ class SpirvModule {
     void AddFunction(std::span<const std::uint32_t> words);
     void BeginHelperFunction();
     void EndHelperFunction();
+    [[nodiscard]] std::size_t AddLoopMerge(std::uint32_t merge, std::uint32_t continuation);
+    void PatchLoopContinue(std::size_t offset, std::uint32_t continuation);
+    [[nodiscard]] std::size_t FunctionInstructionOffset() const { return functionInstructions.size(); }
+    [[nodiscard]] bool FunctionHasTerminationSince(std::size_t offset) const;
     [[nodiscard]] SpirvDeferredPhi AddDeferredPhi(std::uint32_t type, std::uint32_t result, std::size_t incomingCount);
     void PatchDeferredPhi(SpirvDeferredPhi phi, std::size_t incoming, std::uint32_t value, std::uint32_t parent);
 

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

@@ -185,15 +185,15 @@ std::uint32_t EmitBranchCondition(SpirvValueEmitContext& ctx, const BlockInfo& i
     return result;
 }
 
-void EmitStructuredTerminator(SpirvValueEmitContext& ctx, const IrProgram& program, const BlockInfo& info) {
+void EmitStructuredTerminator(SpirvValueEmitContext& ctx, StructuredFunctionState& functionState, const IrProgram& program, const BlockInfo& info) {
     auto& state = ctx.state;
     const Terminator& term = info.terminator;
     const auto emitMerge = [&]() {
         if (term.loopHeader) {
             const IrBlock* merge = TargetBlock(program, term.mergeBlock);
             const IrBlock* cont = TargetBlock(program, term.continueBlock);
             if (merge != nullptr && cont != nullptr) {
-                state.module.AddFunction(spv::OpLoopMerge, ctx.Label(merge), ctx.Label(cont), spv::LoopControlMaskNone);
+                functionState.loopMergeOffsets[cont].push_back(state.module.AddLoopMerge(ctx.Label(merge), ctx.Label(cont)));
             }
         } else if (term.kind == TerminatorKind::ConditionalBranch && term.mergeBlock != InvalidControlFlowId) {
             if (const IrBlock* merge = TargetBlock(program, term.mergeBlock); merge != nullptr) {
@@ -832,9 +832,17 @@ void EmitControlFlow(SpirvValueEmitContext& context, StructuredFunctionState& fu
         const bool stops = state.bdaStopsInvocations;
         state.continueTarget = IsContinueTarget(program, info->id);
         state.bdaStopsInvocations = stops && !state.continueTarget;
+        const auto instructionOffset = state.module.FunctionInstructionOffset();
         EmitStructuredBlock(context, functionState, block);
+        if (const auto merges = functionState.loopMergeOffsets.find(block);
+            merges != functionState.loopMergeOffsets.end() && state.module.FunctionHasTerminationSince(instructionOffset)) {
+            const auto continuation = state.module.AllocateId();
+            for (const auto offset : merges->second) state.module.PatchLoopContinue(offset, continuation);
+            state.module.AddFunction(spv::OpBranch, continuation);
+            EmitLabel(state, continuation);
+        }
         functionState.blockExitLabels.emplace(block, state.currentLabel);
-        EmitStructuredTerminator(context, program, *info);
+        EmitStructuredTerminator(context, functionState, program, *info);
         state.bdaStopsInvocations = stops;
         state.continueTarget = false;
     }

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

@@ -230,6 +230,39 @@ void SpirvModule::appendString(std::vector<std::uint32_t>& words, const std::str
     }
 }
 
+std::size_t SpirvModule::AddLoopMerge(std::uint32_t merge, std::uint32_t continuation) {
+    const auto offset = functionInstructions.size();
+    AddFunction(spv::OpLoopMerge, merge, continuation, spv::LoopControlMaskNone);
+    return offset;
+}
+
+void SpirvModule::PatchLoopContinue(std::size_t offset, std::uint32_t continuation) {
+    if (offset > functionInstructions.size() || functionInstructions.size() - offset < 4u ||
+        functionInstructions[offset] != ((4u << spv::WordCountShift) | spv::OpLoopMerge) || continuation == 0u) {
+        throw std::runtime_error("SpirvModule::PatchLoopContinue received an invalid loop merge or a zero id");
+    }
+    functionInstructions[offset + 2u] = continuation;
+}
+
+bool SpirvModule::FunctionHasTerminationSince(std::size_t offset) const {
+    if (offset > functionInstructions.size()) throw std::runtime_error("invalid function instruction offset");
+    bool boundary = offset == functionInstructions.size();
+    bool terminates = false;
+    for (std::size_t cursor = 0u; cursor < functionInstructions.size();) {
+        const auto count = functionInstructions[cursor] >> spv::WordCountShift;
+        if (count == 0u || count > functionInstructions.size() - cursor) throw std::runtime_error("invalid function instruction length");
+        boundary |= cursor == offset;
+        const auto opcode = functionInstructions[cursor] & spv::OpCodeMask;
+        if (cursor >= offset) {
+            terminates |= opcode == spv::OpUnreachable || opcode == spv::OpReturn || opcode == spv::OpReturnValue ||
+                opcode == spv::OpKill || opcode == spv::OpTerminateInvocation;
+        }
+        cursor += count;
+    }
+    if (!boundary) throw std::runtime_error("function instruction offset is not an instruction boundary");
+    return terminates;
+}
+
 SpirvDeferredPhi SpirvModule::AddDeferredPhi(std::uint32_t type, std::uint32_t result, std::size_t incomingCount) {
     std::vector<std::uint32_t> words = {spv::OpPhi, type, result};
     words.resize(words.size() + incomingCount * 2u);