Implement permuting for mul_mat and cpy

Author: reeselevine

ggml/src/ggml-webgpu/ggml-webgpu.cpp

@@ -22,8 +22,8 @@
 #define WEBGPU_MAX_BUFFERS 32
 
 #define WEBGPU_MUL_MAT_WG_SIZE 64
-#define WEBGPU_MUL_MAT_PARAMS_SIZE (7 * sizeof(uint32_t)) // M, N, K, batch sizes, broadcasts
-#define WEBGPU_CPY_PARAMS_SIZE (3 * sizeof(uint32_t)) // number of elements to copy, alignments
+#define WEBGPU_MUL_MAT_PARAMS_SIZE (13 * sizeof(uint32_t)) // M, N, K, batch sizes, broadcasts
+#define WEBGPU_CPY_PARAMS_SIZE (15 * sizeof(uint32_t)) // strides and offsets
 #define WEBGPU_STORAGE_BUF_BINDING_MULT 4 // a storage buffer binding size must be a multiple of 4
 
 /* End Constants */
@@ -266,10 +266,26 @@ static bool ggml_webgpu_encode_node(webgpu_context ctx, ggml_tensor * node){
             ggml_backend_webgpu_map_buffer(ctx, ctx->cpy_params_host_buf, 
                 wgpu::MapMode::Write, 0, ctx->cpy_params_host_buf.GetSize());
             uint32_t * params = (uint32_t *) ctx->cpy_params_host_buf.GetMappedRange();
-            uint32_t ne = (uint32_t)ggml_nelements(node); // number of elements to copy
+            uint32_t ne = (uint32_t)ggml_nelements(node);
             params[0] = ne;
             params[1] = src_misalignment;
             params[2] = dst_misalignment;
+
+            // Convert byte-strides to element-strides
+            params[3] = (uint32_t)src->nb[0]/ggml_type_size(src->type);
+            params[4] = (uint32_t)src->nb[1]/ggml_type_size(src->type);
+            params[5] = (uint32_t)src->nb[2]/ggml_type_size(src->type);
+            params[6] = (uint32_t)src->nb[3]/ggml_type_size(src->type);
+            params[7] = (uint32_t)node->nb[0]/ggml_type_size(node->type);
+            params[8] = (uint32_t)node->nb[1]/ggml_type_size(node->type);
+            params[9] = (uint32_t)node->nb[2]/ggml_type_size(node->type);
+            params[10] = (uint32_t)node->nb[3]/ggml_type_size(node->type);
+            // Logical shape — same for both tensors even if permuted
+            params[11] = (uint32_t)(src->ne[0]);
+            params[12] = (uint32_t)(src->ne[1]);
+            params[13] = (uint32_t)(src->ne[2]);
+            params[14] = (uint32_t)(src->ne[3]);
+
             ctx->cpy_params_host_buf.Unmap();
 
             wgpu::BindGroupEntry entries[3];
@@ -338,10 +354,18 @@ static bool ggml_webgpu_encode_node(webgpu_context ctx, ggml_tensor * node){
             params[0] = (uint32_t)node->ne[1]; // number of rows in result (M)
             params[1] = (uint32_t)node->ne[0]; // number of columns in result (N)
             params[2] = (uint32_t)src0->ne[0]; // number of columns in src0/src1 (K)
-            params[3] = (uint32_t)src0->ne[2]; // batch size in dimension 2
-            params[4] = (uint32_t)src0->ne[3]; // batch size in dimension 3
-            params[5] = (uint32_t)(src1->ne[2]/src0->ne[2]); // broadcast in dimension 2
-            params[6] = (uint32_t)(src1->ne[3]/src0->ne[3]); // broadcast in dimension 3
+
+            params[3] = (uint32_t)src0->nb[1]/ggml_type_size(src0->type); // stride (elements) of src0 in dimension 1
+            params[4] = (uint32_t)src1->nb[1]/ggml_type_size(src1->type); // stride (elements) of src1 in dimension 1
+            params[5] = (uint32_t)src0->nb[2]/ggml_type_size(src0->type); // stride (elements) of src0 in dimension 2
+            params[6] = (uint32_t)src1->nb[2]/ggml_type_size(src1->type); // stride (elements) of src1 in dimension 2
+            params[7] = (uint32_t)src0->nb[3]/ggml_type_size(src0->type); // stride (elements) of src0 in dimension 3
+            params[8] = (uint32_t)src1->nb[3]/ggml_type_size(src1->type); // stride (elements) of src1 in dimension 3
+
+            params[9] = (uint32_t)src0->ne[2]; // batch size in dimension 2
+            params[10] = (uint32_t)src0->ne[3]; // batch size in dimension 3
+            params[11] = (uint32_t)(src1->ne[2]/src0->ne[2]); // broadcast in dimension 2
+            params[12] = (uint32_t)(src1->ne[3]/src0->ne[3]); // broadcast in dimension 3
 
             ctx->mul_mat_params_host_buf.Unmap();
 

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

@@ -7,26 +7,54 @@ var<storage, read_write> src: array<f32>;
 var<storage, read_write> dst: array<f16>;
 
 struct Params {
-    ne: u32, // number of elements
-    src_offset: u32, // src offset in bytes
-    dst_offset: u32 // dst offset in bytes
+    ne: u32,            // total number of elements
+    src_offset: u32,    // in bytes
+    dst_offset: u32,    // in bytes
+
+    // 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;
-const elems_per_thread: u32 = 4;
-
 @compute @workgroup_size(wg_size)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x * elems_per_thread;
-    // chunked loop
-    for (var j: u32 = 0u; j < elems_per_thread; j = j + 1u) {
-        let i = idx + j;
-        if (i < params.ne) {
-            // Convert f32 to f16
-            dst[dst_offset/2 + i] = f16(src[src_offset/4 + i]);
-        }
+    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.dst_offset / 2 + dst_idx] = f16(src[params.src_offset / 4 + src_idx]);
 }

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

@@ -2,6 +2,13 @@ struct MulMatParams {
     m: u32,
     n: u32,
     k: u32,
+    stride_01: u32,
+    stride_11: u32,
+    stride_02: u32,
+    stride_12: u32,
+    stride_03: u32,
+    stride_13: u32,
+
     bs02: u32,
     bs03: u32,
     broadcast2: u32,
@@ -21,12 +28,6 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
         return;
     }
 
-    let src02_stride = params.n * params.k;
-    let src03_stride = src02_stride * params.bs02;
-
-    let src12_stride = params.m * params.k;
-    let src13_stride = src12_stride * params.bs02 * params.broadcast2;
-
     let dst2_stride = params.m * params.n;
     let dst3_stride = dst2_stride * params.bs02 * params.broadcast2;
 
@@ -37,7 +38,7 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
     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;
+    let src12_idx = dst2_idx; // src1 is not broadcast
 
     let dst2_rem = dst3_rem % dst2_stride; 
 
@@ -46,8 +47,8 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
     var sum = 0.0;
     for (var i: u32 = 0u; i < params.k; i = i + 1u) {
-        let src0_idx = src03_idx * src03_stride + src02_idx * src02_stride + col * params.k + i;
-        let src1_idx = src13_idx * src13_stride + src12_idx * src12_stride + row * params.k + i;
+        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;