From ac9a065d3144b2dff9a0a0de1b075bf7ff421a0e Mon Sep 17 00:00:00 2001 From: Mason M Date: Sun, 30 Jun 2024 00:02:41 -0300 Subject: [PATCH] Remove Python dependency from Vulkan build --- .devops/nix/package.nix | 1 - Makefile | 14 +- README.md | 3 +- ggml/ggml_vk_generate_shaders.py | 237 -------- ggml/src/CMakeLists.txt | 9 +- ggml/src/vulkan-gen-shaders/CMakeLists.txt | 5 + .../vulkan-gen-shaders/vulkan-gen-shaders.cpp | 525 ++++++++++++++++++ 7 files changed, 544 insertions(+), 250 deletions(-) delete mode 100644 ggml/ggml_vk_generate_shaders.py create mode 100644 ggml/src/vulkan-gen-shaders/CMakeLists.txt create mode 100644 ggml/src/vulkan-gen-shaders/vulkan-gen-shaders.cpp diff --git a/.devops/nix/package.nix b/.devops/nix/package.nix index 6e88f844d..d75326e65 100644 --- a/.devops/nix/package.nix +++ b/.devops/nix/package.nix @@ -134,7 +134,6 @@ let vulkan-headers vulkan-loader shaderc - python3 ]; in diff --git a/Makefile b/Makefile index 60cfc65a5..a0b69392a 100644 --- a/Makefile +++ b/Makefile @@ -186,6 +186,10 @@ ifdef GGML_RPC BUILD_TARGETS += rpc-server endif +ifdef GGML_VULKAN + BUILD_TARGETS += vulkan-gen-shaders +endif + default: $(BUILD_TARGETS) test: $(TEST_TARGETS) @@ -711,9 +715,8 @@ ifdef GGML_VULKAN_RUN_TESTS MK_CPPFLAGS += -DGGML_VULKAN_RUN_TESTS endif -PYTHON_CMD = python GLSLC_CMD = glslc -_llama_vk_genshaders_cmd = $(PYTHON_CMD) ggml/ggml_vk_generate_shaders.py +_llama_vk_genshaders_cmd = $(shell pwd)/vulkan-gen-shaders _llama_vk_header = ggml/src/ggml-vulkan-shaders.hpp _llama_vk_source = ggml/src/ggml-vulkan-shaders.cpp _llama_vk_input_dir = ggml/src/vulkan-shaders @@ -724,13 +727,16 @@ ggml/src/ggml-vulkan.o: ggml/src/ggml-vulkan.cpp ggml/include/ggml-vulkan.h $(_l $(_llama_vk_header): $(_llama_vk_source) -$(_llama_vk_source): $(_llama_vk_shader_deps) +$(_llama_vk_source): $(_llama_vk_shader_deps) vulkan-gen-shaders $(_llama_vk_genshaders_cmd) \ --glslc $(GLSLC_CMD) \ --input-dir $(_llama_vk_input_dir) \ --target-hpp $(_llama_vk_header) \ --target-cpp $(_llama_vk_source) +vulkan-gen-shaders: ggml/src/vulkan-gen-shaders/vulkan-gen-shaders.cpp + $(CXX) $(CXXFLAGS) -o $@ $(LDFLAGS) ggml/src/vulkan-gen-shaders/vulkan-gen-shaders.cpp + endif # GGML_VULKAN ifdef GGML_HIPBLAS @@ -1100,7 +1106,7 @@ clean: rm -vrf ggml/src/ggml-cuda/template-instances/*.o rm -rvf $(BUILD_TARGETS) rm -rvf $(TEST_TARGETS) - rm -f ggml/src/ggml-vulkan-shaders.hpp ggml/src/ggml-vulkan-shaders.cpp + rm -f vulkan-gen-shaders ggml/src/ggml-vulkan-shaders.hpp ggml/src/ggml-vulkan-shaders.cpp find examples pocs -type f -name "*.o" -delete # diff --git a/README.md b/README.md index cb6f83e70..d161ed55c 100644 --- a/README.md +++ b/README.md @@ -586,8 +586,7 @@ Building the program with BLAS support may lead to some performance improvements mingw-w64-ucrt-x86_64-gcc \ mingw-w64-ucrt-x86_64-cmake \ mingw-w64-ucrt-x86_64-vulkan-devel \ - mingw-w64-ucrt-x86_64-shaderc \ - mingw-w64-ucrt-x86_64-python3 + mingw-w64-ucrt-x86_64-shaderc ``` Switch into `llama.cpp` directory and build using CMake. ```sh diff --git a/ggml/ggml_vk_generate_shaders.py b/ggml/ggml_vk_generate_shaders.py deleted file mode 100644 index ecd11f25d..000000000 --- a/ggml/ggml_vk_generate_shaders.py +++ /dev/null @@ -1,237 +0,0 @@ -#!/usr/bin/env python - -import logging -import argparse -import asyncio -import os -from tempfile import gettempdir -import sys - -logger = logging.getLogger("ggml-vk-generate-shaders") - -GLSLC = "glslc" - -type_names = [ - "f32", - "f16", - "q4_0", - "q4_1", - "q5_0", - "q5_1", - "q8_0", - "q2_k", - "q3_k", - "q4_k", - "q5_k", - "q6_k", -] - -ASYNCIO_CONCURRENCY = 64 - -lock = asyncio.Lock() -shader_fnames = [] - - -async def string_to_spv(name, in_fname, defines, fp16=True): - name = f"{name}{'_fp32' if not fp16 else ''}" - out_fname = os.path.join(output_dir, f"{name}.spv") - - in_path = os.path.join(input_dir, in_fname) - - cmd = [GLSLC, "-fshader-stage=compute", "--target-env=vulkan1.2", "-O", in_path, "-o", out_fname] - - cmd.extend([f"-D{key}={value}" for key, value in defines.items()]) - - proc = await asyncio.create_subprocess_exec(*cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE) - - stdout, stderr = await proc.communicate() - - stdout = stdout.decode() - error = stderr.decode() - - if proc.returncode: - cmd = " ".join(cmd) - logger.error(f"cannot compile {name}\n\n{cmd}\n\n{error}") - return - - async with lock: - shader_fnames.append((name, out_fname)) - - -def matmul_shaders(tasks, fp16, matmul_id): - if fp16: - load_vec = "8" - aligned_b_type_f32 = "mat2x4" - aligned_b_type_f16 = "f16mat2x4" - else: - load_vec = "4" - aligned_b_type_f32 = "vec4" - aligned_b_type_f16 = "f16vec4" - - base_dict = {"FLOAT_TYPE": "float" if not fp16 else "float16_t"} - shader_name = "matmul" - - if matmul_id: - base_dict["MUL_MAT_ID"] = "1" - shader_name = "matmul_id" - - if fp16: - base_dict["FLOAT16"] = "1" - - # Shaders with f16 B_TYPE - tasks.append(string_to_spv(f"{shader_name}_f32_f16", "mul_mm.comp", base_dict | {"DATA_A_F32": "1", "B_TYPE": "float16_t", "D_TYPE": "float"}, fp16)) - tasks.append(string_to_spv(f"{shader_name}_f32_f16_aligned", "mul_mm.comp", 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"}, fp16)) - - tasks.append(string_to_spv(f"{shader_name}_f16", "mul_mm.comp", base_dict | {"DATA_A_F16": "1", "B_TYPE": "float16_t", "D_TYPE": "float"}, fp16)) - tasks.append(string_to_spv(f"{shader_name}_f16_aligned", "mul_mm.comp", 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"}, fp16)) - - for tname in type_names: - data_a_key = f"DATA_A_{tname.upper()}" - load_vec_a = load_vec if tname in ("f32", "f16") else "2" - tasks.append(string_to_spv(f"{shader_name}_{tname}_f32", "mul_mm.comp", base_dict | {data_a_key: "1", "B_TYPE": "float", "D_TYPE": "float"}, fp16)) - tasks.append(string_to_spv(f"{shader_name}_{tname}_f32_aligned", "mul_mm.comp", base_dict | {data_a_key: "2", "LOAD_VEC_A": load_vec_a, "LOAD_VEC_B": load_vec, "B_TYPE": aligned_b_type_f32, "D_TYPE": "float"}, fp16)) - - -async def main(): - logger.info("ggml_vulkan: Generating and compiling shaders to SPIR-V") - - tasks = [] - - for fp16 in (False, True): - # MUL_MAT - matmul_shaders(tasks, fp16, False) - # MUL_MAT_ID - matmul_shaders(tasks, fp16, True) - - for tname in type_names: - base_dict = {"FLOAT_TYPE": "float"} - - # mul mat vec - data_a_key = f"DATA_A_{tname.upper()}" - shader = f"mul_mat_vec_{tname}.comp" if tname.endswith("_k") else "mul_mat_vec.comp" - - tasks.append(string_to_spv(f"mul_mat_vec_{tname}_f32_f32", shader, base_dict | {data_a_key: "1", "B_TYPE": "float", "D_TYPE": "float"})) - tasks.append(string_to_spv(f"mul_mat_vec_{tname}_f16_f32", shader, base_dict | {data_a_key: "1", "B_TYPE": "float16_t", "D_TYPE": "float"})) - - tasks.append(string_to_spv(f"mul_mat_vec_id_{tname}_f32", shader, base_dict | {"MUL_MAT_ID": "1", data_a_key: "1", "B_TYPE": "float", "D_TYPE": "float"})) - - # Dequant shaders - if tname != "f16": - tasks.append(string_to_spv(f"dequant_{tname}", f"dequant_{tname}.comp", base_dict | {data_a_key: "1", "D_TYPE": "float16_t"})) - - # get_rows - if not tname.endswith("_k"): - shader = "get_rows.comp" if tname in ("f32", "f16") else "get_rows_quant.comp" - - if tname == "f16": - tasks.append(string_to_spv(f"get_rows_{tname}", shader, {data_a_key: "1", "B_TYPE": "int", "D_TYPE": "float16_t", "OPTIMIZATION_ERROR_WORKAROUND": "1"})) - else: - tasks.append(string_to_spv(f"get_rows_{tname}", shader, {data_a_key: "1", "B_TYPE": "int", "D_TYPE": "float16_t"})) - tasks.append(string_to_spv(f"get_rows_{tname}_f32", shader, {data_a_key: "1", "B_TYPE": "int", "D_TYPE": "float"})) - - tasks.append(string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {"A_TYPE": "float16_t", "B_TYPE": "float", "D_TYPE": "float"})) - tasks.append(string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {"A_TYPE": "float16_t", "B_TYPE": "float", "D_TYPE": "float"})) - - # Norms - tasks.append(string_to_spv("norm_f32", "norm.comp", base_dict | {"A_TYPE": "float", "D_TYPE": "float"})) - tasks.append(string_to_spv("rms_norm_f32", "rms_norm.comp", base_dict | {"A_TYPE": "float", "D_TYPE": "float"})) - - tasks.append(string_to_spv("cpy_f32_f32", "copy.comp", {"A_TYPE": "float", "D_TYPE": "float"})) - tasks.append(string_to_spv("cpy_f32_f16", "copy.comp", {"A_TYPE": "float", "D_TYPE": "float16_t"})) - tasks.append(string_to_spv("cpy_f16_f16", "copy.comp", {"A_TYPE": "float16_t", "D_TYPE": "float16_t", "OPTIMIZATION_ERROR_WORKAROUND": "1"})) - - tasks.append(string_to_spv("add_f32", "add.comp", {"A_TYPE": "float", "B_TYPE": "float", "D_TYPE": "float", "FLOAT_TYPE": "float"})) - - tasks.append(string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {})) - - tasks.append(string_to_spv("mul_f32", "mul.comp", {"A_TYPE": "float", "B_TYPE": "float", "D_TYPE": "float", "FLOAT_TYPE": "float"})) - - tasks.append(string_to_spv("div_f32", "div.comp", {"A_TYPE": "float", "B_TYPE": "float", "D_TYPE": "float", "FLOAT_TYPE": "float"})) - - tasks.append(string_to_spv("scale_f32", "scale.comp", {"A_TYPE": "float", "D_TYPE": "float", "FLOAT_TYPE": "float"})) - - tasks.append(string_to_spv("sqr_f32", "square.comp", {"A_TYPE": "float", "D_TYPE": "float", "FLOAT_TYPE": "float"})) - - tasks.append(string_to_spv("clamp_f32", "clamp.comp", {"A_TYPE": "float", "D_TYPE": "float", "FLOAT_TYPE": "float"})) - - tasks.append(string_to_spv("gelu_f32", "gelu.comp", {"A_TYPE": "float", "D_TYPE": "float"})) - tasks.append(string_to_spv("silu_f32", "silu.comp", {"A_TYPE": "float", "D_TYPE": "float"})) - tasks.append(string_to_spv("relu_f32", "relu.comp", {"A_TYPE": "float", "D_TYPE": "float"})) - - tasks.append(string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {"A_TYPE": "float", "D_TYPE": "float"})) - - tasks.append(string_to_spv("soft_max_f32", "soft_max.comp", base_dict | {"A_TYPE": "float", "B_TYPE": "float", "D_TYPE": "float"})) - tasks.append(string_to_spv("soft_max_f32_f16", "soft_max.comp", base_dict | {"A_TYPE": "float", "B_TYPE": "float16_t", "D_TYPE": "float"})) - - tasks.append(string_to_spv("rope_norm_f32", "rope_norm.comp", {"A_TYPE": "float", "D_TYPE": "float"})) - tasks.append(string_to_spv("rope_norm_f16", "rope_norm.comp", {"A_TYPE": "float16_t", "D_TYPE": "float16_t"})) - - tasks.append(string_to_spv("rope_neox_f32", "rope_neox.comp", {"A_TYPE": "float", "D_TYPE": "float"})) - tasks.append(string_to_spv("rope_neox_f16", "rope_neox.comp", {"A_TYPE": "float16_t", "D_TYPE": "float16_t"})) - - tasks.append(string_to_spv("argsort_f32", "argsort.comp", {"A_TYPE": "float"})) - - tasks.append(string_to_spv("sum_rows_f32", "sum_rows.comp", base_dict | {"A_TYPE": "float", "D_TYPE": "float"})) - - # Helper to decorate tasks with semaphore acquisition. - async def withSemaphore(sem, task): - async with sem: - return await task - - # Run tasks concurrently guarded by a concurrency limit. - sem = asyncio.Semaphore(ASYNCIO_CONCURRENCY) - await asyncio.gather(*(withSemaphore(sem, task) for task in tasks)) - - with open(target_hpp, "w") as hdr, open(target_cpp, "w") as src: - hdr.write("#include \n\n") - src.write(f"#include \"{os.path.basename(target_hpp)}\"\n\n") - for name, path in sorted(shader_fnames): - with open(path, "rb") as spv: - counter = 0 - newline_counter = 0 - data = "" - for val in spv.read(): - data += f"0x{val:02x}," - newline_counter += 1 - counter += 1 - if newline_counter >= 12: - newline_counter = 0 - data += "\n" - hdr.write(f"extern unsigned char {name}_data[{counter}];\n") - hdr.write(f"const uint64_t {name}_len = {counter};\n\n") - src.write(f"unsigned char {name}_data[{counter}] = {{\n{data}\n}};\n\n") - if not no_clean: - os.remove(path) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="GGML Vulkan Shader Generator") - - parser.add_argument("--glslc", help="Path to glslc") - parser.add_argument("--input-dir", default="vulkan-shaders", help="Directory containing shader sources") - parser.add_argument("--output-dir", default=gettempdir(), help="Directory for containing SPIR-V output") - parser.add_argument("--target-hpp", default="ggml-vulkan-shaders.hpp", help="Path to generated header file") - parser.add_argument("--target-cpp", default="ggml-vulkan-shaders.cpp", help="Path to generated cpp file") - parser.add_argument("--no-clean", action="store_true", help="Keep temporary SPIR-V files in output-dir after build") - parser.add_argument("--verbose", action="store_true", help="increase output verbosity") - - args = parser.parse_args() - - logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) - - if args.glslc: - GLSLC = args.glslc - - input_dir = args.input_dir - if not os.path.isdir(input_dir): - sys.exit(f"\"{input_dir}\" must be a valid directory containing shader sources") - - output_dir = args.output_dir - if not os.path.isdir(output_dir): - os.makedirs(output_dir) - - target_hpp = args.target_hpp - target_cpp = args.target_cpp - no_clean = args.no_clean - - asyncio.run(main()) diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 1a6c72928..ea7dfa773 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -525,11 +525,6 @@ if (GGML_RPC) endif() if (GGML_VULKAN) - find_package(Python COMPONENTS Interpreter) - if (NOT Python_FOUND) - message(FATAL_ERROR "python is required to compile Vulkan shaders") - endif() - find_package(Vulkan) if (Vulkan_FOUND) @@ -563,7 +558,9 @@ if (GGML_VULKAN) add_compile_definitions(GGML_VULKAN_RUN_TESTS) endif() - set (_llama_vk_genshaders_cmd ${Python_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/../ggml_vk_generate_shaders.py) + add_subdirectory(vulkan-gen-shaders) + + set (_llama_vk_genshaders_cmd vulkan-gen-shaders) set (_llama_vk_header ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.hpp) set (_llama_vk_source ${CMAKE_CURRENT_BINARY_DIR}/ggml-vulkan-shaders.cpp) set (_llama_vk_input_dir ${CMAKE_CURRENT_SOURCE_DIR}/vulkan-shaders) diff --git a/ggml/src/vulkan-gen-shaders/CMakeLists.txt b/ggml/src/vulkan-gen-shaders/CMakeLists.txt new file mode 100644 index 000000000..6c7d1f92a --- /dev/null +++ b/ggml/src/vulkan-gen-shaders/CMakeLists.txt @@ -0,0 +1,5 @@ + +set(TARGET vulkan-gen-shaders) +add_executable(${TARGET} vulkan-gen-shaders.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_compile_features(${TARGET} PRIVATE cxx_std_11) diff --git a/ggml/src/vulkan-gen-shaders/vulkan-gen-shaders.cpp b/ggml/src/vulkan-gen-shaders/vulkan-gen-shaders.cpp new file mode 100644 index 000000000..c4eafffcd --- /dev/null +++ b/ggml/src/vulkan-gen-shaders/vulkan-gen-shaders.cpp @@ -0,0 +1,525 @@ + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef _WIN32 + #include + #include // For _mkdir on Windows +#else + #include + #include + #include +#endif + +#define ASYNCIO_CONCURRENCY 64 + +std::mutex lock; +std::vector> shader_fnames; + +std::string GLSLC = "glslc"; +std::string input_dir = "vulkan-shaders"; +std::string output_dir = "/tmp"; +std::string target_hpp = "ggml-vulkan-shaders.hpp"; +std::string target_cpp = "ggml-vulkan-shaders.cpp"; +bool no_clean = false; + +const std::vector type_names = { + "f32", + "f16", + "q4_0", + "q4_1", + "q5_0", + "q5_1", + "q8_0", + "q2_k", + "q3_k", + "q4_k", + "q5_k", + "q6_k" +}; + +void execute_command(const std::string& command, std::string& stdout_str, std::string& stderr_str) { +#ifdef _WIN32 + HANDLE stdout_read, stdout_write; + HANDLE stderr_read, stderr_write; + SECURITY_ATTRIBUTES sa = { sizeof(SECURITY_ATTRIBUTES), NULL, TRUE }; + + if (!CreatePipe(&stdout_read, &stdout_write, &sa, 0) || + !SetHandleInformation(stdout_read, HANDLE_FLAG_INHERIT, 0)) { + throw std::runtime_error("Failed to create stdout pipe"); + } + + if (!CreatePipe(&stderr_read, &stderr_write, &sa, 0) || + !SetHandleInformation(stderr_read, HANDLE_FLAG_INHERIT, 0)) { + throw std::runtime_error("Failed to create stderr pipe"); + } + + PROCESS_INFORMATION pi; + STARTUPINFOA si = { sizeof(STARTUPINFOA) }; + si.dwFlags = STARTF_USESTDHANDLES; + si.hStdOutput = stdout_write; + si.hStdError = stderr_write; + + std::vector cmd(command.begin(), command.end()); + cmd.push_back('\0'); + + if (!CreateProcessA(NULL, cmd.data(), NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) { + throw std::runtime_error("Failed to create process"); + } + + CloseHandle(stdout_write); + CloseHandle(stderr_write); + + std::array buffer; + DWORD bytes_read; + + while (ReadFile(stdout_read, buffer.data(), buffer.size(), &bytes_read, NULL) && bytes_read > 0) { + stdout_str.append(buffer.data(), bytes_read); + } + + while (ReadFile(stderr_read, buffer.data(), buffer.size(), &bytes_read, NULL) && bytes_read > 0) { + stderr_str.append(buffer.data(), bytes_read); + } + + CloseHandle(stdout_read); + CloseHandle(stderr_read); + WaitForSingleObject(pi.hProcess, INFINITE); + CloseHandle(pi.hProcess); + CloseHandle(pi.hThread); +#else +int stdout_pipe[2]; + int stderr_pipe[2]; + + if (pipe(stdout_pipe) != 0 || pipe(stderr_pipe) != 0) { + throw std::runtime_error("Failed to create pipes"); + } + + pid_t pid = fork(); + if (pid < 0) { + throw std::runtime_error("Failed to fork process"); + } + + if (pid == 0) { + close(stdout_pipe[0]); + close(stderr_pipe[0]); + dup2(stdout_pipe[1], STDOUT_FILENO); + dup2(stderr_pipe[1], STDERR_FILENO); + close(stdout_pipe[1]); + close(stderr_pipe[1]); + execl("/bin/sh", "sh", "-c", command.c_str(), (char*) nullptr); + _exit(EXIT_FAILURE); + } else { + close(stdout_pipe[1]); + close(stderr_pipe[1]); + + std::array buffer; + ssize_t bytes_read; + + while ((bytes_read = read(stdout_pipe[0], buffer.data(), buffer.size())) > 0) { + stdout_str.append(buffer.data(), bytes_read); + } + + while ((bytes_read = read(stderr_pipe[0], buffer.data(), buffer.size())) > 0) { + stderr_str.append(buffer.data(), bytes_read); + } + + close(stdout_pipe[0]); + close(stderr_pipe[0]); + waitpid(pid, nullptr, 0); + } +#endif +} + +bool directory_exists(const std::string& path) { + struct stat info; + if (stat(path.c_str(), &info) != 0) { + return false; // Path doesn't exist or can't be accessed + } + return (info.st_mode & S_IFDIR) != 0; // Check if it is a directory +} + +bool create_directory(const std::string& path) { +#ifdef _WIN32 + return _mkdir(path.c_str()) == 0 || errno == EEXIST; // EEXIST means the directory already exists +#else + return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST; // 0755 is the directory permissions +#endif +} + +std::string to_uppercase(const std::string& input) { + std::string result = input; + for (char& c : result) { + c = std::toupper(c); + } + return result; +} + +bool string_ends_with(const std::string& str, const std::string& suffix) { + if (suffix.size() > str.size()) { + return false; + } + return std::equal(suffix.rbegin(), suffix.rend(), str.rbegin()); +} + +#ifdef _WIN32 + static const char path_separator = '\\'; +#else + static const char path_separator = '/'; +#endif + +std::string join_paths(const std::string& path1, const std::string& path2) { + return path1 + path_separator + path2; +} + +std::string basename(const std::string &path) { + return path.substr(path.find_last_of("/\\") + 1); +} + +void string_to_spv(const std::string& _name, const std::string& in_fname, const std::map& defines, bool fp16 = true) { + std::string name = _name + (fp16 ? "" : "_fp32"); + std::string out_fname = join_paths(output_dir, name + ".spv"); + std::string in_path = join_paths(input_dir, in_fname); + + std::vector cmd = {GLSLC, "-fshader-stage=compute", "--target-env=vulkan1.2", "-O", in_path, "-o", out_fname}; + for (const auto& define : defines) { + cmd.push_back("-D" + define.first + "=" + define.second); + } + + std::string command; + for (const auto& part : cmd) { + command += part + " "; + } + + std::string stdout_str, stderr_str; + try { + // std::cout << "Executing command: "; + // for (const auto& part : cmd) { + // std::cout << part << " "; + // } + // std::cout << std::endl; + + execute_command(command, stdout_str, stderr_str); + if (!stderr_str.empty()) { + std::cerr << "cannot compile " << name << "\n\n" << command << "\n\n" << stderr_str << std::endl; + return; + } + + std::lock_guard guard(lock); + shader_fnames.push_back(std::make_pair(name, out_fname)); + } catch (const std::exception& e) { + std::cerr << "Error executing command for " << name << ": " << e.what() << std::endl; + } +} + +std::map merge_maps(const std::map& a, const std::map& b) { + std::map result = a; + result.insert(b.begin(), b.end()); + return result; +} + +void matmul_shaders(std::vector>& tasks, bool fp16, bool matmul_id) { + std::string load_vec = fp16 ? "8" : "4"; + std::string aligned_b_type_f32 = fp16 ? "mat2x4" : "vec4"; + std::string aligned_b_type_f16 = fp16 ? "f16mat2x4" : "f16vec4"; + + std::map base_dict = {{"FLOAT_TYPE", fp16 ? "float16_t" : "float"}}; + std::string shader_name = "matmul"; + + if (matmul_id) { + base_dict["MUL_MAT_ID"] = "1"; + shader_name = "matmul_id"; + } + + if (fp16) { + base_dict["FLOAT16"] = "1"; + } + + // Shaders with f16 B_TYPE + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv(shader_name + "_f32_f16", "mul_mm.comp", merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}), fp16); + })); + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv(shader_name + "_f32_f16_aligned", "mul_mm.comp", 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"}}), fp16); + })); + + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv(shader_name + "_f16", "mul_mm.comp", merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}), fp16); + })); + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv(shader_name + "_f16_aligned", "mul_mm.comp", 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"}}), fp16); + })); + + for (const auto& tname : type_names) { + std::string data_a_key = "DATA_A_" + to_uppercase(tname); + std::string load_vec_a = (tname == "f32" || tname == "f16") ? load_vec : "2"; + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv(shader_name + "_" + tname + "_f32", "mul_mm.comp", merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}), fp16); + })); + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv(shader_name + "_" + tname + "_f32_aligned", "mul_mm.comp", merge_maps(base_dict, {{data_a_key, "2"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"D_TYPE", "float"}}), fp16); + })); + } +} + +void process_shaders(std::vector>& tasks) { + std::cout << "ggml_vulkan: Generating and compiling shaders to SPIR-V" << std::endl; + std::map base_dict = {{"FLOAT_TYPE", "float"}}; + + for (const auto& fp16 : {false, true}) { + matmul_shaders(tasks, fp16, false); + matmul_shaders(tasks, fp16, true); + } + + for (const auto& tname : type_names) { + // mul mat vec + std::string data_a_key = "DATA_A_" + to_uppercase(tname); + std::string shader = (string_ends_with(tname, "_k")) ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; + + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); + })); + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}})); + })); + + tasks.push_back(std::async(std::launch::async, [=] { + 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"}, {"D_TYPE", "float"}})); + })); + + // Dequant shaders + if (tname != "f16") { + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}})); + })); + } + + if (!string_ends_with(tname, "_k")) { + shader = (tname == "f32" || tname == "f16") ? "get_rows.comp" : "get_rows_quant.comp"; + + if (tname == "f16") { + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("get_rows_" + tname, shader, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + })); + } else { + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("get_rows_" + tname, shader, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}}); + })); + } + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("get_rows_" + tname + "_f32", shader, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float"}}); + })); + } + } + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + })); + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); + })); + + // Norms + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + })); + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("cpy_f32_f32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + })); + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); + })); + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("add_f32", "add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("mul_f32", "mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("div_f32", "div.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("scale_f32", "scale.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("sqr_f32", "square.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("clamp_f32", "clamp.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("gelu_f32", "gelu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + })); + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("silu_f32", "silu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + })); + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("relu_f32", "relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + })); + + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("soft_max_f32", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); + })); + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("soft_max_f32_f16", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}})); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("rope_norm_f32", "rope_norm.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + })); + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("rope_norm_f16", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("rope_neox_f32", "rope_neox.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + })); + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("rope_neox_f16", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}}); + })); + + tasks.push_back(std::async(std::launch::async, [] { + string_to_spv("argsort_f32", "argsort.comp", {{"A_TYPE", "float"}}); + })); + + tasks.push_back(std::async(std::launch::async, [=] { + string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + })); +} + +void write_output_files() { + FILE* hdr = fopen(target_hpp.c_str(), "w"); + FILE* src = fopen(target_cpp.c_str(), "w"); + + fprintf(hdr, "#include \n\n"); + fprintf(src, "#include \"%s\"\n\n", basename(target_hpp).c_str()); + + for (const auto& pair : shader_fnames) { + const std::string& name = pair.first; + const std::string& path = pair.second; + FILE* spv = fopen(path.c_str(), "rb"); + if (!spv) { + std::cerr << "Error opening SPIR-V file: " << path << "\n"; + continue; + } + + fseek(spv, 0, SEEK_END); + size_t size = ftell(spv); + fseek(spv, 0, SEEK_SET); + + std::vector data(size); + size_t read_size = fread(data.data(), 1, size, spv); + fclose(spv); + if (read_size != size) { + std::cerr << "Error reading SPIR-V file: " << path << "\n"; + continue; + } + + fprintf(hdr, "extern unsigned char %s_data[%zu];\n", name.c_str(), size); + fprintf(hdr, "const uint64_t %s_len = %zu;\n\n", name.c_str(), size); + + fprintf(src, "unsigned char %s_data[%zu] = {\n", name.c_str(), size); + for (size_t i = 0; i < size; ++i) { + fprintf(src, "0x%02x,", data[i]); + if ((i + 1) % 12 == 0) fprintf(src, "\n"); + } + fprintf(src, "\n};\n\n"); + + if (!no_clean) { + std::remove(path.c_str()); + } + } + + fclose(hdr); + fclose(src); +} + +int main(int argc, char** argv) { + std::map args; + for (int i = 1; i < argc; i += 2) { + if (i + 1 < argc) { + args[argv[i]] = argv[i + 1]; + } + } + + if (args.find("--glslc") != args.end()) { + GLSLC = args["--glslc"]; // Path to glslc + } + if (args.find("--input-dir") != args.end()) { + input_dir = args["--input-dir"]; // Directory containing shader sources + } + if (args.find("--output-dir") != args.end()) { + output_dir = args["--output-dir"]; // Directory for containing SPIR-V output + } + if (args.find("--target-hpp") != args.end()) { + target_hpp = args["--target-hpp"]; // Path to generated header file + } + if (args.find("--target-cpp") != args.end()) { + target_cpp = args["--target-cpp"]; // Path to generated cpp file + } + if (args.find("--no-clean") != args.end()) { + no_clean = true; // Keep temporary SPIR-V files in output-dir after build + } + + if (!directory_exists(input_dir)) { + std::cerr << "\"" << input_dir << "\" must be a valid directory containing shader sources" << std::endl; + return EXIT_FAILURE; + } + + if (!directory_exists(output_dir)) { + if (!create_directory(output_dir)) { + std::cerr << "Error creating output directory: " << output_dir << "\n"; + return EXIT_FAILURE; + } + } + + std::vector> tasks; + process_shaders(tasks); + + for (auto& task : tasks) { + task.get(); + } + + write_output_files(); + + return EXIT_SUCCESS; +} +