Add LZ4 compression to WebGPU shaders

Flax.flaxproj

@@ -4,7 +4,7 @@
     "Major": 1,
     "Minor": 12,
     "Revision": 0,
-    "Build": 6906
+    "Build": 6907
   },
   "Company": "Flax",
   "Copyright": "Copyright (c) 2012-2026 Wojciech Figat. All rights reserved.",

Source/Engine/GraphicsDevice/Vulkan/Types.h

@@ -127,6 +127,11 @@ struct SpirvShaderHeader
         /// The raw WGSL shader code.
         /// </summary>
         WGSL = 1,
+
+        /// <summary>
+        /// The WGSL shader code compressed with LZ4.
+        /// </summary>
+        WGSL_LZ4 = 2,
     };
 
     /// <summary>

Source/Engine/GraphicsDevice/WebGPU/GPUShaderWebGPU.cpp

@@ -9,6 +9,7 @@
 #include "Engine/Core/Types/DataContainer.h"
 #include "Engine/GraphicsDevice/Vulkan/Types.h"
 #include "Engine/Serialization/MemoryReadStream.h"
+#include <ThirdParty/LZ4/lz4.h>
 
 GPUConstantBufferWebGPU::GPUConstantBufferWebGPU(GPUDeviceWebGPU* device, uint32 size, const StringView& name) noexcept
     : GPUResourceWebGPU(device, name)
@@ -25,8 +26,28 @@ GPUShaderProgram* GPUShaderWebGPU::CreateGPUShaderProgram(ShaderStage type, cons
 
     // Extract the WGSL shader
     BytesContainer wgsl;
-    ASSERT(header->Type == SpirvShaderHeader::Types::WGSL);
+    switch (header->Type)
-    wgsl.Link(bytecode);
+    {
+    case SpirvShaderHeader::Types::WGSL:
+        wgsl.Link(bytecode);
+        break;
+    case SpirvShaderHeader::Types::WGSL_LZ4:
+    {
+        int32 originalSize = *(int32*)bytecode.Get();
+        bytecode = bytecode.Slice(sizeof(int32));
+        wgsl.Allocate(originalSize);
+        const int32 res = LZ4_decompress_safe((const char*)bytecode.Get(), (char*)wgsl.Get(), bytecode.Length(), originalSize);
+        if (res <= 0)
+        {
+            LOG(Error, "Failed to decompress shader");
+            return nullptr;
+        }
+        break;
+    }
+    default:
+        LOG(Error, "Invalid shader program format");
+        return nullptr;
+    }
 
     // Create a shader module
     WGPUShaderSourceWGSL shaderCodeDesc = WGPU_SHADER_SOURCE_WGSL_INIT;

Source/Engine/GraphicsDevice/WebGPU/GraphicsDeviceWebGPU.Build.cs

@@ -20,5 +20,6 @@ public class GraphicsDeviceWebGPU : GraphicsDeviceBaseModule
         options.OutputFiles.Add(port);
         options.PublicDefinitions.Add("GRAPHICS_API_WEBGPU");
         options.PrivateIncludePaths.Add(Path.Combine(EmscriptenSdk.Instance.EmscriptenPath, "emscripten/cache/ports/emdawnwebgpu/emdawnwebgpu_pkg/webgpu/include"));
+        options.PrivateDependencies.Add("lz4");
     }
 }

Source/Engine/ShadersCompilation/WebGPU/ShaderCompilerWebGPU.Build.cs

@@ -33,6 +33,9 @@ public class ShaderCompilerWebGPU : ShaderCompiler
 
         options.PublicDefinitions.Add("COMPILE_WITH_WEBGPU_SHADER_COMPILER");
         options.PublicDependencies.Add("ShaderCompilerVulkan");
+        options.PrivateDependencies.Add("glslang");
+        options.PrivateDependencies.Add("spirv-tools");
+        options.PrivateDependencies.Add("lz4");
 
         // Deploy tint executable as a dependency for the shader compilation from SPIR-V into WGSL
         // Tint compiler from: https://github.com/google/dawn/releases

Source/Engine/ShadersCompilation/WebGPU/ShaderCompilerWebGPU.cpp

@@ -12,6 +12,7 @@
 #include "Engine/Platform/FileSystem.h"
 #include <ThirdParty/glslang/SPIRV/SpvTools.h>
 #include <ThirdParty/spirv-tools/libspirv.hpp>
+#include <ThirdParty/LZ4/lz4.h>
 
 ShaderCompilerWebGPU::ShaderCompilerWebGPU(ShaderProfile profile)
     : ShaderCompilerVulkan(profile)
@@ -90,6 +91,21 @@ bool ShaderCompilerWebGPU::Write(ShaderCompilationContext* context, ShaderFuncti
     FileSystem::DeleteFile(inputFile);
     FileSystem::DeleteFile(outputFile);
 
+    // Compress
+    const int32 srcSize = wgsl.Length() + 1;
+    const int32 maxSize = LZ4_compressBound(srcSize);
+    Array<byte> wgslCompressed;
+    wgslCompressed.Resize(maxSize + sizeof(int32));
+    const int32 dstSize = LZ4_compress_default(wgsl.Get(), (char*)wgslCompressed.Get() + sizeof(int32), srcSize, maxSize);
+    if (dstSize > 0)
+    {
+        wgslCompressed.Resize(dstSize + sizeof(int32));
+        *(int32*)wgslCompressed.Get() = srcSize; // Store original size in the beginning to decompress it
+        header.Type = SpirvShaderHeader::Types::WGSL_LZ4;
+        return WriteShaderFunctionPermutation(_context, meta, permutationIndex, bindings, &header, sizeof(header), wgslCompressed.Get(), wgslCompressed.Count());
+    }
+
+
     header.Type = SpirvShaderHeader::Types::WGSL;
     return WriteShaderFunctionPermutation(_context, meta, permutationIndex, bindings, &header, sizeof(header), wgsl.Get(), wgsl.Length() + 1);
 }