[WIP] Shader compress

CMakeLists.txt

@@ -106,6 +106,10 @@ else()
     option(NCNN_BUILD_EXAMPLES "build examples" ON)
 endif()
 
+if(NCNN_VULKAN)
+    option(NCNN_SHADER_COMPRESS "compress shader" ON)
+endif()
+
 if(NCNN_SHARED_LIB)
     if(NCNN_ENABLE_LTO)
         # enable global link time optimization

cmake/ncnn_add_shader.cmake

@@ -3,13 +3,24 @@ macro(ncnn_add_shader NCNN_SHADER_SRC)
     get_filename_component(NCNN_SHADER_SRC_NAME_WE ${NCNN_SHADER_SRC} NAME_WE)
     set(NCNN_SHADER_COMP_HEADER ${CMAKE_CURRENT_BINARY_DIR}/layer/vulkan/shader/${NCNN_SHADER_SRC_NAME_WE}.comp.hex.h)
 
-    add_custom_command(
-        OUTPUT ${NCNN_SHADER_COMP_HEADER}
-        COMMAND ${CMAKE_COMMAND} -DSHADER_SRC=${NCNN_SHADER_SRC} -DSHADER_COMP_HEADER=${NCNN_SHADER_COMP_HEADER} -P "${CMAKE_CURRENT_SOURCE_DIR}/../cmake/ncnn_generate_shader_comp_header.cmake"
-        DEPENDS ${NCNN_SHADER_SRC}
-        COMMENT "Preprocessing shader source ${NCNN_SHADER_SRC_NAME_WE}.comp"
-        VERBATIM
-    )
+    if(NCNN_SHADER_COMPRESS)
+        add_custom_command(
+                OUTPUT ${NCNN_SHADER_COMP_HEADER}
+                COMMAND ${CMAKE_COMMAND} -DSHADER_SRC=${NCNN_SHADER_SRC} -DSHADER_COMP_HEADER=${NCNN_SHADER_COMP_HEADER} -DNCNN_SHADER_COMPRESS=ON -P "${CMAKE_CURRENT_SOURCE_DIR}/../cmake/ncnn_generate_shader_comp_header.cmake"
+                DEPENDS ${NCNN_SHADER_SRC}
+                COMMENT "Preprocessing shader source ${NCNN_SHADER_SRC_NAME_WE}.comp"
+                VERBATIM
+        )
+    else()
+        add_custom_command(
+                OUTPUT ${NCNN_SHADER_COMP_HEADER}
+                COMMAND ${CMAKE_COMMAND} -DSHADER_SRC=${NCNN_SHADER_SRC} -DSHADER_COMP_HEADER=${NCNN_SHADER_COMP_HEADER} -P "${CMAKE_CURRENT_SOURCE_DIR}/../cmake/ncnn_generate_shader_comp_header.cmake"
+                DEPENDS ${NCNN_SHADER_SRC}
+                COMMENT "Preprocessing shader source ${NCNN_SHADER_SRC_NAME_WE}.comp"
+                VERBATIM
+        )
+    endif()
+
     set_source_files_properties(${NCNN_SHADER_COMP_HEADER} PROPERTIES GENERATED TRUE)
 
     get_filename_component(NCNN_SHADER_COMP_HEADER_NAME ${NCNN_SHADER_COMP_HEADER} NAME)

cmake/ncnn_generate_shader_comp_header.cmake

@@ -1,14 +1,58 @@
-
 # must define SHADER_COMP_HEADER SHADER_SRC
+cmake_policy(SET CMP0007 OLD)
 
 file(READ ${SHADER_SRC} comp_data)
 
+if(NCNN_SHADER_COMPRESS)
+    set(WORDLIST_FILE "${CMAKE_SOURCE_DIR}/../../src/layer/vulkan/shader/shader_compress_dict.in")
+    if(EXISTS ${WORDLIST_FILE})
+        file(STRINGS ${WORDLIST_FILE} RAW_LIST)
+
+        # 处理转义序列并创建新列表
+        set(PROCESSED_LIST)
+        foreach(line IN LISTS RAW_LIST)
+            string(REPLACE "\\n" "\n" processed_line "${line}")
+            list(APPEND PROCESSED_LIST "${processed_line}")
+        endforeach()
+
+        # 按长度降序排序 (优先匹配长字符串)
+        list(SORT PROCESSED_LIST COMPARE STRING ORDER DESCENDING)
+        set(WORD_LIST ${PROCESSED_LIST})
+    else()
+        message(FATAL_ERROR "shader_compress_dict is missing: ${WORDLIST_FILE}")
+    endif()
+endif()
+
 # skip leading comment
 string(FIND "${comp_data}" "#version" version_start)
 if(NOT ${version_start} EQUAL -1)
     string(SUBSTRING "${comp_data}" ${version_start} -1 comp_data)
 endif()
 
+if(NCNN_SHADER_COMPRESS)
+    # use
+    list(LENGTH WORD_LIST word_count)
+    if(word_count GREATER 128)
+        message(WARNING "词表过大(超过128词)，将截断")
+        math(EXPR word_count "128")
+    endif()
+
+    set(idx 0)
+    foreach(word IN LISTS WORD_LIST)
+        if(idx LESS 128)  # 确保不超过128个词
+            # 生成替换字符(128 + idx)
+            math(EXPR char_code "128 + ${idx}")
+            string(ASCII ${char_code} replace_char)
+
+            # 执行全局替换
+            string(REPLACE "${word}" "${replace_char}" comp_data "${comp_data}")
+
+            # 索引递增
+            math(EXPR idx "${idx} + 1")
+        endif()
+    endforeach()
+endif ()
+
 # remove whitespace
 string(REGEX REPLACE "\n +" "\n" comp_data "${comp_data}")
 
@@ -24,4 +68,8 @@ string(REGEX REPLACE "([0-9a-f][0-9a-f])" "0x\\1," comp_data_hex ${comp_data_hex
 string(FIND "${comp_data_hex}" "," tail_comma REVERSE)
 string(SUBSTRING "${comp_data_hex}" 0 ${tail_comma} comp_data_hex)
 
-file(WRITE ${SHADER_COMP_HEADER} "static const char ${SHADER_SRC_NAME_WE}_comp_data[] = {${comp_data_hex}};\n")
+if(NCNN_SHADER_COMPRESS)
+    file(WRITE ${SHADER_COMP_HEADER} "static const unsigned char ${SHADER_SRC_NAME_WE}_comp_data[] = {${comp_data_hex}};\n")
+else ()
+    file(WRITE ${SHADER_COMP_HEADER} "static const char ${SHADER_SRC_NAME_WE}_comp_data[] = {${comp_data_hex}};\n")
+endif ()

src/CMakeLists.txt

@@ -174,6 +174,30 @@ ncnn_add_layer(InverseSpectrogram)
 if(NCNN_VULKAN)
     ncnn_add_shader(${CMAKE_CURRENT_SOURCE_DIR}/convert_ycbcr.comp)
     ncnn_add_shader(${CMAKE_CURRENT_SOURCE_DIR}/layer/vulkan/shader/vulkan_activation.comp)
+    if(NCNN_SHADER_COMPRESS)
+        set(WORDLIST_FILE "${CMAKE_CURRENT_SOURCE_DIR}/layer/vulkan/shader/shader_compress_dict.in")
+        if(EXISTS ${WORDLIST_FILE})
+            file(STRINGS ${WORDLIST_FILE} RAW_LIST)
+            set(PROCESSED_LIST)
+            foreach(line IN LISTS RAW_LIST)
+                string(REPLACE "\\n" "\n" processed_line "${line}")
+                list(APPEND PROCESSED_LIST "${processed_line}")
+            endforeach()
+            list(SORT PROCESSED_LIST COMPARE STRING ORDER DESCENDING)
+            set(WORD_LIST ${PROCESSED_LIST})
+
+            # 导出 WORD_LIST 到文件
+            set(OUTPUT_FILE "${CMAKE_CURRENT_BINARY_DIR}/shader_compress_dict.in")
+            file(WRITE "${OUTPUT_FILE}" "")
+            foreach(line IN LISTS WORD_LIST)
+                string(REPLACE "\\" "\\\\" escaped_line "${line}")
+                string(REPLACE "\"" "\\\"" escaped_line "${escaped_line}")
+                file(APPEND "${OUTPUT_FILE}" "\"${escaped_line}\",\n")
+            endforeach()
+        else()
+            message(FATAL_ERROR "shader_compress_dict is missing: ${WORDLIST_FILE}")
+        endif()
+    endif()
 endif()
 
 add_custom_target(ncnn-generate-spirv DEPENDS ${NCNN_SHADER_SPV_HEX_FILES})

src/gpu.cpp

@@ -87,11 +87,46 @@ static GpuInfo* g_gpu_infos[NCNN_MAX_GPU_COUNT] = {0};
 static Mutex g_default_vkdev_lock;
 static VulkanDevice* g_default_vkdev[NCNN_MAX_GPU_COUNT] = {0};
 
+#if NCNN_SHADER_COMPRESS
+struct layer_shader_registry_entry
+{
+    const unsigned char* comp_data;
+    int comp_data_size;
+};
+
+void decompress_layer_shader(const unsigned char* comp_data, const int comp_data_size, char*& decompressed_data, int& decompressed_data_size)
+{
+    static const char* dict[] = {
+#include "shader_compress_dict.in"
+    };
+
+    std::vector<char> buffer;
+    for (int i = 0; i < comp_data_size; i++)
+    {
+        if (comp_data[i] < 0x80)
+        {
+            buffer.push_back(comp_data[i]);
+        }
+        else
+        {
+            for (char* c = (char*)dict[comp_data[i] - 0x80]; *c; c++)
+            {
+                buffer.push_back(*c);
+            }
+        }
+    }
+    decompressed_data_size = (int)buffer.size();
+    decompressed_data = new char[decompressed_data_size];
+    memcpy(decompressed_data, buffer.data(), decompressed_data_size);
+}
+
+#else
 struct layer_shader_registry_entry
 {
     const char* comp_data;
     int comp_data_size;
 };
+#endif
 
 #include "layer_shader_spv_data.h"
 
@@ -4409,9 +4444,19 @@ class VulkanShaderIncluder : public glslang::TShader::Includer
     {
         if (strcmp(headerName, "vulkan_activation.comp") == 0)
         {
+#if NCNN_SHADER_COMPRESS
+            char* headerData;
+            int headerLength;
+
+            decompress_layer_shader(vulkan_activation_comp_data, sizeof(vulkan_activation_comp_data), headerData, headerLength);
+#else
             const char* const headerData = vulkan_activation_comp_data;
             const size_t headerLength = sizeof(vulkan_activation_comp_data);
+#endif
+
             glslang::TShader::Includer::IncludeResult* r = new glslang::TShader::Includer::IncludeResult(headerName, headerData, headerLength, 0);
+
+            delete[] headerData;
             return r;
         }
 
@@ -5536,11 +5581,22 @@ int compile_spirv_module(int shader_type_index, const Option& opt, std::vector<u
         NCNN_LOGE("no such shader module %d", shader_type_index);
         return -1;
     }
-
+#if NCNN_SHADER_COMPRESS
+    char* comp_data;
+    int comp_data_size;
+    decompress_layer_shader(layer_shader_registry[shader_type_index].comp_data, layer_shader_registry[shader_type_index].comp_data_size, comp_data, comp_data_size);
+#else
     const char* comp_data = layer_shader_registry[shader_type_index].comp_data;
     int comp_data_size = layer_shader_registry[shader_type_index].comp_data_size;
+#endif
+
+    int ret = compile_spirv_module(comp_data, comp_data_size, opt, spirv);
+
+#if NCNN_SHADER_COMPRESS
+    delete[] comp_data;
+#endif
 
-    return compile_spirv_module(comp_data, comp_data_size, opt, spirv);
+    return ret;
 }
 
 int resolve_shader_info(const uint32_t* spv_data, size_t spv_data_size, ShaderInfo& shader_info)

src/layer/vulkan/shader/shader_compress_dict.in

@@ -0,0 +1,66 @@
+#define shape_constant_id_offset
+layout (constant_id = shape_constant_id_offset + 
+layout (constant_id = 
+layout (binding = 
+layout (push_constant) uniform parameter
+layout
+void main()
+if (gx >= psc(w) || gy >= psc(h) || gz >= psc(c))
+int gx = int(gl_GlobalInvocationID.x);
+int gy = int(gl_GlobalInvocationID.y);
+int gz = int(gl_GlobalInvocationID.z);
+const int gi = gz * psc(cstep) + gy * psc(w) + gx;
+gl_GlobalInvocationID
+afp v = buffer_ld1(bottom_top_blob_data, gi);
+buffer_st1(bottom_top_blob_data, gi, v);
+afpvec4 v = buffer_ld4(bottom_top_blob_data, gi);
+buffer_st4(bottom_top_blob_data, gi, v);
+bottom_top_blob_data
+bottom_top_blob
+top_blob_data
+dims
+cstep
+out
+scale
+axis
+buffer_ld8(
+buffer_ld4(
+buffer_ld1(
+buffer
+const
+int
+uint
+afpvec4
+sfpvec8
+sfpvec4
+sfp
+afp
+readonly
+writeonly
+sum
+bias
+offset
+psc
+if
+else
+for
+extension
+return
+#version 450
+matrix
+fcoopmatNV
+vec4
+coop
+fp16
+order
+type
+blob
+data
+v_offset
+gl_ScopeSubgroup
+activation_type
+CooperativeMatrix
+max
+min
+abs
+slope

src/platform.h.in

@@ -57,6 +57,7 @@
 #cmakedefine01 NCNN_INT8
 #cmakedefine01 NCNN_BF16
 #cmakedefine01 NCNN_FORCE_INLINE
+#cmakedefine01 NCNN_SHADER_COMPRESS
 
 #cmakedefine NCNN_VERSION_STRING "@NCNN_VERSION_STRING@"