feat(recompiler): load prebuilt shader packs before the writable cache

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

@@ -1,5 +1,6 @@
 #include "ShaderDiskCache.hpp"
 #include "CacheKey.hpp"
+#include "ShaderCachePack.hpp"
 #include "VertexInputSpecialization.hpp"
 #include "SpirvBackend/SpirvSpecialization.hpp"
 #include "Optimization/ResourceMaterializer.hpp"
@@ -377,6 +378,171 @@ void verifyEntryRoundTrip() {
     require(ShaderDiskCache::DecodeEntry(incompatibleFile, key, decoded) == ShaderDiskCache::LoadStatus::Rejected, "an entry with an incompatible runtime ABI loads");
 }
 
+std::vector<std::byte> packBytes(const std::vector<std::pair<std::vector<std::byte>, std::vector<std::byte>>>& entries) {
+    using namespace ShaderRecompiler::ShaderDiskCache;
+    std::vector<PackRecord> records;
+    std::uint64_t offset = sizeof(PackHeader) + entries.size() * sizeof(PackRecord);
+    for (const auto& [key, entry] : entries) {
+        records.push_back({PackKeyHash(key), offset, entry.size()});
+        offset += entry.size();
+    }
+    std::sort(records.begin(), records.end(), [](const PackRecord& left, const PackRecord& right) { return left.hash < right.hash; });
+    PackHeader header;
+    header.sourceVersion = SourceVersion();
+    header.fileBytes = offset;
+    header.entryCount = entries.size();
+    header.indexHash = ShaderRecompiler::HashBytes(std::as_bytes(std::span(records)));
+    std::vector<std::byte> bytes(offset);
+    std::memcpy(bytes.data(), &header, sizeof(header));
+    std::memcpy(bytes.data() + sizeof(header), records.data(), records.size() * sizeof(PackRecord));
+    std::size_t at = sizeof(header) + records.size() * sizeof(PackRecord);
+    for (const auto& [key, entry] : entries) {
+        std::memcpy(bytes.data() + at, entry.data(), entry.size());
+        at += entry.size();
+    }
+    return bytes;
+}
+
+void refreshIndexHash(std::vector<std::byte>& pack, std::size_t entries) {
+    using namespace ShaderRecompiler::ShaderDiskCache;
+    const auto index = std::span(pack).subspan(sizeof(PackHeader), entries * sizeof(PackRecord));
+    const auto hash = ShaderRecompiler::HashBytes(index);
+    std::memcpy(pack.data() + offsetof(PackHeader, indexHash), &hash, sizeof(hash));
+}
+
+void verifyPrebuiltPack(const std::filesystem::path& directory) {
+    using namespace ShaderRecompiler::ShaderDiskCache;
+    SampleRequest sample;
+    const auto key = sample.Key();
+    const auto variant = sampleVariant();
+    const auto file = EncodeEntry(key, variant);
+    auto otherKey = key;
+    otherKey.back() ^= std::byte{1};
+    const auto otherFile = EncodeEntry(otherKey, variant);
+    std::filesystem::create_directories(directory);
+    const auto writePack = [&](const char* name, const std::vector<std::byte>& bytes) {
+        const auto path = directory / name;
+        require(WriteFileAtomically(path, bytes), "cannot write the test pack");
+        return path;
+    };
+    PackLoadTimings timings;
+    CompiledVariant decoded;
+    {
+        CachePack pack(writePack("valid.pack", packBytes({{key, file}, {otherKey, otherFile}})));
+        require(pack.Valid(), "a valid pack is rejected: " + pack.Error());
+        require(pack.Count() == 2, "the pack does not hold both entries");
+        require(pack.Load(key, decoded, timings) == LoadStatus::Loaded, "a pack entry does not load");
+        requireSameVariant(variant, decoded, "pack round trip");
+        auto missingKey = key;
+        missingKey.front() ^= std::byte{1};
+        require(pack.Load(missingKey, decoded, timings) == LoadStatus::Absent, "an absent key loads from the pack");
+    }
+    {
+        std::vector<std::pair<std::vector<std::byte>, std::vector<std::byte>>> mismatched;
+        mismatched.emplace_back(otherKey, file);
+        CachePack pack(writePack("mismatch.pack", packBytes(mismatched)));
+        require(pack.Valid(), "the key mismatch pack is rejected: " + pack.Error());
+        require(pack.Load(otherKey, decoded, timings) != LoadStatus::Loaded, "an entry stored under another key loads");
+    }
+    {
+        auto bytes = packBytes({{key, file}});
+        const auto truncated = static_cast<std::uint64_t>(file.size() / 2);
+        std::memcpy(bytes.data() + sizeof(PackHeader) + offsetof(PackRecord, size), &truncated, sizeof(truncated));
+        refreshIndexHash(bytes, 1);
+        CachePack pack(writePack("truncated.pack", bytes));
+        require(pack.Valid(), "the truncated-entry pack is rejected: " + pack.Error());
+        require(pack.Load(key, decoded, timings) == LoadStatus::Rejected, "a truncated pack entry loads");
+    }
+    {
+        auto bytes = packBytes({{key, file}});
+        bytes[sizeof(PackHeader)] ^= std::byte{0x10};
+        CachePack pack(writePack("corrupt-index.pack", bytes));
+        require(!pack.Valid(), "a pack with a damaged index is accepted");
+    }
+    {
+        auto bytes = packBytes({{key, file}});
+        const auto outside = static_cast<std::uint64_t>(bytes.size()) + 16;
+        std::memcpy(bytes.data() + sizeof(PackHeader) + offsetof(PackRecord, offset), &outside, sizeof(outside));
+        refreshIndexHash(bytes, 1);
+        CachePack pack(writePack("corrupt-record.pack", bytes));
+        require(!pack.Valid(), "a pack with an out-of-range record is accepted");
+    }
+    {
+        auto bytes = packBytes({{key, file}});
+        const auto wrong = SourceVersion() + 1;
+        std::memcpy(bytes.data() + offsetof(PackHeader, sourceVersion), &wrong, sizeof(wrong));
+        CachePack pack(writePack("wrong-version.pack", bytes));
+        require(!pack.Valid(), "a pack from another compiler source version is accepted");
+    }
+    {
+        auto bytes = packBytes({{key, file}});
+        bytes.front() ^= std::byte{0x10};
+        CachePack pack(writePack("unknown-format.pack", bytes));
+        require(!pack.Valid(), "a pack with an unknown format is accepted");
+    }
+}
+
+SampleRequest prebuiltSample() {
+    SampleRequest sample;
+    sample.code[3] ^= 0x5a5a5a5au;
+    return sample;
+}
+
+int runPrebuiltLoadingProcess() {
+    using namespace ShaderRecompiler::ShaderDiskCache;
+    SampleRequest sample = prebuiltSample();
+    const auto key = sample.Key();
+    const auto expected = sampleVariant();
+    CompiledVariant loaded;
+    require(Load(key, loaded), "the prebuilt pack entry does not load");
+    requireSameVariant(expected, loaded, "prebuilt load");
+    require(Totals().prebuiltHits == 1, "the prebuilt load was not counted");
+    const auto afterPack = Totals();
+    require(afterPack.prebuiltBytesRead > 0 && afterPack.prebuiltLoadNanoseconds > 0, "the prebuilt load was not timed");
+    require(afterPack.hits == 0 && afterPack.loadFailures == 0 && afterPack.bytesRead == 0 && afterPack.loadNanoseconds == 0, "the prebuilt load reached the writable totals");
+    auto looseKey = key;
+    looseKey.back() ^= std::byte{1};
+    CompiledVariant loose;
+    require(Load(looseKey, loose), "the loose prebuilt entry does not load");
+    requireSameVariant(expected, loose, "loose prebuilt load");
+    require(Totals().prebuiltHits == 2, "the loose prebuilt load was not counted");
+    require(Totals().hits == 0 && Totals().loadFailures == 0 && Totals().bytesRead == 0, "the loose prebuilt load reached the writable totals");
+    auto corruptKey = looseKey;
+    corruptKey.front() ^= std::byte{1};
+    CompiledVariant corrupt;
+    require(!Load(corruptKey, corrupt), "a corrupt loose prebuilt entry loads");
+    require(Totals().prebuiltRejected >= 1, "the corrupt loose prebuilt entry was not counted as rejected");
+    require(Totals().loadFailures == 0, "the corrupt loose prebuilt entry reached the writable failures");
+    auto otherKey = looseKey;
+    otherKey.front() ^= std::byte{1};
+    CompiledVariant other;
+    require(!Load(otherKey, other), "a key absent from prebuilt storage loads");
+    Store(otherKey, std::make_shared<const CompiledVariant>(sampleVariant()));
+    Flush();
+    require(Load(otherKey, other), "the writable cache fall-through does not load");
+    std::cout << "prebuilt process loaded from the pack and fell through to the writable cache\n";
+    return 0;
+}
+
+void verifyPrebuiltAcrossProcesses(const char* self, const std::filesystem::path& directory) {
+    using namespace ShaderRecompiler::ShaderDiskCache;
+    SampleRequest sample = prebuiltSample();
+    const auto prebuilt = directory / "prebuilt";
+    std::filesystem::create_directories(prebuilt);
+    const auto bytes = packBytes({{sample.Key(), EncodeEntry(sample.Key(), sampleVariant())}});
+    require(WriteFileAtomically(prebuilt / "test.pack", bytes), "cannot write the prebuilt pack");
+    auto looseKey = sample.Key();
+    looseKey.back() ^= std::byte{1};
+    require(WriteFileAtomically(prebuilt / EntryName(looseKey), EncodeEntry(looseKey, sampleVariant())), "cannot write the loose prebuilt entry");
+    auto corruptKey = looseKey;
+    corruptKey.front() ^= std::byte{1};
+    auto corruptEntry = EncodeEntry(corruptKey, sampleVariant());
+    corruptEntry.resize(corruptEntry.size() / 2);
+    require(WriteFileAtomically(prebuilt / EntryName(corruptKey), corruptEntry), "cannot write the corrupt loose prebuilt entry");
+    const std::string command = "\"" + std::string(self) + "\" --load-prebuilt";
+    require(std::system(command.c_str()) == 0, "the prebuilt loading process failed");
+}
+
 void verifyArtifactStorageIsolation() {
     SampleRequest sample;
     const auto key = sample.Key();
@@ -1906,6 +2072,7 @@ void verifyDefaultDirectory(const char* self) {
 int main(int argc, char** argv) {
     try {
         if (argc == 2 && std::string_view(argv[1]) == "--load") return runLoadingProcess();
+        if (argc == 2 && std::string_view(argv[1]) == "--load-prebuilt") return runPrebuiltLoadingProcess();
         if (argc == 2 && std::string_view(argv[1]) == "--no-failure-memo") return runWithoutFailureMemo();
         if (argc == 2 && std::string_view(argv[1]).starts_with("--load-plan")) return runPlanLoadingProcess(argv[1]);
         const auto directory = std::filesystem::temp_directory_path() / ("aps5-shader-disk-cache-test-" + std::to_string(std::random_device{}()));
@@ -1916,6 +2083,7 @@ int main(int argc, char** argv) {
         verifyAtomicReplacement(directory);
         verifyResultRoundTrip();
         verifyEntryRoundTrip();
+        verifyPrebuiltPack(directory / "packs");
         verifyArtifactStorageIsolation();
         verifyKeySensitivity();
         verifyFloatModeSerialization();
@@ -1924,6 +2092,7 @@ int main(int argc, char** argv) {
         verifyStore();
         verifyPersist();
         verifyAcrossProcesses(argv[0]);
+        verifyPrebuiltAcrossProcesses(argv[0], directory);
         verifyPlanAcrossProcesses(argv[0]);
         verifyEmissionFailureMemo();
         verifyFailureMemoSwitch(argv[0]);

core/shader/recompiler/CMakeLists.txt

@@ -36,6 +36,7 @@ function(add_shader_recompiler target runtimeTarget)
     add_library(${target} STATIC EXCLUDE_FROM_ALL
             ${CMAKE_CURRENT_FUNCTION_LIST_DIR}/Recompiler.cpp
             ${CMAKE_CURRENT_FUNCTION_LIST_DIR}/ShaderDiskCache.cpp
+            ${CMAKE_CURRENT_FUNCTION_LIST_DIR}/ShaderCachePack.cpp
             ${CMAKE_CURRENT_FUNCTION_LIST_DIR}/SpirvBackend/src/RectListShaders.cpp
             ${CMAKE_CURRENT_FUNCTION_LIST_DIR}/SpirvBackend/src/BarycentricGeometryShader.cpp
             ${CMAKE_CURRENT_FUNCTION_LIST_DIR}/RdnaDecoder/src/RdnaOpcode.cpp

core/shader/recompiler/ShaderCachePack.cpp

@@ -0,0 +1,115 @@
+#include "ShaderCachePack.hpp"
+#include "ShaderCacheDirectory.hpp"
+#include <algorithm>
+#include <bit>
+#include <chrono>
+#include <cstring>
+#include <limits>
+#ifdef _WIN32
+#ifndef NOMINMAX
+#define NOMINMAX
+#endif
+#ifndef WIN32_LEAN_AND_MEAN
+#define WIN32_LEAN_AND_MEAN
+#endif
+#include <windows.h>
+#else
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <unistd.h>
+#endif
+
+namespace ShaderRecompiler::ShaderDiskCache {
+namespace {
+using Clock = std::chrono::steady_clock;
+std::uint64_t elapsed(Clock::time_point start) {
+    return std::chrono::duration_cast<std::chrono::nanoseconds>(Clock::now() - start).count();
+}
+}
+
+std::array<std::uint64_t, 2> PackKeyHash(std::span<const std::byte> key) {
+    return {HashBytes(key, 0x5eed0001ull), HashBytes(key, 0x5eed0002ull)};
+}
+
+std::vector<std::filesystem::path> PrebuiltCachePaths(const std::filesystem::path& cacheRoot) {
+    if (cacheRoot.empty()) return {};
+    return {cacheRoot / "prebuilt"};
+}
+
+CachePack::CachePack(const std::filesystem::path& path) {
+    static_assert(std::endian::native == std::endian::little);
+#ifdef _WIN32
+    const auto handle = CreateFileW(path.c_str(), GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_DELETE, nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
+    if (handle == INVALID_HANDLE_VALUE) { error = "cannot open"; return; }
+    LARGE_INTEGER size{};
+    if (!GetFileSizeEx(handle, &size) || size.QuadPart < static_cast<LONGLONG>(sizeof(PackHeader)) || static_cast<std::uint64_t>(size.QuadPart) > std::numeric_limits<std::size_t>::max()) {
+        CloseHandle(handle); error = "invalid size"; return;
+    }
+    mapping = CreateFileMappingW(handle, nullptr, PAGE_READONLY, 0, 0, nullptr);
+    fileHandle = handle;
+    if (!mapping) { error = "cannot map"; return; }
+    const auto* data = static_cast<const std::byte*>(MapViewOfFile(mapping, FILE_MAP_READ, 0, 0, 0));
+    if (!data) { error = "cannot map view"; return; }
+    file = {data, static_cast<std::size_t>(size.QuadPart)};
+#else
+    const auto fd = open(path.c_str(), O_RDONLY | O_CLOEXEC);
+    if (fd < 0) { error = "cannot open"; return; }
+    struct stat size{};
+    if (fstat(fd, &size) != 0 || size.st_size < static_cast<off_t>(sizeof(PackHeader)) || static_cast<std::uint64_t>(size.st_size) > std::numeric_limits<std::size_t>::max()) {
+        close(fd); error = "invalid size"; return;
+    }
+    auto* data = mmap(nullptr, size.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
+    close(fd);
+    if (data == MAP_FAILED) { error = "cannot map"; return; }
+    file = {static_cast<const std::byte*>(data), static_cast<std::size_t>(size.st_size)};
+#endif
+    Validate();
+}
+
+CachePack::~CachePack() {
+#ifdef _WIN32
+    if (!file.empty()) UnmapViewOfFile(file.data());
+    if (mapping) CloseHandle(mapping);
+    if (fileHandle) CloseHandle(fileHandle);
+#else
+    if (!file.empty()) munmap(const_cast<std::byte*>(file.data()), file.size());
+#endif
+}
+
+void CachePack::Validate() {
+    PackHeader header;
+    std::memcpy(&header, file.data(), sizeof(header));
+    if (header.magic != PackMagic) { error = "unknown pack format"; return; }
+    if (header.sourceVersion != SourceVersion()) { error = "different compiler source version"; return; }
+    if (header.fileBytes != file.size() || header.entryCount > (file.size() - sizeof(header)) / sizeof(PackRecord)) { error = "truncated index"; return; }
+    const auto indexBytes = static_cast<std::size_t>(header.entryCount) * sizeof(PackRecord);
+    const auto index = file.subspan(sizeof(header), indexBytes);
+    if (HashBytes(index) != header.indexHash) { error = "index checksum mismatch"; return; }
+    const std::span<const PackRecord> checked(reinterpret_cast<const PackRecord*>(index.data()), header.entryCount);
+    std::array<std::uint64_t, 2> previous{};
+    for (const auto& record : checked) {
+        if (record.hash < previous || record.offset < sizeof(header) + indexBytes || record.offset > file.size() || record.size > file.size() - record.offset || record.size < 48) { error = "invalid index record"; return; }
+        previous = record.hash;
+    }
+    records = checked;
+}
+
+LoadStatus CachePack::Load(std::span<const std::byte> key, CompiledVariant& variant, PackLoadTimings& timings) const {
+    const auto start = Clock::now();
+    const auto hash = PackKeyHash(key);
+    const auto first = std::lower_bound(records.begin(), records.end(), hash, [](const auto& record, const auto& value) { return record.hash < value; });
+    timings.lookupNanoseconds += elapsed(start);
+    auto status = LoadStatus::Absent;
+    for (auto it = first; it != records.end() && it->hash == hash; ++it) {
+        const auto data = file.subspan(it->offset, it->size);
+        const auto decode = Clock::now();
+        const auto result = DecodeEntry(data, key, variant);
+        timings.decodeNanoseconds += elapsed(decode);
+        timings.bytes += data.size();
+        if (result == LoadStatus::Loaded) return result;
+        if (result == LoadStatus::Rejected) status = result;
+    }
+    return status;
+}
+}

core/shader/recompiler/ShaderCachePack.hpp

@@ -0,0 +1,51 @@
+#pragma once
+#include "ShaderDiskCache.hpp"
+#include <array>
+
+namespace ShaderRecompiler::ShaderDiskCache {
+
+inline constexpr std::uint64_t PackMagic = 0x314b434150535041ull;
+struct PackHeader {
+    std::uint64_t magic = PackMagic;
+    std::uint64_t sourceVersion = 0;
+    std::uint64_t fileBytes = 0;
+    std::uint64_t entryCount = 0;
+    std::uint64_t indexHash = 0;
+};
+struct PackRecord {
+    std::array<std::uint64_t, 2> hash{};
+    std::uint64_t offset = 0;
+    std::uint64_t size = 0;
+};
+static_assert(sizeof(PackHeader) == 40 && sizeof(PackRecord) == 32);
+
+struct PackLoadTimings {
+    std::uint64_t lookupNanoseconds = 0;
+    std::uint64_t decodeNanoseconds = 0;
+    std::uint64_t bytes = 0;
+};
+
+[[nodiscard]] std::array<std::uint64_t, 2> PackKeyHash(std::span<const std::byte> key);
+[[nodiscard]] std::vector<std::filesystem::path> PrebuiltCachePaths(const std::filesystem::path& cacheRoot);
+
+class CachePack {
+public:
+    explicit CachePack(const std::filesystem::path& path);
+    ~CachePack();
+    CachePack(const CachePack&) = delete;
+    CachePack& operator=(const CachePack&) = delete;
+    [[nodiscard]] bool Valid() const { return error.empty(); }
+    [[nodiscard]] const std::string& Error() const { return error; }
+    [[nodiscard]] std::size_t Count() const { return records.size(); }
+    [[nodiscard]] LoadStatus Load(std::span<const std::byte> key, CompiledVariant& variant, PackLoadTimings& timings) const;
+
+private:
+    void Validate();
+    std::span<const std::byte> file;
+    std::span<const PackRecord> records;
+    std::string error;
+    void* mapping = nullptr;
+    void* fileHandle = nullptr;
+};
+
+}