ggml: Add initial WebGPU backend (#14521)

* Minimal setup of webgpu backend with dawn. Just prints out the adapter and segfaults

* Initialize webgpu device

* Making progress on setting up the backend

* Finish more boilerplate/utility functions

* Organize file and work on alloc buffer

* Add webgpu_context to prepare for actually running some shaders

* Work on memset and add shader loading

* Work on memset polyfill

* Implement set_tensor as webgpu WriteBuffer, remove host_buffer stubs since webgpu doesn't support it

* Implement get_tensor and buffer_clear

* Finish rest of setup

* Start work on compute graph

* Basic mat mul working

* Work on emscripten build

* Basic WebGPU backend instructions

* Use EMSCRIPTEN flag

* Work on passing ci, implement 4d tensor multiplication

* Pass thread safety test

* Implement permuting for mul_mat and cpy

* minor cleanups

* Address feedback

* Remove division by type size in cpy op

* Fix formatting and add github action workflows for vulkan and metal (m-series) webgpu backends

* Fix name

* Fix macos dawn prefix path

ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl

@@ -0,0 +1,60 @@
+enable f16;
+
+@group(0) @binding(0)
+var<storage, read_write> src: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> dst: array<f16>;
+
+struct Params {
+    ne: u32,            // total number of elements
+    offset_src: u32,    // in elements
+    offset_dst: u32,    // in elements
+
+    // Strides (in elements) — may be permuted
+    stride_src0: u32,
+    stride_src1: u32,
+    stride_src2: u32,
+    stride_src3: u32,
+
+    stride_dst0: u32,
+    stride_dst1: u32,
+    stride_dst2: u32,
+    stride_dst3: u32,
+
+    // Logical shape (same for both tensors)
+    ne0: u32,
+    ne1: u32,
+    ne2: u32,
+    ne3: u32,
+};
+
+@group(0) @binding(2)
+var<uniform> params: Params;
+
+override wg_size: u32;
+@compute @workgroup_size(wg_size)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    if (gid.x >= params.ne) {
+        return;
+    }
+
+    var i = gid.x;
+
+    let i3 = i / (params.ne2 * params.ne1 * params.ne0);
+    i = i % (params.ne2 * params.ne1 * params.ne0);
+
+    let i2 = i / (params.ne1 * params.ne0);
+    i = i % (params.ne1 * params.ne0);
+
+    let i1 = i / params.ne0;
+    let i0 = i % params.ne0;
+
+    let src_idx = i0 * params.stride_src0 + i1 * params.stride_src1 +
+                  i2 * params.stride_src2 + i3 * params.stride_src3;
+
+    let dst_idx = i0 * params.stride_dst0 + i1 * params.stride_dst1 +
+                  i2 * params.stride_dst2 + i3 * params.stride_dst3;
+
+    dst[params.offset_dst + dst_idx] = f16(src[params.offset_src + src_idx]);
+}

ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py

@@ -0,0 +1,35 @@
+import os
+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'
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input', required=True)
+    parser.add_argument('--output', required=True)
+    args = parser.parse_args()
+
+    with open(args.output, 'w', encoding='utf-8') as out:
+        out.write("// Auto-generated shader embedding \n\n")
+        for fname in sorted(os.listdir(args.input)):
+            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')
+
+
+if __name__ == '__main__':
+    main()

ggml/src/ggml-webgpu/wgsl-shaders/memset.wgsl

@@ -0,0 +1,40 @@
+@group(0) @binding(0)
+var<storage, read_write> output_buffer: array<u32>;
+
+struct Params {
+    offset: u32, // in bytes
+    size: u32,   // in bytes
+    value: u32,  // 4 8-bit values, which are either repeating (memset_tensor) or may be separate (cleaning up unaligned set_tensor operations)
+};
+
+@group(0) @binding(1)
+var<uniform> params: Params;
+
+override wg_size: u32;
+override bytes_per_thread: u32;
+
+@compute @workgroup_size(wg_size)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x * bytes_per_thread;
+    let start = params.offset;
+    let end = params.offset + params.size;
+
+    for (var j: u32 = 0u; j < bytes_per_thread; j = j + 1u) {
+        let byte_index = start + i + j;
+        if (byte_index + 4u <= end) {
+            output_buffer[(byte_index >> 2u)] = params.value;
+        } else {
+            // Handle tail (unaligned)
+            for (var k: u32 = 0u; k < 4u; k = k + 1u) {
+                let idx = byte_index + k;
+                if (idx < end) {
+                    let word_idx = idx >> 2u;
+                    let byte_offset = (idx & 3u) * 8u;
+                    let mask = ~(0xffu << byte_offset);
+                    let existing = output_buffer[word_idx];
+                    output_buffer[word_idx] = (existing & mask) | ((params.value & 0xffu) << byte_offset);
+                }
+            }
+        }
+    }
+}

ggml/src/ggml-webgpu/wgsl-shaders/mul_mat.wgsl

@@ -0,0 +1,56 @@
+struct MulMatParams {
+    m: u32,
+    n: u32,
+    k: u32,
+    // all strides are in elements
+    stride_01: u32,
+    stride_11: u32,
+    stride_02: u32,
+    stride_12: u32,
+    stride_03: u32,
+    stride_13: u32,
+
+    bs02: u32,
+    bs03: u32,
+    broadcast2: u32,
+    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(2) var<storage, read_write> dst: array<f32>; // M rows, N columns
+
+@group(0) @binding(3) var<uniform> params: MulMatParams;
+
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let total = params.m * params.n * params.bs02 * params.broadcast2 * params.bs03 * params.broadcast3;
+    if (global_id.x >= total) {
+        return;
+    }
+
+    let dst2_stride = params.m * params.n;
+    let dst3_stride = dst2_stride * params.bs02 * params.broadcast2;
+
+    let dst3_idx = global_id.x / dst3_stride;
+    let src03_idx = dst3_idx / params.broadcast3; // src0 may be broadcast along the third dimension
+    let src13_idx = dst3_idx; // src1 is not broadcast
+    let dst3_rem = global_id.x % dst3_stride;
+
+    let dst2_idx = dst3_rem / dst2_stride;
+    let src02_idx = dst2_idx / params.broadcast2; // src0 may also be broadcast along the second dimension
+    let src12_idx = dst2_idx; // src1 is not broadcast
+
+    let dst2_rem = dst3_rem % dst2_stride;
+
+    let row = dst2_rem / params.n; // output row
+    let col = dst2_rem % params.n; // output column
+
+    var sum = 0.0;
+    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];
+    }
+    dst[dst3_idx * dst3_stride + dst2_idx * dst2_stride + row * params.n + col] = sum;
+}