KBaba7/llama.cpp

1 2 3#include <iostream>4#include <fstream>5#include <sstream>6#include <string>7#include <stdexcept>8#include <array>9#include <vector>10#include <map>11#include <thread>12#include <mutex>13#include <future>14#include <queue>15#include <condition_variable>16#include <cstdio>17#include <cstring>18#include <cstdlib>19#include <cassert>20#include <algorithm>21#include <sys/stat.h>22#include <sys/types.h>23 24#ifdef _WIN3225    #include <windows.h>26    #include <direct.h> // For _mkdir on Windows27#else28    #include <unistd.h>29    #include <sys/wait.h>30    #include <fcntl.h>31#endif32 33#define ASYNCIO_CONCURRENCY 6434 35std::mutex lock;36std::vector<std::pair<std::string, std::string>> shader_fnames;37 38std::string GLSLC = "glslc";39std::string input_dir = "vulkan-shaders";40std::string output_dir = "/tmp";41std::string target_hpp = "ggml-vulkan-shaders.hpp";42std::string target_cpp = "ggml-vulkan-shaders.cpp";43bool no_clean = false;44 45const std::vector<std::string> type_names = {46    "f32",47    "f16",48    "q4_0",49    "q4_1",50    "q5_0",51    "q5_1",52    "q8_0",53    "q2_k",54    "q3_k",55    "q4_k",56    "q5_k",57    "q6_k",58    "iq2_xxs",59    "iq2_xs",60    "iq2_s",61    "iq3_xxs",62    "iq3_s",63    "iq4_xs",64    "iq4_nl"65};66 67namespace {68void execute_command(const std::string& command, std::string& stdout_str, std::string& stderr_str) {69#ifdef _WIN3270    HANDLE stdout_read, stdout_write;71    HANDLE stderr_read, stderr_write;72    SECURITY_ATTRIBUTES sa = { sizeof(SECURITY_ATTRIBUTES), NULL, TRUE };73 74    if (!CreatePipe(&stdout_read, &stdout_write, &sa, 0) ||75        !SetHandleInformation(stdout_read, HANDLE_FLAG_INHERIT, 0)) {76        throw std::runtime_error("Failed to create stdout pipe");77    }78 79    if (!CreatePipe(&stderr_read, &stderr_write, &sa, 0) ||80        !SetHandleInformation(stderr_read, HANDLE_FLAG_INHERIT, 0)) {81        throw std::runtime_error("Failed to create stderr pipe");82    }83 84    PROCESS_INFORMATION pi;85    STARTUPINFOA si = {};86    si.cb = sizeof(STARTUPINFOA);87    si.dwFlags = STARTF_USESTDHANDLES;88    si.hStdOutput = stdout_write;89    si.hStdError = stderr_write;90 91    std::vector<char> cmd(command.begin(), command.end());92    cmd.push_back('\0');93 94    if (!CreateProcessA(NULL, cmd.data(), NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) {95        throw std::runtime_error("Failed to create process");96    }97 98    CloseHandle(stdout_write);99    CloseHandle(stderr_write);100 101    std::array<char, 128> buffer;102    DWORD bytes_read;103 104    while (ReadFile(stdout_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) {105        stdout_str.append(buffer.data(), bytes_read);106    }107 108    while (ReadFile(stderr_read, buffer.data(), (DWORD)buffer.size(), &bytes_read, NULL) && bytes_read > 0) {109        stderr_str.append(buffer.data(), bytes_read);110    }111 112    CloseHandle(stdout_read);113    CloseHandle(stderr_read);114    WaitForSingleObject(pi.hProcess, INFINITE);115    CloseHandle(pi.hProcess);116    CloseHandle(pi.hThread);117#else118int stdout_pipe[2];119    int stderr_pipe[2];120 121    if (pipe(stdout_pipe) != 0 || pipe(stderr_pipe) != 0) {122        throw std::runtime_error("Failed to create pipes");123    }124 125    pid_t pid = fork();126    if (pid < 0) {127        throw std::runtime_error("Failed to fork process");128    }129 130    if (pid == 0) {131        close(stdout_pipe[0]);132        close(stderr_pipe[0]);133        dup2(stdout_pipe[1], STDOUT_FILENO);134        dup2(stderr_pipe[1], STDERR_FILENO);135        close(stdout_pipe[1]);136        close(stderr_pipe[1]);137        execl("/bin/sh", "sh", "-c", command.c_str(), (char*) nullptr);138        _exit(EXIT_FAILURE);139    } else {140        close(stdout_pipe[1]);141        close(stderr_pipe[1]);142 143        std::array<char, 128> buffer;144        ssize_t bytes_read;145 146        while ((bytes_read = read(stdout_pipe[0], buffer.data(), buffer.size())) > 0) {147            stdout_str.append(buffer.data(), bytes_read);148        }149 150        while ((bytes_read = read(stderr_pipe[0], buffer.data(), buffer.size())) > 0) {151            stderr_str.append(buffer.data(), bytes_read);152        }153 154        close(stdout_pipe[0]);155        close(stderr_pipe[0]);156        waitpid(pid, nullptr, 0);157    }158#endif159}160 161bool directory_exists(const std::string& path) {162    struct stat info;163    if (stat(path.c_str(), &info) != 0) {164        return false; // Path doesn't exist or can't be accessed165    }166    return (info.st_mode & S_IFDIR) != 0; // Check if it is a directory167}168 169bool create_directory(const std::string& path) {170#ifdef _WIN32171    return _mkdir(path.c_str()) == 0 || errno == EEXIST; // EEXIST means the directory already exists172#else173    return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST; // 0755 is the directory permissions174#endif175}176 177std::string to_uppercase(const std::string& input) {178    std::string result = input;179    for (char& c : result) {180        c = std::toupper(c);181    }182    return result;183}184 185bool string_ends_with(const std::string& str, const std::string& suffix) {186    if (suffix.size() > str.size()) {187        return false;188    }189    return std::equal(suffix.rbegin(), suffix.rend(), str.rbegin());190}191 192static const char path_separator = '/';193 194std::string join_paths(const std::string& path1, const std::string& path2) {195    return path1 + path_separator + path2;196}197 198std::string basename(const std::string &path) {199    return path.substr(path.find_last_of("/\\") + 1);200}201 202// variables to track number of compiles in progress203static uint32_t compile_count = 0;204static std::mutex compile_count_mutex;205static std::condition_variable compile_count_cond;206 207void string_to_spv_func(const std::string& _name, const std::string& in_fname, const std::map<std::string, std::string>& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) {208    std::string name = _name + (f16acc ? "_f16acc" : "") + (coopmat ? "_coopmat" : "") + (coopmat2 ? "_cm2" : (fp16 ? "" : "_fp32"));209    std::string out_fname = join_paths(output_dir, name + ".spv");210    std::string in_path = join_paths(input_dir, in_fname);211 212    std::string target_env = (name.find("_cm2") != std::string::npos) ? "--target-env=vulkan1.3" : "--target-env=vulkan1.2";213 214    // disable spirv-opt for coopmat shaders for https://github.com/ggerganov/llama.cpp/issues/10734215    std::string opt_level = coopmat ? "" : "-O";216 217    #ifdef _WIN32218        std::vector<std::string> cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, "\"" + in_path + "\"", "-o", "\"" + out_fname + "\""};219    #else220        std::vector<std::string> cmd = {GLSLC, "-fshader-stage=compute", target_env, opt_level, in_path, "-o",  out_fname};221    #endif222 223    #ifdef GGML_VULKAN_SHADER_DEBUG_INFO224        cmd.push_back("-g");225    #endif226 227    for (const auto& define : defines) {228        cmd.push_back("-D" + define.first + "=" + define.second);229    }230 231    std::string command;232    for (const auto& part : cmd) {233        command += part + " ";234    }235 236    std::string stdout_str, stderr_str;237    try {238        // std::cout << "Executing command: ";239        // for (const auto& part : cmd) {240        //     std::cout << part << " ";241        // }242        // std::cout << std::endl;243 244        execute_command(command, stdout_str, stderr_str);245        if (!stderr_str.empty()) {246            std::cerr << "cannot compile " << name << "\n\n" << command << "\n\n" << stderr_str << std::endl;247            return;248        }249 250        std::lock_guard<std::mutex> guard(lock);251        shader_fnames.push_back(std::make_pair(name, out_fname));252    } catch (const std::exception& e) {253        std::cerr << "Error executing command for " << name << ": " << e.what() << std::endl;254    }255    {256        std::lock_guard<std::mutex> guard(compile_count_mutex);257        assert(compile_count > 0);258        compile_count--;259    }260    compile_count_cond.notify_all();261}262 263std::map<std::string, std::string> merge_maps(const std::map<std::string, std::string>& a, const std::map<std::string, std::string>& b) {264    std::map<std::string, std::string> result = a;265    result.insert(b.begin(), b.end());266    return result;267}268 269static std::vector<std::future<void>> compiles;270void string_to_spv(const std::string& _name, const std::string& in_fname, const std::map<std::string, std::string>& defines, bool fp16 = true, bool coopmat = false, bool coopmat2 = false, bool f16acc = false) {271    {272        // wait until fewer than N compiles are in progress.273        // 16 is an arbitrary limit, the goal is to avoid "failed to create pipe" errors.274        uint32_t N = 16;275        std::unique_lock<std::mutex> guard(compile_count_mutex);276        while (compile_count >= N) {277            compile_count_cond.wait(guard);278        }279        compile_count++;280    }281    compiles.push_back(std::async(string_to_spv_func, _name, in_fname, defines, fp16, coopmat, coopmat2, f16acc));282}283 284void matmul_shaders(bool fp16, bool matmul_id, bool coopmat, bool coopmat2, bool f16acc) {285    std::string load_vec = coopmat2 ? "1" : fp16 ? "8" : "4";286    std::string aligned_b_type_f32 = coopmat2 ? "float" : fp16 ? "mat2x4" : "vec4";287    std::string aligned_b_type_f16 = coopmat2 ? "float16_t" : fp16 ? "f16mat2x4" : "f16vec4";288 289    std::map<std::string, std::string> base_dict = {{"FLOAT_TYPE", (coopmat2 || fp16) ? "float16_t" : "float"}};290    std::string shader_name = "matmul";291 292    if (matmul_id) {293        base_dict["MUL_MAT_ID"] = "1";294        shader_name = "matmul_id";295    }296 297    if (fp16) {298        base_dict["FLOAT16"] = "1";299    }300 301    base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float";302 303    if (coopmat) {304        base_dict["COOPMAT"] = "1";305    }306 307    base_dict["ACC_TYPE"] = f16acc ? "float16_t" : "float";308 309    std::string source_name = coopmat2 ? "mul_mm_cm2.comp" : "mul_mm.comp";310 311    // Shaders with f16 B_TYPE312    string_to_spv(shader_name + "_f32_f16", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}, }), fp16, coopmat, coopmat2, f16acc);313    string_to_spv(shader_name + "_f32_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc);314 315    string_to_spv(shader_name + "_f16_aligned", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc);316    string_to_spv(shader_name + "_f16", source_name, merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);317 318    for (const auto& tname : type_names) {319        std::string data_a_key = "DATA_A_" + to_uppercase(tname);320        // For unaligned, load one at a time for f32/f16, or two at a time for quants321        std::string load_vec_a_unaligned = (coopmat2 || tname == "f32" || tname == "f16") ? "1" : "2";322        // For aligned matmul loads323        std::string load_vec_a = (coopmat2 || tname == "f32" || tname == "f16") ? load_vec : "2";324 325        // don't generate f32 variants for coopmat2326        if (!coopmat2) {327            string_to_spv(shader_name + "_" + tname + "_f32", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc);328            string_to_spv(shader_name + "_" + tname + "_f32_aligned", source_name, merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc);329        }330 331        if (tname != "f16" && tname != "f32") {332            string_to_spv(shader_name + "_" + tname + "_f16", source_name,          merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned},                           {"B_TYPE", "float16_t"},        {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}}), fp16, coopmat, coopmat2, f16acc);333            string_to_spv(shader_name + "_" + tname + "_f16_aligned", source_name,  merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a},           {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc);334        }335    }336}337 338void process_shaders() {339    std::cout << "ggml_vulkan: Generating and compiling shaders to SPIR-V" << std::endl;340    std::map<std::string, std::string> base_dict = {{"FLOAT_TYPE", "float"}};341 342    // matmul343    for (const auto& matmul_id : {false, true}) {344        // No coopmats345        // fp32346        matmul_shaders(false, matmul_id, false, false, false);347 348        // fp16, fp32acc and fp16acc349        matmul_shaders(true, matmul_id, false, false, false);350        matmul_shaders(true, matmul_id, false, false, true);351 352#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)353        // Coopmat, fp32acc and fp16acc354        matmul_shaders(true, matmul_id, true, false, false);355        matmul_shaders(true, matmul_id, true, false, true);356#endif357 358#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT)359        // Coopmat2, fp32acc and fp16acc360        matmul_shaders(true, matmul_id, false, true, false);361        matmul_shaders(true, matmul_id, false, true, true);362#endif363    }364 365#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT)366    // flash attention367    for (const auto& f16acc : {false, true}) {368        std::string acctype = f16acc ? "float16_t" : "float";369 370        for (const auto& tname : type_names) {371            if (tname == "f32") {372                continue;373            }374 375            if (tname == "f16") {376                string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp",377                    merge_maps(base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}}), true, false, true, f16acc);378            } else {379                std::string data_a_key = "DATA_A_" + to_uppercase(tname);380                string_to_spv("flash_attn_f32_f16_" + tname, "flash_attn_cm2.comp",381                    merge_maps(base_dict, {{data_a_key, "1"}, {"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"ACC_TYPE", acctype}, {"DEQUANTFUNC", "dequantFunc"+to_uppercase(tname) }, {"BLOCK_SIZE", "QUANT_K_"+to_uppercase(tname) }}), true, false, true, f16acc);382            }383        }384    }385#endif386 387    for (const auto& tname : type_names) {388        // mul mat vec389        std::string data_a_key = "DATA_A_" + to_uppercase(tname);390        std::string shader = (string_ends_with(tname, "_k")) ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp";391 392        string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}}));393        string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPE_VEC2", "f16vec2"}, {"B_TYPE_VEC4", "f16vec4"}, {"D_TYPE", "float"}}));394 395        string_to_spv("mul_mat_vec_id_" + tname + "_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPE_VEC2", "vec2"}, {"B_TYPE_VEC4", "vec4"}, {"D_TYPE", "float"}}));396 397        // Dequant shaders398        if (tname != "f16") {399            string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}}));400        }401 402        if (!string_ends_with(tname, "_k")) {403            shader = (tname == "f32" || tname == "f16") ? "get_rows.comp" : "get_rows_quant.comp";404 405            if (tname == "f16") {406                string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}));407            } else {408                string_to_spv("get_rows_" + tname, shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}}));409            }410            string_to_spv("get_rows_" + tname + "_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float"}}));411        }412    }413 414    string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});415    string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});416 417    // Norms418    string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));419    string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));420    string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));421 422    string_to_spv("cpy_f32_f32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});423    string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}});424    string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}});425    string_to_spv("contig_cpy_f32_f32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});426    string_to_spv("contig_cpy_f32_f16", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}});427    string_to_spv("contig_cpy_f16_f16", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}});428 429    for (std::string t : {"q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {430        string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});431        string_to_spv("cpy_" + t + "_f32", "copy_from_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});432    }433 434    string_to_spv("add_f32", "add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});435    string_to_spv("add_f16_f32_f16", "add.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"FLOAT_TYPE", "float"}});436 437    string_to_spv("acc_f32", "acc.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});438 439    string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {});440 441    string_to_spv("mul_f32", "mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});442 443    string_to_spv("div_f32", "div.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});444 445    string_to_spv("repeat_f32", "repeat.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});446 447    string_to_spv("scale_f32", "scale.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});448 449    string_to_spv("sqr_f32", "square.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});450 451    string_to_spv("sin_f32", "sin.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});452 453    string_to_spv("cos_f32", "cos.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});454 455    string_to_spv("clamp_f32", "clamp.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});456 457    string_to_spv("pad_f32", "pad.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});458 459    string_to_spv("concat_f32", "concat.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});460    string_to_spv("concat_f16", "concat.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}});461    string_to_spv("concat_i32", "concat.comp", {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}});462 463    string_to_spv("upscale_f32", "upscale.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});464 465    string_to_spv("gelu_f32", "gelu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});466    string_to_spv("gelu_quick_f32", "gelu_quick.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});467    string_to_spv("silu_f32", "silu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});468    string_to_spv("relu_f32", "relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});469    string_to_spv("leaky_relu_f32", "leaky_relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});470    string_to_spv("tanh_f32", "tanh.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});471 472    string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});473 474    string_to_spv("soft_max_f32", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));475    string_to_spv("soft_max_f32_f16", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}));476 477    string_to_spv("rope_norm_f32", "rope_norm.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});478    string_to_spv("rope_norm_f16", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});479    string_to_spv("rope_norm_f16_rte", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}});480 481    string_to_spv("rope_neox_f32", "rope_neox.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});482    string_to_spv("rope_neox_f16", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});483    string_to_spv("rope_neox_f16_rte", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}});484 485    string_to_spv("argsort_f32", "argsort.comp", {{"A_TYPE", "float"}});486 487    string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));488 489    string_to_spv("im2col_f32", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));490    string_to_spv("im2col_f32_f16", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}));491    string_to_spv("im2col_f32_f16_rte", "im2col.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"RTE16", "1"}}));492 493    string_to_spv("timestep_embedding_f32", "timestep_embedding.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));494 495    string_to_spv("pool2d_f32", "pool2d.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));496 497    string_to_spv("rwkv_wkv6_f32", "wkv6.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));498 499    for (auto &c : compiles) {500        c.wait();501    }502}503 504void write_output_files() {505    FILE* hdr = fopen(target_hpp.c_str(), "w");506    FILE* src = fopen(target_cpp.c_str(), "w");507 508    fprintf(hdr, "#include <cstdint>\n\n");509    fprintf(src, "#include \"%s\"\n\n", basename(target_hpp).c_str());510 511    std::sort(shader_fnames.begin(), shader_fnames.end());512    for (const auto& pair : shader_fnames) {513        const std::string& name = pair.first;514        #ifdef _WIN32515            std::string path = pair.second;516            std::replace(path.begin(), path.end(), '/', '\\' );517        #else518            const std::string& path = pair.second;519        #endif520 521        FILE* spv = fopen(path.c_str(), "rb");522        if (!spv) {523            std::cerr << "Error opening SPIR-V file: " << path << " (" << strerror(errno) << ")\n";524            continue;525        }526 527        fseek(spv, 0, SEEK_END);528        size_t size = ftell(spv);529        fseek(spv, 0, SEEK_SET);530 531        std::vector<unsigned char> data(size);532        size_t read_size = fread(data.data(), 1, size, spv);533        fclose(spv);534        if (read_size != size) {535            std::cerr << "Error reading SPIR-V file: " << path << " (" << strerror(errno) << ")\n";536            continue;537        }538 539        fprintf(hdr, "extern unsigned char %s_data[%zu];\n", name.c_str(), size);540        fprintf(hdr, "const uint64_t %s_len = %zu;\n\n", name.c_str(), size);541 542        fprintf(src, "unsigned char %s_data[%zu] = {\n", name.c_str(), size);543        for (size_t i = 0; i < size; ++i) {544            fprintf(src, "0x%02x,", data[i]);545            if ((i + 1) % 12 == 0) fprintf(src, "\n");546        }547        fprintf(src, "\n};\n\n");548 549        if (!no_clean) {550            std::remove(path.c_str());551        }552    }553 554    fclose(hdr);555    fclose(src);556}557}558 559int main(int argc, char** argv) {560    std::map<std::string, std::string> args;561    for (int i = 1; i < argc; ++i) {562        std::string arg = argv[i];563        if (arg.rfind("--", 0) == 0) {564            if (i + 1 < argc && argv[i + 1][0] != '-') {565                args[arg] = argv[i + 1];566                ++i;567            } else {568                args[arg] = "";569            }570        }571    }572 573    if (args.find("--glslc") != args.end()) {574        GLSLC = args["--glslc"]; // Path to glslc575    }576    if (args.find("--input-dir") != args.end()) {577        input_dir = args["--input-dir"]; // Directory containing shader sources578    }579    if (args.find("--output-dir") != args.end()) {580        output_dir = args["--output-dir"]; // Directory for containing SPIR-V output581    }582    if (args.find("--target-hpp") != args.end()) {583        target_hpp = args["--target-hpp"]; // Path to generated header file584    }585    if (args.find("--target-cpp") != args.end()) {586        target_cpp = args["--target-cpp"]; // Path to generated cpp file587    }588    if (args.find("--no-clean") != args.end()) {589        no_clean = true; // Keep temporary SPIR-V files in output-dir after build590    }591 592    if (!directory_exists(input_dir)) {593        std::cerr << "\"" << input_dir << "\" must be a valid directory containing shader sources" << std::endl;594        return EXIT_FAILURE;595    }596 597    if (!directory_exists(output_dir)) {598        if (!create_directory(output_dir)) {599            std::cerr << "Error creating output directory: " << output_dir << "\n";600            return EXIT_FAILURE;601        }602    }603 604    process_shaders();605 606    write_output_files();607 608    return EXIT_SUCCESS;609}610