volvoxai 0.1.0

package/js/ops/where.js

@@ -0,0 +1,15 @@
+export function _cpuWhere(node) {
+    const cond = node.inputs.cond || node.inputs.condition;
+    const a = node.inputs.x || node.inputs.a;
+    const b = node.inputs.y || node.inputs.b;
+    const outBuf = node.outputs.out.buffer;
+    
+    const condBuf = cond.buffer;
+    const aBuf = a.buffer;
+    const bBuf = b.buffer;
+    
+    const elements = outBuf.length;
+    for (let i = 0; i < elements; i++) {
+        outBuf[i] = condBuf[i] !== 0 ? aBuf[i] : bBuf[i];
+    }
+}

package/package.json

@@ -0,0 +1,33 @@
+{
+  "name": "volvoxai",
+  "version": "0.1.0",
+  "description": "A Zero-Dependency, Bare-Metal Deep Learning Engine for the Browser and Node.js",
+  "main": "js/index.js",
+  "type": "module",
+  "bin": {
+    "volvox": "./bin/volvox.js"
+  },
+  "files": [
+    "bin/",
+    "dist/",
+    "docs/",
+    "js/",
+    "shaders/",
+    "volvoxai.wasm"
+  ],
+  "scripts": {
+    "build": "esbuild js/index.js --bundle --format=esm --loader:.wgsl=text --outfile=dist/volvoxai.js",
+    "build:min": "esbuild js/index.js --bundle --minify --format=esm --loader:.wgsl=text --outfile=dist/volvoxai.min.js",
+    "build:all": "npm run build && npm run build:min",
+    "test": "echo 'Run tests manually for now'",
+    "start": "node ./bin/volvox.js"
+  },
+  "dependencies": {
+    "commander": "^11.1.0"
+  },
+  "devDependencies": {
+    "esbuild": "^0.28.1"
+  },
+  "author": "",
+  "license": "MIT"
+}

package/shaders/add.wgsl

@@ -0,0 +1,13 @@
+@group(0) @binding(0) var<storage, read> a : array<f32>;
+            @group(0) @binding(1) var<storage, read> b : array<f32>;
+            @group(0) @binding(2) var<storage, read_write> output : array<f32>;
+            
+            struct Params { size : u32 }
+            @group(0) @binding(3) var<uniform> params : Params;
+            
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let idx = global_id.x;
+                if (idx >= params.size) { return; }
+                output[idx] = a[idx] + b[idx];
+            }

package/shaders/add3Relu.wgsl

@@ -0,0 +1,23 @@
+@group(0) @binding(0) var<storage, read> a : array<f32>;
+@group(0) @binding(1) var<storage, read> b : array<f32>;
+@group(0) @binding(2) var<storage, read> c : array<f32>;
+@group(0) @binding(3) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    size : u32,
+    relu : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let i = gid.x;
+    if (i >= params.size) { return; }
+    var v = a[i] + b[i] + c[i];
+    if (params.relu == 1u) {
+        v = max(v, 0.0);
+    } else if (params.relu >= 2u) {
+        v = min(max(v, 0.0), 6.0);
+    }
+    output[i] = v;
+}

package/shaders/addRelu.wgsl

@@ -0,0 +1,22 @@
+@group(0) @binding(0) var<storage, read> a : array<f32>;
+@group(0) @binding(1) var<storage, read> b : array<f32>;
+@group(0) @binding(2) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    size : u32,
+    relu : u32,
+}
+@group(0) @binding(3) var<uniform> params : Params;
+
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let i = gid.x;
+    if (i >= params.size) { return; }
+    var v = a[i] + b[i];
+    if (params.relu == 1u) {
+        v = max(v, 0.0);
+    } else if (params.relu >= 2u) {
+        v = min(max(v, 0.0), 6.0);
+    }
+    output[i] = v;
+}

package/shaders/averagePool2D.wgsl

@@ -0,0 +1,24 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { b: u32, in_h: u32, in_w: u32, c: u32, out_h: u32, out_w: u32, kh: u32, kw: u32, sh: u32, sw: u32, ph: u32, pw: u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(8, 8, 1)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let x = global_id.x; let y = global_id.y; let z = global_id.z;
+            let b = z / params.c;
+            let c = z - b * params.c;
+            if (x >= params.out_w || y >= params.out_h || b >= params.b) { return; }
+            var sum = 0.0; var count = 0u;
+            for (var ky = 0u; ky < params.kh; ky = ky + 1u) {
+                for (var kx = 0u; kx < params.kw; kx = kx + 1u) {
+                    let in_y = i32(y * params.sh) - i32(params.ph) + i32(ky);
+                    let in_x = i32(x * params.sw) - i32(params.pw) + i32(kx);
+                    if (in_y >= 0 && in_y < i32(params.in_h) && in_x >= 0 && in_x < i32(params.in_w)) {
+                        sum = sum + input[((b * params.in_h + u32(in_y)) * params.in_w + u32(in_x)) * params.c + c];
+                        count = count + 1u;
+                    }
+                }
+            }
+            if (count == 0u) { count = 1u; }
+            output[((b * params.out_h + y) * params.out_w + x) * params.c + c] = sum / f32(count);
+        }

package/shaders/batchNorm2D.wgsl

@@ -0,0 +1,21 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read> weight : array<f32>;
+        @group(0) @binding(2) var<storage, read> bias : array<f32>;
+        @group(0) @binding(3) var<storage, read> running_mean : array<f32>;
+        @group(0) @binding(4) var<storage, read> running_var : array<f32>;
+        @group(0) @binding(5) var<storage, read_write> output : array<f32>;
+        struct Params { b : u32, c : u32, h : u32, w : u32, eps : f32 }
+        @group(0) @binding(6) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            let total = params.b * params.c * params.h * params.w;
+            if (idx >= total) { return; }
+            let c = idx % params.c;
+            let mean = running_mean[c];
+            let var_val = running_var[c];
+            let gamma = weight[c];
+            let beta = bias[c];
+            let inv_std = 1.0 / sqrt(var_val + params.eps);
+            output[idx] = (input[idx] - mean) * inv_std * gamma + beta;
+        }

package/shaders/binaryBroadcast.wgsl

@@ -0,0 +1,34 @@
+@group(0) @binding(0) var<storage, read> a : array<f32>;
+        @group(0) @binding(1) var<storage, read> b : array<f32>;
+        @group(0) @binding(2) var<storage, read_write> output : array<f32>;
+        
+        struct Params { size : u32, is_b_scalar : u32, b_size : u32, a_size: u32, is_a_scalar: u32 }
+        @group(0) @binding(3) var<uniform> params : Params;
+        
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            
+            var a_val : f32 = 0.0;
+            if (params.is_a_scalar == 1u) {
+                a_val = a[0];
+            } else if (params.a_size < params.size && params.a_size > 0u) {
+                a_val = a[idx % params.a_size];
+            } else {
+                a_val = a[idx];
+            }
+            
+            var b_val : f32 = 0.0;
+            if (params.is_b_scalar == 1u) {
+                b_val = b[0];
+            } else if (params.b_size < params.size && params.b_size > 0u) {
+                b_val = b[idx % params.b_size];
+            } else {
+                b_val = b[idx];
+            }
+            
+            var out_val : f32 = 0.0;
+            undefined
+            output[idx] = out_val;
+        }

package/shaders/broadcastBinary.wgsl

@@ -0,0 +1,26 @@
+@group(0) @binding(0) var<storage, read> a : array<f32>;
+        @group(0) @binding(1) var<storage, read> b : array<f32>;
+        @group(0) @binding(2) var<storage, read_write> output : array<f32>;
+        @group(0) @binding(3) var<storage, read> md : array<u32>;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+            let idx = gid.x;
+            let total = md[0];
+            if (idx >= total) { return; }
+            let rank = md[1];
+            var rem = idx;
+            var a_idx = 0u;
+            var b_idx = 0u;
+            for (var d = 0u; d < rank; d = d + 1u) {
+                let os = md[2u + d];
+                let coord = rem / os;
+                rem = rem - coord * os;
+                a_idx = a_idx + coord * md[2u + rank + d];
+                b_idx = b_idx + coord * md[2u + 2u * rank + d];
+            }
+            let av = a[a_idx];
+            let bv = b[b_idx];
+            var out_val = 0.0;
+            //__BINOP__
+            output[idx] = out_val;
+        }

package/shaders/clip.wgsl

@@ -0,0 +1,10 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32, min_v : f32, max_v : f32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            output[idx] = clamp(input[idx], params.min_v, params.max_v);
+        }

package/shaders/concat2.wgsl

@@ -0,0 +1,16 @@
+@group(0) @binding(0) var<storage, read> a : array<f32>;
+            @group(0) @binding(1) var<storage, read> b : array<f32>;
+            @group(0) @binding(2) var<storage, read_write> output : array<f32>;
+            
+            struct Params { a_size : u32, b_size : u32 }
+            @group(0) @binding(3) var<uniform> params : Params;
+            
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let idx = global_id.x;
+                if (idx < params.a_size) {
+                    output[idx] = a[idx];
+                } else if (idx < params.a_size + params.b_size) {
+                    output[idx] = b[idx - params.a_size];
+                }
+            }

package/shaders/concatCopy.wgsl

@@ -0,0 +1,10 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32, offset : u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+            let i = gid.x;
+            if (i >= params.size) { return; }
+            output[params.offset + i] = input[i];
+        }

package/shaders/concatSigmoidCopy.wgsl

@@ -0,0 +1,16 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    size : u32,
+    offset : u32,
+}
+@group(0) @binding(2) var<uniform> params : Params;
+
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let i = gid.x;
+    if (i >= params.size) { return; }
+    let v = input[i];
+    output[params.offset + i] = 1.0 / (1.0 + exp(-v));
+}

package/shaders/conv1D.wgsl

@@ -0,0 +1,37 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+            @group(0) @binding(1) var<storage, read> weight : array<f32>;
+            @group(0) @binding(2) var<storage, read> bias : array<f32>;
+            @group(0) @binding(3) var<storage, read_write> output : array<f32>;
+            
+            struct Params {
+                in_c : u32, in_l : u32,
+                out_c : u32, k : u32,
+                stride : u32, pad : u32, relu : u32
+            }
+            @group(0) @binding(4) var<uniform> params : Params;
+            
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let x = global_id.x;
+                let oc = global_id.y;
+                let out_l = (params.in_l + 2u * params.pad - params.k) / params.stride + 1u;
+                
+                if (x >= out_l || oc >= params.out_c) { return; }
+                
+                var sum = bias[oc];
+                let w_base = oc * params.in_c * params.k;
+                
+                for (var ic = 0u; ic < params.in_c; ic = ic + 1u) {
+                    let in_base = ic * params.in_l;
+                    let w_ic_base = w_base + ic * params.k;
+                    for (var k = 0u; k < params.k; k = k + 1u) {
+                        let in_x = i32(x * params.stride + k) - i32(params.pad);
+                        if (in_x >= 0 && in_x < i32(params.in_l)) {
+                            sum = sum + input[in_base + u32(in_x)] * weight[w_ic_base + k];
+                        }
+                    }
+                }
+                
+                if (params.relu == 1u && sum < 0.0) { sum = 0.0; }
+                output[oc * out_l + x] = sum;
+            }

package/shaders/conv2D.wgsl

@@ -0,0 +1,80 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read> weight : array<f32>;
+@group(0) @binding(2) var<storage, read> bias : array<f32>;
+@group(0) @binding(3) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    n : u32,
+    in_h : u32,
+    in_w : u32,
+    in_c : u32,
+    out_c : u32,
+    out_h : u32,
+    out_w : u32,
+    kh : u32,
+    kw : u32,
+    sy : u32,
+    sx : u32,
+    pt : u32,
+    pl : u32,
+    groups : u32,
+    relu : u32,
+    dy : u32,
+    dx : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let z = gid.z;
+    let nb = z / params.out_c;
+    let oc = z - nb * params.out_c;
+    if (nb >= params.n || ox >= params.out_w || oy >= params.out_h) { return; }
+
+    var sum = 0.0;
+    if (params.groups == params.in_c) {
+        let mult = params.out_c / params.in_c;
+        let ic = oc / mult;
+        let m = oc - ic * mult;
+        for (var yy = 0u; yy < params.kh; yy = yy + 1u) {
+            let iy = i32(oy * params.sy + yy * params.dy) - i32(params.pt);
+            if (iy < 0 || iy >= i32(params.in_h)) { continue; }
+            for (var xx = 0u; xx < params.kw; xx = xx + 1u) {
+                let ix = i32(ox * params.sx + xx * params.dx) - i32(params.pl);
+                if (ix < 0 || ix >= i32(params.in_w)) { continue; }
+                let ii = ((nb * params.in_h + u32(iy)) * params.in_w + u32(ix)) * params.in_c + ic;
+                let wi = (((yy * params.kw + xx) * params.in_c + ic) * mult) + m;
+                sum = sum + input[ii] * weight[wi];
+            }
+        }
+    } else {
+        let out_per_g = params.out_c / params.groups;
+        let in_per_g = params.in_c / params.groups;
+        let g = oc / out_per_g;
+        let ic0 = g * in_per_g;
+        for (var icl = 0u; icl < in_per_g; icl = icl + 1u) {
+            let ic = ic0 + icl;
+            for (var yy = 0u; yy < params.kh; yy = yy + 1u) {
+                let iy = i32(oy * params.sy + yy * params.dy) - i32(params.pt);
+                if (iy < 0 || iy >= i32(params.in_h)) { continue; }
+                for (var xx = 0u; xx < params.kw; xx = xx + 1u) {
+                    let ix = i32(ox * params.sx + xx * params.dx) - i32(params.pl);
+                    if (ix < 0 || ix >= i32(params.in_w)) { continue; }
+                    let ii = ((nb * params.in_h + u32(iy)) * params.in_w + u32(ix)) * params.in_c + ic;
+                    let wi = (((yy * params.kw + xx) * params.in_c + ic) * params.out_c) + oc;
+                    sum = sum + input[ii] * weight[wi];
+                }
+            }
+        }
+    }
+
+    var v = sum + bias[oc];
+    if (params.relu == 1u) {
+        v = max(v, 0.0);
+    } else if (params.relu >= 2u) {
+        v = min(max(v, 0.0), 6.0);
+    }
+    output[((nb * params.out_h + oy) * params.out_w + ox) * params.out_c + oc] = v;
+}

package/shaders/conv2DDepthwise4.wgsl

@@ -0,0 +1,74 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read> weight : array<f32>;
+@group(0) @binding(2) var<storage, read> bias : array<f32>;
+@group(0) @binding(3) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    n : u32,
+    in_h : u32,
+    in_w : u32,
+    in_c : u32,
+    out_c : u32,
+    out_h : u32,
+    out_w : u32,
+    kh : u32,
+    kw : u32,
+    sy : u32,
+    sx : u32,
+    pt : u32,
+    pl : u32,
+    groups : u32,
+    relu : u32,
+    dy : u32,
+    dx : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+fn apply_relu(v_in : f32) -> f32 {
+    var v = v_in;
+    if (params.relu == 1u) {
+        v = max(v, 0.0);
+    } else if (params.relu >= 2u) {
+        v = min(max(v, 0.0), 6.0);
+    }
+    return v;
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let c4_count = (params.out_c + 3u) / 4u;
+    let nb = gid.z / c4_count;
+    let ch0 = (gid.z - nb * c4_count) * 4u;
+    if (nb >= params.n || ox >= params.out_w || oy >= params.out_h || ch0 >= params.out_c) { return; }
+
+    let has1 = ch0 + 1u < params.out_c;
+    let has2 = ch0 + 2u < params.out_c;
+    let has3 = ch0 + 3u < params.out_c;
+    var sum0 = 0.0;
+    var sum1 = 0.0;
+    var sum2 = 0.0;
+    var sum3 = 0.0;
+
+    for (var yy = 0u; yy < params.kh; yy = yy + 1u) {
+        let iy = i32(oy * params.sy + yy * params.dy) - i32(params.pt);
+        if (iy < 0 || iy >= i32(params.in_h)) { continue; }
+        for (var xx = 0u; xx < params.kw; xx = xx + 1u) {
+            let ix = i32(ox * params.sx + xx * params.dx) - i32(params.pl);
+            if (ix < 0 || ix >= i32(params.in_w)) { continue; }
+            let input_base = ((nb * params.in_h + u32(iy)) * params.in_w + u32(ix)) * params.in_c + ch0;
+            let weight_base = ((yy * params.kw + xx) * params.in_c) + ch0;
+            sum0 = sum0 + input[input_base] * weight[weight_base];
+            if (has1) { sum1 = sum1 + input[input_base + 1u] * weight[weight_base + 1u]; }
+            if (has2) { sum2 = sum2 + input[input_base + 2u] * weight[weight_base + 2u]; }
+            if (has3) { sum3 = sum3 + input[input_base + 3u] * weight[weight_base + 3u]; }
+        }
+    }
+
+    let output_base = ((nb * params.out_h + oy) * params.out_w + ox) * params.out_c + ch0;
+    output[output_base] = apply_relu(sum0 + bias[ch0]);
+    if (has1) { output[output_base + 1u] = apply_relu(sum1 + bias[ch0 + 1u]); }
+    if (has2) { output[output_base + 2u] = apply_relu(sum2 + bias[ch0 + 2u]); }
+    if (has3) { output[output_base + 3u] = apply_relu(sum3 + bias[ch0 + 3u]); }
+}

package/shaders/conv2DDepthwise8.wgsl

@@ -0,0 +1,66 @@
+@group(0) @binding(0) var<storage, read> input : array<vec4<f32>>;
+@group(0) @binding(1) var<storage, read> weight : array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> bias : array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read_write> output : array<vec4<f32>>;
+
+struct Params {
+    n : u32,
+    in_h : u32,
+    in_w : u32,
+    in_c : u32,
+    out_c : u32,
+    out_h : u32,
+    out_w : u32,
+    kh : u32,
+    kw : u32,
+    sy : u32,
+    sx : u32,
+    pt : u32,
+    pl : u32,
+    groups : u32,
+    relu : u32,
+    dy : u32,
+    dx : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+fn apply_relu4(v_in : vec4<f32>) -> vec4<f32> {
+    var v = v_in;
+    if (params.relu == 1u) {
+        v = max(v, vec4<f32>(0.0));
+    } else if (params.relu >= 2u) {
+        v = min(max(v, vec4<f32>(0.0)), vec4<f32>(6.0));
+    }
+    return v;
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let c8_count = (params.out_c + 7u) / 8u;
+    let nb = gid.z / c8_count;
+    let ch0 = (gid.z - nb * c8_count) * 8u;
+    if (nb >= params.n || ox >= params.out_w || oy >= params.out_h || ch0 >= params.out_c) { return; }
+
+    var sum0 = vec4<f32>(0.0);
+    var sum1 = vec4<f32>(0.0);
+
+    for (var yy = 0u; yy < params.kh; yy = yy + 1u) {
+        let iy = i32(oy * params.sy + yy * params.dy) - i32(params.pt);
+        if (iy < 0 || iy >= i32(params.in_h)) { continue; }
+        for (var xx = 0u; xx < params.kw; xx = xx + 1u) {
+            let ix = i32(ox * params.sx + xx * params.dx) - i32(params.pl);
+            if (ix < 0 || ix >= i32(params.in_w)) { continue; }
+            let input_base = (((nb * params.in_h + u32(iy)) * params.in_w + u32(ix)) * params.in_c + ch0) / 4u;
+            let weight_base = (((yy * params.kw + xx) * params.in_c) + ch0) / 4u;
+            sum0 = sum0 + input[input_base] * weight[weight_base];
+            sum1 = sum1 + input[input_base + 1u] * weight[weight_base + 1u];
+        }
+    }
+
+    let output_base = (((nb * params.out_h + oy) * params.out_w + ox) * params.out_c + ch0) / 4u;
+    let bias_base = ch0 / 4u;
+    output[output_base] = apply_relu4(sum0 + bias[bias_base]);
+    output[output_base + 1u] = apply_relu4(sum1 + bias[bias_base + 1u]);
+}

package/shaders/conv2DPointwise16.wgsl

@@ -0,0 +1,67 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read> weight : array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> bias : array<vec4<f32>>;
+@group(0) @binding(3) var<storage, read_write> output : array<vec4<f32>>;
+
+struct Params {
+    n : u32,
+    in_h : u32,
+    in_w : u32,
+    in_c : u32,
+    out_c : u32,
+    out_h : u32,
+    out_w : u32,
+    kh : u32,
+    kw : u32,
+    sy : u32,
+    sx : u32,
+    pt : u32,
+    pl : u32,
+    groups : u32,
+    relu : u32,
+    dy : u32,
+    dx : u32,
+}
+@group(0) @binding(4) var<uniform> params : Params;
+
+fn apply_relu4(v_in : vec4<f32>) -> vec4<f32> {
+    var v = v_in;
+    if (params.relu == 1u) {
+        v = max(v, vec4<f32>(0.0));
+    } else if (params.relu >= 2u) {
+        v = min(max(v, vec4<f32>(0.0)), vec4<f32>(6.0));
+    }
+    return v;
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let oc16_count = (params.out_c + 15u) / 16u;
+    let nb = gid.z / oc16_count;
+    let oc0 = (gid.z - nb * oc16_count) * 16u;
+    if (nb >= params.n || ox >= params.out_w || oy >= params.out_h || oc0 >= params.out_c) { return; }
+
+    var sum0 = vec4<f32>(0.0);
+    var sum1 = vec4<f32>(0.0);
+    var sum2 = vec4<f32>(0.0);
+    var sum3 = vec4<f32>(0.0);
+
+    let input_base = ((nb * params.in_h + oy) * params.in_w + ox) * params.in_c;
+    for (var ic = 0u; ic < params.in_c; ic = ic + 1u) {
+        let x = input[input_base + ic];
+        let wbase = (ic * params.out_c + oc0) / 4u;
+        sum0 = sum0 + x * weight[wbase];
+        sum1 = sum1 + x * weight[wbase + 1u];
+        sum2 = sum2 + x * weight[wbase + 2u];
+        sum3 = sum3 + x * weight[wbase + 3u];
+    }
+
+    let output_base = (((nb * params.out_h + oy) * params.out_w + ox) * params.out_c + oc0) / 4u;
+    let bias_base = oc0 / 4u;
+    output[output_base] = apply_relu4(sum0 + bias[bias_base]);
+    output[output_base + 1u] = apply_relu4(sum1 + bias[bias_base + 1u]);
+    output[output_base + 2u] = apply_relu4(sum2 + bias[bias_base + 2u]);
+    output[output_base + 3u] = apply_relu4(sum3 + bias[bias_base + 3u]);
+}