char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Abhijit Rameshggml-webgpu: port all AOT operators to JIT (#20728)825eb91a6

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.3 KiB89 linesraw
1import os
2import re
3import argparse
4
5
6def expand_includes(shader, input_dir):
7    """
8    Replace #include "file" lines in the text with the contents of that file.
9    Searches for files relative to input_dir.
10    """
11    include_pattern = re.compile(r'^\s*#include\s+"([^"]+)"\s*$', re.MULTILINE)
12
13    def replacer(match):
14        fname = match.group(1)
15        file_path = os.path.join(input_dir, fname)
16        if not os.path.exists(file_path):
17            raise FileNotFoundError(f"Included file not found: {file_path}")
18        with open(file_path, "r", encoding="utf-8") as f:
19            included_code = f.read()
20        # Recursively expand includes inside the included file
21        return expand_includes(included_code, input_dir)
22
23    return include_pattern.sub(replacer, shader)
24
25
26def chunk_shader(shader_code, max_chunk_len=60000):
27    """Split shader_code into safe raw-string sized chunks."""
28    return [shader_code[i : i + max_chunk_len] for i in range(0, len(shader_code), max_chunk_len)]
29
30
31def raw_delim(shader_code):
32    """Pick a raw-string delimiter that does not appear in the shader."""
33    delim = "wgsl"
34    while f"){delim}\"" in shader_code:
35        delim += "_x"
36    return delim
37
38
39def write_shader(shader_name, shader_code, output_dir, outfile, input_dir):
40    shader_code = expand_includes(shader_code, input_dir)
41
42    if output_dir:
43        wgsl_filename = os.path.join(output_dir, f"{shader_name}.wgsl")
44        with open(wgsl_filename, "w", encoding="utf-8") as f_out:
45            f_out.write(shader_code)
46
47    delim = raw_delim(shader_code)
48    chunks = chunk_shader(shader_code)
49
50    if len(chunks) == 1:
51        outfile.write(f'const char* wgsl_{shader_name} = R"{delim}({shader_code}){delim}";\n\n')
52    else:
53        for idx, chunk in enumerate(chunks):
54            outfile.write(f'static const char wgsl_{shader_name}_part{idx}[] = R"{delim}({chunk}){delim}";\n\n')
55        outfile.write(f'static const std::string& wgsl_{shader_name}_str() {{\n')
56        outfile.write('    static const std::string s = []{\n')
57        outfile.write('        std::string tmp;\n')
58        outfile.write(f'        tmp.reserve({len(shader_code)});\n')
59        for idx in range(len(chunks)):
60            outfile.write(f'        tmp.append(wgsl_{shader_name}_part{idx});\n')
61        outfile.write('        return tmp;\n')
62        outfile.write('    }();\n')
63        outfile.write('    return s;\n')
64        outfile.write('}\n')
65        outfile.write(f'const char* wgsl_{shader_name} = wgsl_{shader_name}_str().c_str();\n\n')
66
67
68def main():
69    parser = argparse.ArgumentParser()
70    parser.add_argument("--input_dir", required=True)
71    parser.add_argument("--output_file", required=True)
72    args = parser.parse_args()
73
74    with open(args.output_file, "w", encoding="utf-8") as out:
75        out.write("// Auto-generated shader embedding\n")
76        out.write("#include <string>\n\n")
77        for fname in sorted(os.listdir(args.input_dir)):
78            if fname.endswith(".wgsl"):
79                shader_path = os.path.join(args.input_dir, fname)
80                shader_name = fname.replace(".wgsl", "")
81
82                with open(shader_path, "r", encoding="utf-8") as f:
83                    shader_code = f.read()
84
85                write_shader(shader_name, shader_code, None, out, args.input_dir)
86
87
88if __name__ == "__main__":
89    main()