feat(agc): store the pipeline journal's modules compressed

Module records are SPIR-V and take almost all of a journal's bytes, so
APS5_PIPELINE_JOURNAL_MAX_MIB fills with modules and later pipelines
are not recorded (#8309). A module record now keeps its identity, stage
and encoded size, followed by a zstd frame (level 1) of the encoded
result; EncodeModule writes it, and loading and replay expand it before
decoding. The journal format becomes 11: main is at 10 (8 for the point
lists of #6801, 9 for the 4_4_4_4 UNORM targets of #8564, 10 for the
colour logic operation of #9348).

Co-Authored-By: Claude Opus 5.5 <[email protected]>

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -436,6 +436,10 @@ target_link_libraries(libSceAgcDriver PRIVATE
 if(WIN32)
     target_link_libraries(libSceAgcDriver PRIVATE pdh)
 endif()
+target_include_directories(libSceAgcDriver SYSTEM PRIVATE ${CMAKE_SOURCE_DIR}/3rdparty/zstd/lib)
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    target_link_options(libSceAgcDriver PRIVATE LINKER:--exclude-libs,libzstd.a)
+endif()
 if(APS5_ENABLE_TRACY)
     target_link_libraries(libSceAgcDriver PUBLIC TracyClient)
 endif()

core/libs/prx/libSceAgcDriver/Graphics/src/PipelineJournal.cpp

@@ -14,12 +14,15 @@
 #include <cstring>
 #include <deque>
 #include <fstream>
+#include <limits>
 #include <map>
 #include <mutex>
 #include <set>
+#include <span>
 #include <system_error>
 #include <thread>
 #include <type_traits>
+#include <zstd.h>
 #ifdef _WIN32
 #include <windows.h>
 #else
@@ -47,7 +50,7 @@ static_assert(sizeof(ShaderRecompiler::TessellationConfiguration) == 40, "Tessel
 constexpr std::uint32_t FileMagic = 0x4c4a5041u;
 // Bumped when the recipe codec, CanonicalPipelineState or the store key change in a way the
 // layout guards above do not catch (a field renumbered or reinterpreted at the same size).
-constexpr std::uint32_t FileFormat = 12u;
+constexpr std::uint32_t FileFormat = 13u;
 constexpr std::uint32_t ModuleRecord = 1u;
 constexpr std::uint32_t GraphicsRecord = 2u;
 constexpr std::uint32_t ComputeRecord = 3u;
@@ -530,8 +533,21 @@ struct LoadedModule {
     std::uint64_t payloadOffset = 0;
     std::uint64_t payloadBytes = 0;
     std::uint64_t payloadHash = 0;
+    std::uint32_t moduleBytes = 0;
 };
 
+constexpr std::size_t ModuleHeaderBytes = 24u;
+
+bool expandModule(std::span<const std::byte> payload, std::vector<std::byte>& module) {
+    if (payload.size() < ModuleHeaderBytes) return false;
+    std::uint32_t bytes = 0;
+    std::memcpy(&bytes, payload.data() + 20, sizeof(bytes));
+    module.resize(bytes);
+    const auto frame = payload.subspan(ModuleHeaderBytes);
+    const auto expanded = ZSTD_decompress(module.data(), module.size(), frame.data(), frame.size());
+    return !ZSTD_isError(expanded) && expanded == bytes;
+}
+
 struct LoadedEntry {
     std::uint32_t kind;
     GraphicsPipelineRecipe graphics;
@@ -909,7 +925,7 @@ private:
                 const auto payload = std::span<const std::byte>(bytes);
                 if (ShaderRecompiler::HashBytes(payload) != record.payloadHash) break;
                 if (record.kind == ModuleRecord) {
-                    if (payload.size() < 20u) break;
+                    if (payload.size() < ModuleHeaderBytes) break;
                     ShaderRecompiler::PipelineIdentity identity;
                     std::memcpy(&identity.low, payload.data(), sizeof(identity.low));
                     std::memcpy(&identity.high, payload.data() + 8, sizeof(identity.high));
@@ -918,8 +934,9 @@ private:
                     bool decoded = !decode;
                     if (decode && !loaded.indexModules.contains(identity)) {
                         ShaderRecompiler::RecompileResult result;
-                        decoded = ShaderRecompiler::ShaderDiskCache::DecodeResult(payload.subspan(20), result) && result.Identity() == identity;
-                        if (decoded) loaded.modules.emplace(identity, LoadedModule{static_cast<ShaderRecompiler::ShaderStage>(stage), payloadOffset, record.payloadBytes, record.payloadHash});
+                        std::vector<std::byte> module;
+                        decoded = expandModule(payload, module) && ShaderRecompiler::ShaderDiskCache::DecodeResult(module, result) && result.Identity() == identity;
+                        if (decoded) loaded.modules.emplace(identity, LoadedModule{static_cast<ShaderRecompiler::ShaderStage>(stage), payloadOffset, record.payloadBytes, record.payloadHash, static_cast<std::uint32_t>(module.size())});
                         else ++dropped;
                     }
                     if (decoded) loaded.indexModules.insert(identity);
@@ -1140,12 +1157,19 @@ private:
     }
 
     std::vector<std::byte> EncodeModule(const PendingModule& module) {
+        std::vector<std::byte> encoded;
+        ShaderRecompiler::ShaderDiskCache::EncodeResult(*module.result, encoded);
+        if (encoded.size() > std::numeric_limits<std::uint32_t>::max()) return {};
         std::vector<std::byte> payload;
         Writer writer(payload);
         writer.U64(module.identity.low);
         writer.U64(module.identity.high);
         writer.U32(static_cast<std::uint32_t>(module.stage));
-        ShaderRecompiler::ShaderDiskCache::EncodeResult(*module.result, payload);
+        writer.U32(static_cast<std::uint32_t>(encoded.size()));
+        payload.resize(ModuleHeaderBytes + ZSTD_compressBound(encoded.size()));
+        const auto compressed = ZSTD_compress(payload.data() + ModuleHeaderBytes, payload.size() - ModuleHeaderBytes, encoded.data(), encoded.size(), 1);
+        if (ZSTD_isError(compressed)) return {};
+        payload.resize(ModuleHeaderBytes + compressed);
         return payload;
     }
 
@@ -1169,6 +1193,7 @@ private:
         for (const auto& identity : missing) {
             const auto module = std::ranges::find_if(pending.modules, [&](const PendingModule& candidate) { return candidate.identity == identity; });
             missingPayloads.push_back(EncodeModule(*module));
+            if (missingPayloads.back().empty()) return true;
         }
         std::uint64_t batchSize = recordBytes(pending.payload.size());
         for (const auto& payload : missingPayloads) batchSize += recordBytes(payload.size());
@@ -1389,10 +1414,12 @@ private:
     }
 
     std::shared_ptr<const ShaderRecompiler::RecompileResult> readModule(const ShaderRecompiler::PipelineIdentity& identity, const LoadedModule& module) const {
-        auto bytes = std::make_shared<std::vector<std::byte>>(static_cast<std::size_t>(module.payloadBytes));
-        if (loaded.file == nullptr || !loaded.file->Read(module.payloadOffset, *bytes) || ShaderRecompiler::HashBytes(*bytes) != module.payloadHash) return nullptr;
+        std::vector<std::byte> payload(static_cast<std::size_t>(module.payloadBytes));
+        if (loaded.file == nullptr || !loaded.file->Read(module.payloadOffset, payload) || ShaderRecompiler::HashBytes(payload) != module.payloadHash) return nullptr;
+        auto bytes = std::make_shared<std::vector<std::byte>>();
+        if (!expandModule(payload, *bytes)) return nullptr;
         auto result = std::make_shared<ShaderRecompiler::RecompileResult>();
-        if (!ShaderRecompiler::ShaderDiskCache::DecodeResult(std::span<const std::byte>(*bytes).subspan(20), *result, bytes) || !(result->Identity() == identity)) return nullptr;
+        if (!ShaderRecompiler::ShaderDiskCache::DecodeResult(*bytes, *result, bytes) || !(result->Identity() == identity)) return nullptr;
         return result;
     }
 
@@ -1421,7 +1448,7 @@ private:
                 counters.rejected.fetch_add(1, std::memory_order_relaxed);
                 return;
             }
-            const Held held(module->second.payloadBytes);
+            const Held held(module->second.moduleBytes);
             const auto result = readModule(module->first, module->second);
             if (result == nullptr) {
                 counters.rejected.fetch_add(1, std::memory_order_relaxed);
@@ -1439,7 +1466,7 @@ private:
                 return;
             }
             modules.push_back(&*module);
-            bytes += module->second.payloadBytes;
+            bytes += module->second.moduleBytes;
         }
         if (modules.empty() || context.descriptorCache == nullptr) {
             counters.rejected.fetch_add(1, std::memory_order_relaxed);

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

@@ -14,6 +14,7 @@
 #include "ShaderCacheDirectory.hpp"
 #include "../NarrowConstantStoreFixture.hpp"
 #include "Recompiler.hpp"
+#include "ShaderDiskCache.hpp"
 #include "Optimization/ResourceProgram.hpp"
 #include "VulkanTestDevice.hpp"
 #ifdef _WIN32
@@ -383,6 +384,14 @@ void HeldModuleTests() {
         for (std::size_t index = 0; index < results.size(); ++index) device->Dispatch(results[index], 1, 1, 1, {}, reinterpret_cast<std::uintptr_t>(codes[index].data()));
         device->WaitIdle();
         AgcDriver::Graphics::FlushPipelineJournal();
+        std::uint64_t encoded = 0;
+        for (const auto& result : results) {
+            std::vector<std::byte> bytes;
+            ShaderRecompiler::ShaderDiskCache::EncodeResult(result, bytes);
+            encoded += bytes.size();
+        }
+        const auto journalBytes = std::filesystem::file_size(AgcDriver::Graphics::PipelineJournalPath());
+        Require(journalBytes < encoded, "the journal is " + std::to_string(journalBytes) + " bytes with 16 of its modules alone encoding to " + std::to_string(encoded) + " bytes, so it does not store them compressed");
     }
     SetEnvironment("APS5_PIPELINE_JOURNAL_THREADS", "1");
     const auto before = Snapshot();
@@ -728,7 +737,9 @@ void OversizedRecordKeepsTheJournal() {
         device->WaitForPipelineJournalReplay();
         const auto code = ComputeCode(2);
         auto result = CompileCompute(*device, code, words, 32);
-        result.shaderDataUserWords.assign(400000u, 0u);
+        result.shaderDataUserWords.resize(400000u);
+        std::uint32_t noise = 0x9e3779b9u;
+        for (auto& word : result.shaderDataUserWords) word = noise = noise * 1664525u + 1013904223u;
         AgcDriver::Graphics::RecordComputePipeline(result, false);
         AgcDriver::Graphics::FlushPipelineJournal();
     }