perf(agc): relocate cached draws across user words that only reach the shader as values

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -154,6 +154,7 @@ set(agcDriverSources
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Draw/DrawPipeline.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Draw/DrawPrecheck.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Draw/DrawRegisterKey.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Draw/DrawRelocate.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Draw/DrawTiming.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Draw/IndirectDraw.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/EntryPoints.cpp

core/libs/prx/libSceAgcDriver/Execution/include/Driver/Dispatch/DispatchCache.hpp

@@ -40,6 +40,14 @@ struct DispatchVariant {
     std::uint32_t flatBinding = NoFlatBinding;
 
     std::uint32_t pushOffset = 0;
+    std::shared_ptr<const ShaderRecompiler::SourceHandle> prepared;
+    std::vector<std::vector<std::pair<std::uint32_t, std::uint32_t>>> bufferWordBuffers;
+    std::vector<std::int8_t> runBases;
+    std::uint64_t relocatableWords = 0;
+    std::uint64_t forwardedWords = 0;
+    std::vector<std::pair<std::uint32_t, std::uint32_t>> shaderDataPatches;
+    std::vector<std::byte> patchedPushConstants;
+    bool dataPatched = false;
     std::shared_ptr<const ShaderRecompiler::ShaderVertexStageInfo> vertexInfo;
 
     std::shared_ptr<ShaderMemory> memory;

core/libs/prx/libSceAgcDriver/Execution/include/Driver/Draw/DrawCache.hpp

@@ -1,6 +1,7 @@
 #ifndef CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_DRIVER_DRAWCACHE_HPP
 #define CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_DRIVER_DRAWCACHE_HPP
 
+#include "UserDataUse.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Dispatch/DispatchCache.hpp"
 #include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp"
 #include <array>
@@ -30,6 +31,7 @@ struct DrawProgram {
     bool merged = false;
     bool mergedPointerRequired = false;
     std::uint32_t mergedPointer = 0;
+    std::uint32_t hiddenUserWords = 0;
 };
 
 struct DrawDecode {
@@ -118,6 +120,60 @@ struct StageCapture {
     std::vector<ShaderRecompiler::MemoryRegion> regions;
     std::uint64_t forgetSerial = 0;
     std::uint32_t pushOffset = 0;
+    ShaderRecompiler::SrtUserDataUse userDataUse;
+    std::shared_ptr<const ShaderRecompiler::SourceHandle> prepared;
+    std::vector<ShaderRecompiler::BufferBaseWord> bufferBaseWords;
+};
+
+struct RelocatableUserWord {
+    std::uint32_t reg;
+    std::vector<std::pair<std::int64_t, std::uint32_t>> reads;
+    std::vector<std::int64_t> dataOffsets;
+    bool forwarded = false;
+};
+
+struct DrawRelocation {
+    std::vector<RelocatableUserWord> words;
+    std::vector<std::uint32_t> banned;
+};
+
+struct DrawUserPatch {
+    std::size_t program;
+    std::uint32_t index, value;
+};
+
+struct DrawAliasKeys {
+    static constexpr std::size_t Stages = 3;
+    std::uint64_t structure = 0;
+    std::array<std::uint64_t, Stages> content{};
+    static std::size_t Stage(std::uint32_t reg) { return reg >= 0x10cu ? 2 : reg >= 0x08cu ? 1 : 0; }
+    std::uint64_t Key(std::uint64_t stageKey, std::uint32_t userDataBase) const { return (stageKey ^ (content[Stage(userDataBase)] + 0x9e3779b97f4a7c15ull)) * 0x100000001b3ull + userDataBase; }
+};
+
+struct DrawStageDonor {
+    std::shared_ptr<DispatchVariant> variant;
+    std::vector<std::uint32_t> userData;
+};
+
+struct DrawDecodeDonor {
+    std::shared_ptr<const DrawDecode> decode;
+    std::vector<std::uint64_t> stageKeys;
+};
+
+struct DrawRelocationPending {
+    std::uint64_t key = 0, identity = 0;
+    DrawAliasKeys aliases;
+    std::shared_ptr<const DrawRelocation> relocation;
+};
+
+struct DrawRelocationTrial {
+    static constexpr std::uint32_t Window = 1024, MinimumHitShare = 16, IdleDraws = 16384, ProbationAdmissions = 2;
+    std::uint32_t lookups = 0, hits = 0, idle = 0, admissions = 0;
+    bool productive = false;
+};
+
+struct DrawRelocationCounters {
+    std::uint64_t lookups = 0, donors = 0, refused = 0, hits = 0, evicted = 0;
 };
 
 bool DecodeProgramsCurrent(const DrawDecode& decode, const ShaderRegistry& registry);

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

@@ -141,8 +141,19 @@ class Driver {
     void maybeReportDrawCache(bool profile);
     void reportDrawCache(DrawEntryCounters& counters, bool profile);
     std::string profileDrawCacheLines(const DrawEntryCounters& counters);
-    static std::uint64_t drawRegisterKey(const QueueState& queue, const ShaderRegistry& registry, std::uint64_t deviceSerial, std::uint64_t* shape = nullptr, std::array<std::uint64_t, 5>* components = nullptr);
+    static std::uint64_t drawRegisterKey(const QueueState& queue, const ShaderRegistry& registry, std::uint64_t deviceSerial, std::uint64_t* shape = nullptr, std::array<std::uint64_t, 5>* components = nullptr, const DrawRelocation* relocation = nullptr, DrawAliasKeys* aliases = nullptr);
     void traceDrawKeyLookup(const std::array<std::uint64_t, 5>& components, bool cacheMiss);
+    static std::uint64_t drawIdentityKey(const QueueState& queue, const ShaderRegistry& registry);
+    static bool drawRelocationEnabled();
+    static std::uint64_t relocatedContentKey(const RelocatableUserWord& word, std::uint32_t value, bool& readable);
+    static void tagVariantRuns(DispatchVariant& variant, std::span<const ShaderRecompiler::MemoryRegion> regions, const ShaderRecompiler::SrtUserDataUse& use, bool eligible, std::uint32_t hiddenWords);
+    static void applyShaderDataPatches(DispatchVariant& variant);
+    static std::shared_ptr<DispatchVariant> cloneVariant(const DispatchVariant& variant);
+    static void tagBufferWords(DispatchVariant& variant, std::span<const ShaderRecompiler::BufferBaseWord> bufferBaseWords);
+    static std::shared_ptr<DispatchVariant> patchBufferWords(const DispatchVariant& variant, std::span<const std::pair<std::uint32_t, std::uint32_t>> live);
+    void learnPendingRelocation();
+    bool relocateDrawEntry(const QueueState& queue, const DrawAliasKeys& aliases, std::shared_ptr<DrawEntry>& entry, std::shared_ptr<const DrawDecode>& decode, std::vector<DrawUserPatch>& patches);
+    void reportDrawRelocation();
     static bool sameVertexInfo(const ShaderRecompiler::ShaderVertexStageInfo& a, const ShaderRecompiler::ShaderVertexStageInfo& b);
     static bool sameDecode(const DrawDecode& a, const DrawDecode& b);
     std::shared_ptr<DrawDecode> decodeDraw(const QueueState& queue, const Submission& submission);
@@ -346,6 +357,13 @@ class Driver {
     std::unordered_map<std::uint64_t, std::uint64_t> drawShapes;
     std::unordered_set<std::uint64_t> traceShapes, traceKeysSeen, traceKeysEvicted;
     std::list<std::uint64_t> drawOrder;
+    std::unordered_map<std::uint64_t, std::shared_ptr<const DrawRelocation>> drawRelocations;
+    std::unordered_map<std::uint64_t, DrawStageDonor> drawStageDonors;
+    std::unordered_map<std::uint64_t, DrawDecodeDonor> drawDecodes;
+    DrawRelocationPending drawRelocationPending;
+    std::unordered_map<std::uint64_t, DrawRelocationTrial> drawRelocationTrials;
+    DrawRelocationCounters drawRelocationCounters;
+    std::vector<std::uint64_t> drawStageDonorKeys;
     std::mutex drawCacheMutex;
     std::uint64_t drawCacheHits = 0, drawCacheEvictions = 0, drawCacheVariants = 0, drawCacheVariantBytes = 0;
     std::optional<DrawKeySightings> drawKeySightings;

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Draw/Draw.cpp

@@ -77,10 +77,28 @@ DrawVerdict Driver::draw(QueueState& queue, std::span<const std::uint32_t> packe
     std::shared_ptr<const DrawDecode> decode;
     std::shared_ptr<DrawEntry> dataEntry;
     std::uint64_t shapeKey = 0;
+    std::uint64_t drawIdentity = 0;
+    DrawAliasKeys drawAliases;
+    std::vector<DrawUserPatch> userPatches;
     if (registerKey) {
         const auto keyStart = profile ? std::chrono::steady_clock::now() : std::chrono::steady_clock::time_point{};
-        drawKey = drawRegisterKey(queue, *submission.shaders, localDevice->Serial(), &shapeKey, traceDrawCache() ? &drawKeyComponents : nullptr);
-        std::lock_guard cacheLock(drawCacheMutex);
+        std::shared_ptr<const DrawRelocation> relocation;
+        bool aliasFirst = false;
+        if (drawRelocationEnabled()) {
+            drawIdentity = drawIdentityKey(queue, *submission.shaders);
+            std::lock_guard cacheLock(drawCacheMutex);
+            learnPendingRelocation();
+            drawStageDonorKeys.clear();
+            const auto known = drawRelocations.find(drawIdentity);
+            if (known != drawRelocations.end()) {
+                auto& trial = drawRelocationTrials[drawIdentity];
+                if (trial.idle != 0) --trial.idle;
+                else relocation = known->second;
+                aliasFirst = relocation != nullptr && trial.productive;
+            }
+        }
+        drawKey = drawRegisterKey(queue, *submission.shaders, localDevice->Serial(), &shapeKey, traceDrawCache() ? &drawKeyComponents : nullptr, aliasFirst ? relocation.get() : nullptr, aliasFirst ? &drawAliases : nullptr);
+        std::unique_lock cacheLock(drawCacheMutex);
         ++drawEntryCounters.lookups;
         ++drawEntryCounters.registerKeyLookups;
         if (profile) drawEntryCounters.keyUs += std::chrono::duration<double, std::micro>(std::chrono::steady_clock::now() - keyStart).count();
@@ -90,14 +108,35 @@ DrawVerdict Driver::draw(QueueState& queue, std::span<const std::uint32_t> packe
             decode = entry->decode;
         } else {
             ++drawEntryCounters.absent;
+            if (relocation != nullptr && !aliasFirst) {
+                cacheLock.unlock();
+                drawRegisterKey(queue, *submission.shaders, localDevice->Serial(), nullptr, nullptr, relocation.get(), &drawAliases);
+                cacheLock.lock();
+            }
+            if (drawIdentity != 0) drawRelocationPending = {drawKey, drawIdentity, drawAliases, relocation};
+            if (drawAliases.structure != 0 && !relocateDrawEntry(queue, drawAliases, entry, decode, userPatches) && drawDataHits() && !drawParameters.indirect && drawShapes.contains(shapeKey)) {
+                entry.reset();
+                decode.reset();
+                userPatches.clear();
+                drawStageDonorKeys.clear();
+            }
+            if (relocation != nullptr) {
+                auto& trial = drawRelocationTrials[drawIdentity];
+                ++trial.lookups;
+                if (trial.lookups >= DrawRelocationTrial::Window) {
+                    trial.productive = trial.hits * DrawRelocationTrial::MinimumHitShare >= trial.lookups;
+                    if (!trial.productive) trial.idle = DrawRelocationTrial::IdleDraws;
+                    trial.lookups = trial.hits = 0;
+                }
+            }
             if (traceDrawCache()) {
                 auto& counters = drawEntryCounters;
                 if (traceKeysEvicted.contains(drawKey)) ++counters.absentEvicted;
                 else if (traceKeysSeen.contains(drawKey)) ++counters.absentNeverInserted;
                 else if (traceShapes.contains(shapeKey)) ++counters.absentUserWords;
                 else ++counters.absentNewRegisters;
             }
-            if (drawDataHits() && !drawParameters.indirect) {
+            if (decode == nullptr && drawDataHits() && !drawParameters.indirect) {
                 if (const auto shape = drawShapes.find(shapeKey); shape != drawShapes.end()) {
                     const auto kept = drawCache.find(shape->second);
                     if (kept != drawCache.end() && kept->second->decode != nullptr && kept->second->shape == shapeKey) dataEntry = kept->second;
@@ -113,6 +152,7 @@ DrawVerdict Driver::draw(QueueState& queue, std::span<const std::uint32_t> packe
         maybeReportDrawCache(profile);
     }
     if (decode != nullptr && !DecodeProgramsCurrent(*decode, *submission.shaders)) {
+        userPatches.clear();
         std::lock_guard cacheLock(drawCacheMutex);
         eraseDrawEntry(drawKey, entry);
         ++drawEntryCounters.absent;
@@ -144,6 +184,7 @@ DrawVerdict Driver::draw(QueueState& queue, std::span<const std::uint32_t> packe
     const auto& graphics = decode->state;
     auto pixel = decode->pixel;
     std::vector<DrawProgram> programs = decode->programs;
+    for (const auto& patch : userPatches) programs[patch.program].userData[patch.index] = patch.value;
     if (const auto skipped = SkippedShader(); skipped != 0 && std::any_of(programs.begin(), programs.end(), [skipped](const DrawProgram& program) { return program.binary.codeAddress == skipped; })) return drawn();
     const auto setMeshIndexBuffer = [&](const Pm4::DrawParameters& parameters) {
         if (!graphics.stages.mesh) return;
@@ -238,6 +279,31 @@ DrawVerdict Driver::draw(QueueState& queue, std::span<const std::uint32_t> packe
     bool verifyHit = false;
     const bool stageEntries = useDrawEntries && !(graphics.stages.tessellation && !drawParameters.indirect && drawParameters.instanceCount > 1u);
     lookupDraw(submission, localDevice, graphics, pixel, programs, roles, vertexInfos, stageEntries, registerKey, profile, drawKey, entry, matched, matchedRegions, drawHit, verifyHit, phaseTiming, phaseMs, dataEntry, dataCandidates, stageData);
+    for (auto& variant : matched) {
+        if (variant != nullptr) applyShaderDataPatches(*variant);
+    }
+    const auto relocatedDecode = decode;
+    if (!userPatches.empty()) {
+        std::lock_guard cacheLock(drawCacheMutex);
+        if (drawHit || verifyHit) {
+            ++drawRelocationCounters.hits;
+            ++drawRelocationTrials[drawIdentity].hits;
+        }
+        if (profile && drawRelocationCounters.lookups >= 100000) reportDrawRelocation();
+    }
+    if (drawHit && std::any_of(matched.begin(), matched.end(), [](const std::shared_ptr<DispatchVariant>& variant) { return variant != nullptr && variant->dataPatched; })) {
+        std::lock_guard cacheLock(drawCacheMutex);
+        for (std::size_t i = 0; i < matched.size() && i < drawStageDonorKeys.size() && i < programs.size(); ++i) {
+            if (matched[i] == nullptr || !matched[i]->dataPatched || drawStageDonorKeys[i] == 0) continue;
+            matched[i]->dataPatched = false;
+            drawStageDonors[drawStageDonorKeys[i]] = {matched[i], programs[i].userData};
+        }
+    }
+    if (!userPatches.empty() && !drawHit) {
+        auto current = std::make_shared<DrawDecode>(*decode);
+        for (const auto& patch : userPatches) current->programs[patch.program].userData[patch.index] = patch.value;
+        decode = std::move(current);
+    }
     const bool patchedStages = std::any_of(stageData.begin(), stageData.end(), [](const StageDataHit& hit) { return hit.compiled != nullptr; });
     std::uint64_t stageDataVerified = 0;
 
@@ -314,7 +380,12 @@ DrawVerdict Driver::draw(QueueState& queue, std::span<const std::uint32_t> packe
         const auto& program = programs[i];
         pushCursorBytes = Graphics::StagePushOffset(pushCursorBytes, program.binary.stage, localDevice->GraphicsPipelineLibraries());
         pushOffsets[i] = pushCursorBytes;
-        if (drawHit) {
+        const bool stageReused = !drawHit && !verifyHit && !(verifyDrawDataHits() && stageData[i].compiled != nullptr) && drawRelocationEnabled() && matched[i] != nullptr;
+        if (stageReused) {
+            fresh[i] = matched[i];
+            reused[i] = true;
+        }
+        if (drawHit || stageReused) {
 
             programResults[i] = stageData[i].compiled != nullptr ? stageData[i].compiled.get() : matched[i]->compiled.get();
             memory.insert(memory.end(), matchedRegions[i].begin(), matchedRegions[i].end());
@@ -353,6 +424,7 @@ DrawVerdict Driver::draw(QueueState& queue, std::span<const std::uint32_t> packe
         const auto slotEnd = (pushCursorBytes / Graphics::PipelinePushSlotBytes + 1u) * Graphics::PipelinePushSlotBytes;
         require(result.pushConstants.size() <= slotEnd - pushCursorBytes, "stage push constants exceed the pipeline push constant block");
         stages.push_back({program.binary.stage, &result, result.pushConstants.empty() ? 0u : pushCursorBytes});
+        if ((drawHit || stageReused) && !matched[i]->patchedPushConstants.empty()) stages.back().pushConstants = matched[i]->patchedPushConstants;
         pushCursorBytes += static_cast<std::uint32_t>(result.pushConstants.size());
     }
 

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Draw/DrawCache.cpp

@@ -102,7 +102,28 @@ bool Driver::admitDrawKey(std::uint64_t key, std::uint64_t shape) {
     std::lock_guard cacheLock(drawCacheMutex);
     if (drawCache.contains(key)) return true;
     if (!drawKeySightings) drawKeySightings.emplace(drawCacheEntries() * DrawKeySightings::Ways);
-    const bool seen = drawKeySightings->SeenBefore(key);
+    bool seen = drawKeySightings->SeenBefore(key);
+    const auto& pending = drawRelocationPending;
+    if (pending.key == key && pending.identity != 0) {
+        if (pending.aliases.structure != 0) {
+            const auto trial = drawRelocationTrials.find(pending.identity);
+            bool aliasSeen = false;
+            for (std::size_t stage = 0; stage < pending.aliases.content.size(); ++stage) {
+                if (pending.aliases.content[stage] == 0) continue;
+                const auto stageKey = (pending.aliases.structure ^ (pending.aliases.content[stage] + 0x9e3779b97f4a7c15ull + stage)) * 0x100000001b3ull;
+                if (drawKeySightings->SeenBefore(stageKey)) aliasSeen = true;
+            }
+            if (aliasSeen && trial != drawRelocationTrials.end()) {
+                auto& state = trial->second;
+                if (state.productive || state.hits != 0 || state.admissions < DrawRelocationTrial::ProbationAdmissions) {
+                    if (!state.productive && state.hits == 0) ++state.admissions;
+                    seen = true;
+                }
+            }
+        } else if (!drawRelocations.contains(pending.identity)) {
+            if (drawKeySightings->SeenBefore(pending.identity ^ 0xc2b2ae3d27d4eb4full)) seen = true;
+        }
+    }
     if (seen || shape == 0 || drawShapes.contains(shape)) return seen;
     if (!drawShapeSightings) drawShapeSightings.emplace(drawCacheEntries() * DrawKeySightings::Ways);
     return drawShapeSightings->SeenBefore(shape);