Work on templating for different types in shaders

reeselevine · reeselevine · commit 6355137c76a1 · 2025-08-07T13:48:09.000-07:00
diff --git a/ggml/src/ggml-webgpu/CMakeLists.txt b/ggml/src/ggml-webgpu/CMakeLists.txt
@@ -20,8 +20,8 @@ add_custom_command(
     COMMAND ${CMAKE_COMMAND} -E make_directory ${SHADER_OUTPUT_DIR}
     COMMAND ${CMAKE_COMMAND} -E env PYTHONIOENCODING=utf-8
         ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/wgsl-shaders/embed_wgsl.py
-            --input "${SHADER_DIR}"
-            --output "${SHADER_HEADER}"
+            --input_dir "${SHADER_DIR}"
+            --output_file "${SHADER_HEADER}"
     DEPENDS ${WGSL_SHADER_FILES} ${CMAKE_CURRENT_SOURCE_DIR}/wgsl-shaders/embed_wgsl.py
     VERBATIM
 )
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -50,6 +50,13 @@ static uint64_t webgpu_tensor_offset(const ggml_tensor * tensor) {
 
 /* Struct definitions */
 
+struct webgpu_pipeline_info {
+    std::string  name;
+    const char * shader_code;
+    ggml_type    src0_type;
+    ggml_type    src1_type;
+};
+
 // Forward reference
 static void ggml_webgpu_create_buffer(wgpu::Device &    device,
                                       wgpu::Buffer &    buffer,
@@ -124,7 +131,8 @@ struct webgpu_context_struct {
     webgpu_buf_pool set_rows_error_buf_pool;
 
     wgpu::ComputePipeline memset_pipeline;
-    wgpu::ComputePipeline mul_mat_pipeline;
+    // [src0 0=fp32,1=fp16][src1 0=fp32,1=fp16]
+    wgpu::ComputePipeline mul_mat_pipeline[2][2];
     wgpu::ComputePipeline set_rows_pipeline;
     wgpu::ComputePipeline cpy_pipeline;
 
@@ -227,6 +235,15 @@ static void ggml_webgpu_create_buffer(wgpu::Device &    device,
 
 /** End WebGPU object initializations */
 
+/** Utility Functions */
+
+size_t ggml_webgpu_binding_size(ggml_tensor * t, size_t misalignment) {
+    return (ggml_nbytes(t) + misalignment + WEBGPU_STORAGE_BUF_BINDING_MULT - 1) &
+           ~(WEBGPU_STORAGE_BUF_BINDING_MULT - 1);
+}
+
+/** End Utility Functions */
+
 /** WebGPU Actions */
 
 // Wait for the queue to finish processing all submitted work
@@ -479,13 +496,11 @@ static void ggml_webgpu_cpy(webgpu_context & ctx, ggml_tensor * src, ggml_tensor
         { .binding = 0,
          .buffer  = ggml_backend_webgpu_tensor_buf(src),
          .offset  = src_offset,
-         .size    = (ggml_nbytes(src) + src_misalignment + WEBGPU_STORAGE_BUF_BINDING_MULT - 1) &
-                  ~(WEBGPU_STORAGE_BUF_BINDING_MULT - 1) },
+         .size    = ggml_webgpu_binding_size(src, src_misalignment) },
         { .binding = 1,
          .buffer  = ggml_backend_webgpu_tensor_buf(dst),
          .offset  = dst_offset,
-         .size    = (ggml_nbytes(dst) + dst_misalignment + WEBGPU_STORAGE_BUF_BINDING_MULT - 1) &
-                  ~(WEBGPU_STORAGE_BUF_BINDING_MULT - 1) }
+         .size   = ggml_webgpu_binding_size(dst, dst_misalignment) }
     };
 
     size_t   max_wg_size = ctx->limits.maxComputeWorkgroupSizeX;
@@ -542,15 +557,15 @@ static void ggml_webgpu_set_rows(webgpu_context & ctx, ggml_tensor * src, ggml_t
         { .binding = 0,
          .buffer  = ggml_backend_webgpu_tensor_buf(src),
          .offset  = ggml_backend_webgpu_tensor_offset(src),
-         .size    = ggml_nbytes(src)                                                                       },
+         .size    = ggml_webgpu_binding_size(src, src_misalignment) },
         { .binding = 1,
          .buffer  = ggml_backend_webgpu_tensor_buf(idx),
          .offset  = ggml_backend_webgpu_tensor_offset(idx),
-         .size    = ggml_nbytes(idx)                                                                       },
+         .size    = ggml_webgpu_binding_size(idx, idx_misalignment) },
         { .binding = 2,
          .buffer  = ggml_backend_webgpu_tensor_buf(dst),
          .offset  = ggml_backend_webgpu_tensor_offset(dst),
-         .size    = ggml_nbytes(dst)                                                                       },
+         .size    = ggml_webgpu_binding_size(dst, dst_misalignment) },
         { .binding = 3, .buffer = error_bufs.dev_buf,    .offset = 0, .size = error_bufs.dev_buf.GetSize() }
     };
 
@@ -564,7 +579,21 @@ static void ggml_webgpu_set_rows(webgpu_context & ctx, ggml_tensor * src, ggml_t
 }
 
 static void ggml_webgpu_mul_mat(webgpu_context & ctx, ggml_tensor * src0, ggml_tensor * src1, ggml_tensor * dst) {
+    size_t src0_offset       = ggml_backend_webgpu_tensor_offset(src0);
+    size_t src0_misalignment = src0_offset & (ctx->limits.minStorageBufferOffsetAlignment - 1);
+    // align to minimum offset alignment
+    src0_offset &= ~(ctx->limits.minStorageBufferOffsetAlignment - 1);
+    size_t src1_offset       = ggml_backend_webgpu_tensor_offset(src1);
+    size_t src1_misalignment = src1_offset & (ctx->limits.minStorageBufferOffsetAlignment - 1);
+    src1_offset &= ~(ctx->limits.minStorageBufferOffsetAlignment - 1);
+    size_t dst_offset       = ggml_backend_webgpu_tensor_offset(dst);
+    size_t dst_misalignment = dst_offset & (ctx->limits.minStorageBufferOffsetAlignment - 1);
+    dst_offset &= ~(ctx->limits.minStorageBufferOffsetAlignment - 1);
+
     std::vector<uint32_t> params = {
+        (uint32_t) (src0_misalignment / ggml_type_size(src0->type)),
+        (uint32_t) (src1_misalignment / ggml_type_size(src1->type)),
+        (uint32_t) (dst_misalignment / ggml_type_size(dst->type)),
         (uint32_t) dst->ne[1],                                  // number of rows in result (M)
         (uint32_t) dst->ne[0],                                  // number of columns in result (N)
         (uint32_t) src0->ne[0],                                 // number of columns in src0/src1 (K)
@@ -584,20 +613,20 @@ static void ggml_webgpu_mul_mat(webgpu_context & ctx, ggml_tensor * src0, ggml_t
         { .binding = 0,
          .buffer  = ggml_backend_webgpu_tensor_buf(src0),
          .offset  = ggml_backend_webgpu_tensor_offset(src0),
-         .size    = ggml_nbytes(src0) },
+         .size    = ggml_webgpu_binding_size(src0, src0_misalignment) },
         { .binding = 1,
          .buffer  = ggml_backend_webgpu_tensor_buf(src1),
          .offset  = ggml_backend_webgpu_tensor_offset(src1),
-         .size    = ggml_nbytes(src1) },
+         .size    = ggml_webgpu_binding_size(src1, src1_misalignment) },
         { .binding = 2,
          .buffer  = ggml_backend_webgpu_tensor_buf(dst),
          .offset  = ggml_backend_webgpu_tensor_offset(dst),
-         .size    = ggml_nbytes(dst)  }
+         .size    = ggml_webgpu_binding_size(dst, dst_misalignment) }
     };
 
     uint32_t wg_x =
         (dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3] + WEBGPU_MUL_MAT_WG_SIZE - 1) / WEBGPU_MUL_MAT_WG_SIZE;
-    ggml_backend_webgpu_build_and_enqueue(ctx, ctx->mul_mat_pipeline, params, entries, wg_x);
+    ggml_backend_webgpu_build_and_enqueue(ctx, ctx->mul_mat_pipeline[src0->type][src1->type], params, entries, wg_x);
 }
 
 // Returns true if node has enqueued work into the queue, false otherwise
@@ -907,7 +936,31 @@ static void ggml_webgpu_init_memset_pipeline(webgpu_context & webgpu_ctx) {
 }
 
 static void ggml_webgpu_init_mul_mat_pipeline(webgpu_context & webgpu_ctx) {
-    ggml_webgpu_create_pipeline(webgpu_ctx->device, webgpu_ctx->mul_mat_pipeline, wgsl_mul_mat, "mul_mat");
+    webgpu_pipeline_info pipeline_infos[4] = {
+        { .name        = "mul_mat_f32_f32",
+         .shader_code = wgsl_mul_mat_f32_f32,
+         .src0_type   = GGML_TYPE_F32,
+         .src1_type   = GGML_TYPE_F32 },
+        { .name        = "mul_mat_f16_f16",
+         .shader_code = wgsl_mul_mat_f16_f16,
+         .src0_type   = GGML_TYPE_F16,
+         .src1_type   = GGML_TYPE_F16 },
+        { .name        = "mul_mat_f32_f16",
+         .shader_code = wgsl_mul_mat_f32_f16,
+         .src0_type   = GGML_TYPE_F32,
+         .src1_type   = GGML_TYPE_F16 },
+        { .name        = "mul_mat_f16_f32",
+         .shader_code = wgsl_mul_mat_f16_f32,
+         .src0_type   = GGML_TYPE_F16,
+         .src1_type   = GGML_TYPE_F32 }
+    };
+
+    for (auto & pipeline_info : pipeline_infos) {
+        ggml_webgpu_create_pipeline(webgpu_ctx->device,
+                                    webgpu_ctx->mul_mat_pipeline[pipeline_info.src0_type][pipeline_info.src1_type],
+                                    pipeline_info.shader_code,
+                                    pipeline_info.name.data());
+    }
 }
 
 static void ggml_webgpu_init_set_rows_pipeline(webgpu_context & webgpu_ctx) {
@@ -1056,7 +1109,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
         case GGML_OP_CPY | GGML_OP_SET_ROWS:
             return op->type == GGML_TYPE_F16 && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_MUL_MAT:
-            return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
+            return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) &&
+                   (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16);
         default:
             return false;
     }
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py b/ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py
@@ -1,35 +1,60 @@
 import os
+import re
+import ast
 import argparse
 
-
-def escape_triple_quotes(wgsl):
-    # Simple defense in case of embedded """
-    return wgsl.replace('"""', '\\"""')
-
-
-def to_cpp_string_literal(varname, content):
-    return f'const char* wgsl_{varname} = R"({content})";\n'
-
+variants_regex = re.compile(r'//\s*Variants:\s*\n(\[.*?\])', re.DOTALL)
+
+def remove_variants_block(template_text):
+    return re.sub(variants_regex, '', template_text)
+
+def extract_variants(template_text):
+    match = re.search(variants_regex, template_text)
+    if not match:
+        return None
+    return ast.literal_eval(match.group(1))
+
+def write_shader(shader_name, shader_code, output_dir, outfile):
+    if output_dir:
+        wgsl_filename = os.path.join(output_dir, f"{shader_name}.wgsl")
+        with open(wgsl_filename, 'w', encoding='utf-8') as f_out:
+            f_out.write(shader_code)
+    outfile.write(f'const char* wgsl_{shader_name} = R"({shader_code})";\n')
+    outfile.write('\n')
+
+def generate_variants(shader_path, output_dir, outfile):
+    shader_base_name = shader_path.split("/")[-1].split(".")[0]
+    with open(shader_path, 'r', encoding='utf-8') as f:
+        shader_code = f.read()
+    variants = extract_variants(shader_code)
+    shader_code = remove_variants_block(shader_code)
+    if not variants:
+        write_shader(shader_base_name, shader_code, output_dir, outfile)
+    else:
+        for variant in variants:
+            shader_variant = shader_code
+            parts = []
+            for key, val in variant.items():
+                parts.append(val)
+                shader_variant = shader_variant.replace(key, val)
+            output_name = f"{shader_base_name}_" + "_".join(parts)
+            write_shader(output_name, shader_variant, output_dir, outfile)
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--input', required=True)
-    parser.add_argument('--output', required=True)
+    parser.add_argument('--input_dir', required=True)
+    parser.add_argument('--output_file', required=True)
+    parser.add_argument('--output_dir')
     args = parser.parse_args()
-
-    with open(args.output, 'w', encoding='utf-8') as out:
+    if args.output_dir:
+        os.makedirs(args.output_dir, exist_ok=True)
+    with open(args.output_file, 'w', encoding='utf-8') as out:
         out.write("// Auto-generated shader embedding \n\n")
-        for fname in sorted(os.listdir(args.input)):
+        for fname in sorted(os.listdir(args.input_dir)):
             if not fname.endswith('.wgsl'):
                 continue
-            shader_path = os.path.join(args.input, fname)
-            varname = os.path.splitext(fname)[0]
-            with open(shader_path, 'r', encoding='utf-8') as f:
-                content = f.read()
-            content = escape_triple_quotes(content)
-            out.write(to_cpp_string_literal(varname, content))
-            out.write('\n')
-
+            shader_path = os.path.join(args.input_dir, fname)
+            generate_variants(shader_path, args.output_dir, out)
 
 if __name__ == '__main__':
     main()
diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat.tmpl.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/mul_mat.tmpl.wgsl
@@ -1,4 +1,30 @@
+// Variants:
+[
+  {
+    "SRC0_TYPE" : "f32",
+    "SRC1_TYPE" : "f32"
+  },
+  {
+    "SRC0_TYPE" : "f16",
+    "SRC1_TYPE" : "f16"
+  },
+  {
+    "SRC0_TYPE" : "f16",
+    "SRC1_TYPE" : "f32"
+  },
+  {
+    "SRC0_TYPE" : "f32",
+    "SRC1_TYPE" : "f16"
+  }
+]
+
+// Shader Template:
+enable f16;
+
 struct MulMatParams {
+    offset_src0: u32, // in elements
+    offset_src1: u32, // in elements
+    offset_dst: u32, // in elements
     m: u32,
     n: u32,
     k: u32,
@@ -16,8 +42,8 @@ struct MulMatParams {
     broadcast3: u32
 };
 
-@group(0) @binding(0) var<storage, read_write> src0: array<f32>; // N rows, K columns
-@group(0) @binding(1) var<storage, read_write> src1: array<f32>; // M rows, K columns (transposed)
+@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>; // N rows, K columns
+@group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>; // M rows, K columns (transposed)
 @group(0) @binding(2) var<storage, read_write> dst: array<f32>; // M rows, N columns
 
 @group(0) @binding(3) var<uniform> params: MulMatParams;
@@ -50,7 +76,8 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
     for (var i: u32 = 0u; i < params.k; i = i + 1u) {
         let src0_idx = src03_idx * params.stride_03 + src02_idx * params.stride_02 + col * params.stride_01 + i;
         let src1_idx = src13_idx * params.stride_13 + src12_idx * params.stride_12 + row * params.stride_11 + i;
-        sum = sum + src0[src0_idx] * src1[src1_idx];
+        sum = sum + f32(src0[params.offset_src0 + src0_idx]) * f32(src1[params.offset_src1 + src1_idx]);
     }
-    dst[dst3_idx * dst3_stride + dst2_idx * dst2_stride + row * params.n + col] = sum;
+    dst[params.offset_dst + dst3_idx * dst3_stride + dst2_idx * dst2_stride + row * params.n + col] = sum;
+
 }
