Save compressed shader cache to a .cpp file.

ShaderRecomp/CMakeLists.txt

@@ -17,11 +17,15 @@ add_executable(ShaderRecomp
 
 find_package(directx-dxc CONFIG REQUIRED)
 find_package(xxhash CONFIG REQUIRED)
+find_package(zstd CONFIG REQUIRED)
 
 target_link_libraries(ShaderRecomp PRIVATE
     Microsoft::DirectXShaderCompiler
-    xxHash::xxhash)
+    xxHash::xxhash
+    $<IF:$<TARGET_EXISTS:zstd::libzstd_shared>,zstd::libzstd_shared,zstd::libzstd_static>)
 
 target_precompile_headers(ShaderRecomp PRIVATE pch.h)
 
 add_compile_definitions(ShaderRecomp _CRT_SECURE_NO_WARNINGS)
+
+file(COPY ${PACKAGE_PREFIX_DIR}/bin/dxil.dll DESTINATION ${CMAKE_CURRENT_BINARY_DIR})

ShaderRecomp/main.cpp

@@ -21,110 +21,157 @@ static void writeAllBytes(const char* filePath, const void* data, size_t dataSiz
     fclose(file);
 }
 
-struct ShaderCacheEntry
+struct RecompiledShader
 {
-    XXH64_hash_t hash = 0;
-    uint32_t dxilOffset = 0;
-    uint32_t dxilSize = 0;
-    uint32_t spirvOffset = 0;
-    uint32_t spirvSize = 0;
-    uint64_t reserved = 0;
-};
-
-struct ShaderHolder
-{
-    std::string filePath;
-    size_t fileSize = 0;
-    std::unique_ptr<uint8_t[]> fileData;
+    uint8_t* data = nullptr;
     IDxcBlob* dxil = nullptr;
     IDxcBlob* spirv = nullptr;
-    ShaderCacheEntry entry;
-};
-
-struct ShaderCacheHeader
-{
-    uint32_t version = 0;
-    uint32_t count = 0;
-    uint32_t reserved0 = 0;
-    uint32_t reserved1 = 0;
 };
 
 int main(int argc, char** argv)
 {
-    if (std::filesystem::is_directory(argv[1]))
-    {
-        std::map<XXH64_hash_t, ShaderHolder> shaders;
+    const char* input =
+#ifdef SHADER_RECOMP_INPUT 
+        SHADER_RECOMP_INPUT
+#else
+        argv[1]
+#endif
+    ;
 
-        for (auto& file : std::filesystem::directory_iterator(argv[1]))
-        {
-            auto extension = file.path().extension();
-            if (extension == ".xpu" || extension == ".xvu")
-            {
-                ShaderHolder shader;
-                shader.filePath = file.path().string();
-                shader.fileData = readAllBytes(shader.filePath.c_str(), shader.fileSize);
-                shader.entry.hash = XXH3_64bits(shader.fileData.get(), shader.fileSize);
+    const char* output =
+#ifdef SHADER_RECOMP_OUTPUT 
+        SHADER_RECOMP_OUTPUT
+#else
+        argv[2]
+#endif
+        ;
 
-                shaders.emplace(shader.entry.hash, std::move(shader));
+    if (std::filesystem::is_directory(input))
+    {
+        std::vector<std::unique_ptr<uint8_t[]>> files;
+        std::map<XXH64_hash_t, RecompiledShader> shaders;
+
+        for (auto& file : std::filesystem::directory_iterator(input))
+        {
+            size_t fileSize = 0;
+            auto fileData = readAllBytes(file.path().string().c_str(), fileSize);
+            bool foundAny = false;
+
+            for (size_t i = 0; fileSize > sizeof(ShaderContainer) && i < fileSize - sizeof(ShaderContainer) - 1;)
+            {
+                auto shaderContainer = reinterpret_cast<const ShaderContainer*>(fileData.get() + i);
+                size_t dataSize = shaderContainer->virtualSize + shaderContainer->physicalSize;
+
+                if ((shaderContainer->flags & 0xFFFFFF00) == 0x102A1100 &&
+                    dataSize < (fileSize - i) &&
+                    shaderContainer->field1C == 0 &&
+                    shaderContainer->field20 == 0)
+                {
+                    XXH64_hash_t hash = XXH3_64bits(shaderContainer, dataSize);
+                    auto shader = shaders.try_emplace(hash);
+                    if (shader.second)
+                    {
+                        shader.first->second.data = fileData.get() + i;
+                        foundAny = true;
+                    }
+
+                    i += dataSize;
+                }
+                else
+                {
+                    i += sizeof(uint32_t);
                 }
             }
 
+            if (foundAny)
+                files.emplace_back(std::move(fileData));
+        }
+
+        std::atomic<uint32_t> progress = 0;
+
         std::for_each(std::execution::par_unseq, shaders.begin(), shaders.end(), [&](auto& hashShaderPair)
             {
                 auto& shader = hashShaderPair.second;
-                printf("%s\n", shader.filePath.c_str());
 
                 thread_local ShaderRecompiler recompiler;
                 recompiler = {};
-                recompiler.recompile(shader.fileData.get());
-                writeAllBytes((shader.filePath + ".hlsl").c_str(), recompiler.out.data(), recompiler.out.size());
+                recompiler.recompile(shader.data);
 
                 thread_local DxcCompiler dxcCompiler;
-
                 shader.dxil = dxcCompiler.compile(recompiler.out, recompiler.isPixelShader, false);
                 shader.spirv = dxcCompiler.compile(recompiler.out, recompiler.isPixelShader, true);
 
                 assert(shader.dxil != nullptr && shader.spirv != nullptr);
                 assert(*(reinterpret_cast<uint32_t*>(shader.dxil->GetBufferPointer()) + 1) != 0 && "DXIL was not signed properly!");
+
+                size_t currentProgress = ++progress;
+                if ((currentProgress % 10) == 0 || (currentProgress == shaders.size() - 1))
+                    std::println("Recompiling shaders... {}%", currentProgress / float(shaders.size()) * 100.0f);
             });
 
-        FILE* file = fopen(argv[2], "wb");
+        std::println("Creating shader cache...");
 
-        ShaderCacheHeader header{};
-        header.version = 0;
-        header.count = shaders.size();
-        fwrite(&header, sizeof(header), 1, file);
+        StringBuffer f;
+        f.println("#include \"shader_cache.h\"");
+        f.println("ShaderCacheEntry g_shaderCacheEntries[] = {{");
 
-        uint32_t headerSize = sizeof(ShaderCacheHeader) + shaders.size() * sizeof(ShaderCacheEntry);
-        uint32_t dxilOffset = headerSize;
-        uint32_t spirvOffset = 0;
+        std::vector<uint8_t> dxil;
+        std::vector<uint8_t> spirv;
 
-        for (auto& [_, shader] : shaders)
+        for (auto& [hash, shader] : shaders)
         {
-            shader.entry.dxilOffset = dxilOffset;
-            shader.entry.dxilSize = uint32_t(shader.dxil->GetBufferSize());
-            shader.entry.spirvOffset = spirvOffset;
-            shader.entry.spirvSize = uint32_t(shader.spirv->GetBufferSize());
-            dxilOffset += shader.entry.dxilSize;
-            spirvOffset += shader.entry.spirvSize;
+            f.println("\t{{ 0x{:X}, {}, {}, {}, {} }},",
+                hash, dxil.size(), shader.dxil->GetBufferSize(), spirv.size(), shader.spirv->GetBufferSize());
+
+            dxil.insert(dxil.end(), reinterpret_cast<uint8_t*>(shader.dxil->GetBufferPointer()),
+                reinterpret_cast<uint8_t*>(shader.dxil->GetBufferPointer()) + shader.dxil->GetBufferSize());    
+            
+            spirv.insert(spirv.end(), reinterpret_cast<uint8_t*>(shader.spirv->GetBufferPointer()),
+                reinterpret_cast<uint8_t*>(shader.spirv->GetBufferPointer()) + shader.spirv->GetBufferSize());
         }
 
-        for (auto& [_, shader] : shaders)
-        {
-            shader.entry.spirvOffset += dxilOffset;
-            fwrite(&shader.entry, sizeof(shader.entry), 1, file);
-        }
+        f.println("}};");
 
-        for (auto& [_, shader] : shaders)
-            fwrite(shader.dxil->GetBufferPointer(), 1, shader.dxil->GetBufferSize(), file);
+        std::println("Compressing DXIL cache...");
 
-        for (auto& [_, shader] : shaders)
-            fwrite(shader.spirv->GetBufferPointer(), 1, shader.spirv->GetBufferSize(), file);
+        int level = ZSTD_maxCLevel();
+        //level = ZSTD_defaultCLevel();
+        std::vector<uint8_t> dxilCompressed(ZSTD_compressBound(dxil.size()));
+        dxilCompressed.resize(ZSTD_compress(dxilCompressed.data(), dxilCompressed.size(), dxil.data(), dxil.size(), level));
 
-        fclose(file);
+        f.print("uint8_t g_compressedDxilCache[] = {{");
+
+        for (auto data : dxilCompressed)
+            f.print("{},", data);
+
+        f.println("}};");
+
+        std::println("Compressing SPIRV cache...");
+
+        std::vector<uint8_t> spirvCompressed(ZSTD_compressBound(spirv.size()));
+        spirvCompressed.resize(ZSTD_compress(spirvCompressed.data(), spirvCompressed.size(), spirv.data(), spirv.size(), level));
+
+        f.print("uint8_t g_compressedSpirvCache[] = {{");
+
+        for (auto data : spirvCompressed)
+            f.print("{},", data);
+
+        f.println("}};");
+
+        f.println("size_t g_shaderCacheEntryCount = {};", shaders.size());
+        f.println("size_t g_dxilCacheCompressedSize = {};", dxilCompressed.size());
+        f.println("size_t g_dxilCacheDecompressedSize = {};", dxil.size());
+        f.println("size_t g_spirvCacheCompressedSize = {};", spirvCompressed.size());
+        f.println("size_t g_spirvCacheDecompressedSize = {};", spirv.size());
+
+        writeAllBytes(output, f.out.data(), f.out.size());
     }
     else
     {
+        ShaderRecompiler recompiler;
+        size_t fileSize;
+        recompiler.recompile(readAllBytes(input, fileSize).get());
+        writeAllBytes(output, recompiler.out.data(), recompiler.out.size());
     }
 
     return 0;

ShaderRecomp/pch.h

@@ -8,11 +8,13 @@
 #include <cstdint>
 #include <execution>
 #include <filesystem>
+#include <print>
 #include <format>
 #include <map>
 #include <string>
 #include <unordered_map>
 #include <xxhash.h>
+#include <zstd.h>
 
 template<typename T>
 struct be

ShaderRecomp/shader_recompiler.h

@@ -3,9 +3,26 @@
 #include "shader.h"
 #include "shader_code.h"
 
-struct ShaderRecompiler
+struct StringBuffer
 {
     std::string out;
+
+    template<class... Args>
+    void print(std::format_string<Args...> fmt, Args&&... args)
+    {
+        std::vformat_to(std::back_inserter(out), fmt.get(), std::make_format_args(args...));
+    }
+
+    template<class... Args>
+    void println(std::format_string<Args...> fmt, Args&&... args)
+    {
+        std::vformat_to(std::back_inserter(out), fmt.get(), std::make_format_args(args...));
+        out += '\n';
+    }
+};
+
+struct ShaderRecompiler : StringBuffer
+{
     uint32_t indentation = 0;
     bool isPixelShader = false;
     const uint8_t* constantTableData = nullptr;
@@ -22,19 +39,6 @@ struct ShaderRecompiler
             out += '\t';
     }
 
-    template<class... Args>
-    void print(std::format_string<Args...> fmt, Args&&... args)
-    {
-        std::vformat_to(std::back_inserter(out), fmt.get(), std::make_format_args(args...));
-    }
-
-    template<class... Args>
-    void println(std::format_string<Args...> fmt, Args&&... args)
-    {
-        std::vformat_to(std::back_inserter(out), fmt.get(), std::make_format_args(args...));
-        out += '\n';
-    }
-
     void printDstSwizzle(uint32_t dstSwizzle, bool operand);
     void printDstSwizzle01(uint32_t dstRegister, uint32_t dstSwizzle);
 

vcpkg.json

@@ -2,6 +2,7 @@
     "builtin-baseline": "e63bd09dc0b7204467705c1c7c71d0e2a3f8860b",
     "dependencies": [
         "directx-dxc",
-        "xxhash"
+        "xxhash",
+        "zstd"
     ]
 }