wgblas 0.1.0 → 0.1.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "wgblas",
-  "version": "0.1.0",
+  "version": "0.1.1",
   "description": "BLAS on WebGPU",
   "type": "module",
   "main": "index.mjs",
@@ -65,6 +65,7 @@
     "index.d.mts",
     "src/**/*.mjs",
     "src/**/*.d.mts",
+    "src/**/*.wgsl",
     "dist/wgblas.browser.js"
   ],
   "repository": {

package/src/shaders/isamax.wgsl

@@ -0,0 +1,59 @@
+// isamax: returns index of element with largest absolute value
+// pass 1 dispatches exactly 2 * WGS workgroups; pass 2 uses reduction/argmax.wgsl.
+
+@group(0) @binding(0) var<storage, read>       x:            array<f32>;
+@group(0) @binding(1) var<storage, read_write> partials_val: array<f32>;
+@group(0) @binding(2) var<storage, read_write> partials_idx: array<u32>;
+@group(0) @binding(3) var<uniform>             params:       Params;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+}
+
+const WGS: u32 = 64;
+
+var<workgroup> tile_val: array<f32, 64>;
+var<workgroup> tile_idx: array<u32, 64>;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid:    vec3u,
+  @builtin(local_invocation_id)  lid:    vec3u,
+  @builtin(workgroup_id)         wgid:   vec3u,
+  @builtin(num_workgroups)       num_wg: vec3u,
+) {
+  // -1.0 is a safe sentinel: any |x[i]| >= 0 beats it,
+  // so workgroups with no elements lose gracefully in the epilogue.
+  var best_val: f32 = -1.0;
+  var best_idx: u32 = 0u;
+
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    let v = abs(x[id * params.x_inc]);
+    if (v > best_val) {
+      best_val = v;
+      best_idx = id;
+    }
+  }
+
+  tile_val[lid.x] = best_val;
+  tile_idx[lid.x] = best_idx;
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    if (lid.x < s) {
+      let a_val = tile_val[lid.x];
+      let b_val = tile_val[lid.x + s];
+      if (b_val > a_val || (b_val == a_val && tile_idx[lid.x + s] < tile_idx[lid.x])) {
+        tile_val[lid.x] = b_val;
+        tile_idx[lid.x] = tile_idx[lid.x + s];
+      }
+    }
+    workgroupBarrier();
+  }
+
+  if (lid.x == 0u) {
+    partials_val[wgid.x] = tile_val[0];
+    partials_idx[wgid.x] = tile_idx[0];
+  }
+}

package/src/shaders/reduction/argmax.wgsl

@@ -0,0 +1,43 @@
+// amax reduction: collapses 2*WGS (value, index) pairs into one index.
+// dispatch: 1 workgroup of WGS threads.
+// partials_val and partials_idx must have exactly 2*WGS entries.
+
+@group(0) @binding(0) var<storage, read>       partials_val: array<f32>;
+@group(0) @binding(1) var<storage, read>       partials_idx: array<u32>;
+@group(0) @binding(2) var<storage, read_write> result:       array<u32>;
+
+const WGS: u32 = 64;
+
+var<workgroup> tile_val: array<f32, 64>;
+var<workgroup> tile_idx: array<u32, 64>;
+
+@compute @workgroup_size(64)
+fn reduce(
+  @builtin(local_invocation_id) lid: vec3u,
+) {
+  let i = lid.x;
+  let a_val = partials_val[i];
+  let b_val = partials_val[i + WGS];
+  if (b_val > a_val || (b_val == a_val && partials_idx[i + WGS] < partials_idx[i])) {
+    tile_val[i] = b_val;
+    tile_idx[i] = partials_idx[i + WGS];
+  } else {
+    tile_val[i] = a_val;
+    tile_idx[i] = partials_idx[i];
+  }
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    if (i < s) {
+      let c_val = tile_val[i];
+      let d_val = tile_val[i + s];
+      if (d_val > c_val || (d_val == c_val && tile_idx[i + s] < tile_idx[i])) {
+        tile_val[i] = d_val;
+        tile_idx[i] = tile_idx[i + s];
+      }
+    }
+    workgroupBarrier();
+  }
+
+  if (i == 0u) { result[0] = tile_idx[0]; }
+}

package/src/shaders/reduction/sum.wgsl

@@ -0,0 +1,26 @@
+// sum reduction: collapses 2*WGS partials into one scalar.
+// dispatch: 1 workgroup of WGS threads.
+// partials must have exactly 2*WGS entries.
+
+@group(0) @binding(0) var<storage, read>       partials: array<f32>;
+@group(0) @binding(1) var<storage, read_write> result:   array<f32>;
+
+const WGS: u32 = 64;
+
+var<workgroup> tile: array<f32, 64>;
+
+@compute @workgroup_size(64)
+fn reduce(
+  @builtin(local_invocation_id) lid: vec3u,
+) {
+  let i = lid.x;
+  tile[i] = partials[i] + partials[i + WGS];
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    if (i < s) { tile[i] += tile[i + s]; }
+    workgroupBarrier();
+  }
+
+  if (i == 0u) { result[0] = tile[0]; }
+}

package/src/shaders/sasum.wgsl

@@ -0,0 +1,37 @@
+// sasum: result = sum(|x[i]|)
+// pass 1 dispatches exactly 2 * WGS workgroups; pass 2 uses reduction/abssum.wgsl.
+
+@group(0) @binding(0) var<storage, read>       x:        array<f32>;
+@group(0) @binding(1) var<storage, read_write> partials: array<f32>;
+@group(0) @binding(2) var<uniform>             params:   Params;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+}
+
+const WGS: u32 = 64;
+
+var<workgroup> tile: array<f32, 64>;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid:    vec3u,
+  @builtin(local_invocation_id)  lid:    vec3u,
+  @builtin(workgroup_id)         wgid:   vec3u,
+  @builtin(num_workgroups)       num_wg: vec3u,
+) {
+  var acc: f32 = 0.0;
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    acc += abs(x[id * params.x_inc]);
+  }
+  tile[lid.x] = acc;
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    if (lid.x < s) { tile[lid.x] += tile[lid.x + s]; }
+    workgroupBarrier();
+  }
+
+  if (lid.x == 0u) { partials[wgid.x] = tile[0]; }
+}

package/src/shaders/saxpy.wgsl

@@ -0,0 +1,25 @@
+// saxpy: y = alpha * x + y
+
+@group(0) @binding(0) var<storage, read>       x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+
+struct Params {
+  n:     u32,
+  alpha: f32,
+  x_inc: u32,
+  y_inc: u32,
+}
+
+@group(0) @binding(2) var<uniform> params: Params;
+
+const WGS: u32 = 64;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid: vec3u,
+  @builtin(num_workgroups) num_wg: vec3u,
+) {
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    y[id * params.y_inc] = params.alpha * x[id * params.x_inc] + y[id * params.y_inc];
+  }
+}

package/src/shaders/scopy.wgsl

@@ -0,0 +1,24 @@
+// scopy: y = x
+
+@group(0) @binding(0) var<storage, read>       x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+  y_inc: u32,
+}
+
+@group(0) @binding(2) var<uniform> params: Params;
+
+const WGS: u32 = 64;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid: vec3u,
+  @builtin(num_workgroups) num_wg: vec3u,
+) {
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    y[id * params.y_inc] = x[id * params.x_inc];
+  }
+}

package/src/shaders/sdot.wgsl

@@ -0,0 +1,39 @@
+// sdot: result = sum(x[i] * y[i])
+// pass 1 dispatches exactly 2 * WGS workgroups; pass 2 uses reduction/sum.wgsl.
+
+@group(0) @binding(0) var<storage, read>       x:        array<f32>;
+@group(0) @binding(1) var<storage, read>       y:        array<f32>;
+@group(0) @binding(2) var<storage, read_write> partials: array<f32>;
+@group(0) @binding(3) var<uniform>             params:   Params;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+  y_inc: u32,
+}
+
+const WGS: u32 = 64;
+
+var<workgroup> tile: array<f32, 64>;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid:    vec3u,
+  @builtin(local_invocation_id)  lid:    vec3u,
+  @builtin(workgroup_id)         wgid:   vec3u,
+  @builtin(num_workgroups)       num_wg: vec3u,
+) {
+  var acc: f32 = 0.0;
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    acc += x[id * params.x_inc] * y[id * params.y_inc];
+  }
+  tile[lid.x] = acc;
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    if (lid.x < s) { tile[lid.x] += tile[lid.x + s]; }
+    workgroupBarrier();
+  }
+
+  if (lid.x == 0u) { partials[wgid.x] = tile[0]; }
+}

package/src/shaders/snrm2.wgsl

@@ -0,0 +1,38 @@
+// snrm2: result = sqrt(sum(x[i] * x[i]))
+// pass 1 dispatches exactly 2 * WGS workgroups; pass 2 uses reduction/sqsum.wgsl.
+
+@group(0) @binding(0) var<storage, read>       x:        array<f32>;
+@group(0) @binding(1) var<storage, read_write> partials: array<f32>;
+@group(0) @binding(2) var<uniform>             params:   Params;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+}
+
+const WGS: u32 = 64;
+
+var<workgroup> tile: array<f32, 64>;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid:    vec3u,
+  @builtin(local_invocation_id)  lid:    vec3u,
+  @builtin(workgroup_id)         wgid:   vec3u,
+  @builtin(num_workgroups)       num_wg: vec3u,
+) {
+  var acc: f32 = 0.0;
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    let v = x[id * params.x_inc];
+    acc += v * v;
+  }
+  tile[lid.x] = acc;
+  workgroupBarrier();
+
+  for (var s = WGS / 2u; s > 0u; s >>= 1u) {
+    if (lid.x < s) { tile[lid.x] += tile[lid.x + s]; }
+    workgroupBarrier();
+  }
+
+  if (lid.x == 0u) { partials[wgid.x] = tile[0]; }
+}

package/src/shaders/srot.wgsl

@@ -0,0 +1,29 @@
+// srot: x = c*x + s*y,  y = -s*x + c*y
+
+@group(0) @binding(0) var<storage, read_write> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+
+struct Params {
+  n:     u32,
+  c:     f32,
+  s:     f32,
+  x_inc: u32,
+  y_inc: u32,
+}
+
+@group(0) @binding(2) var<uniform> params: Params;
+
+const WGS: u32 = 64;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid: vec3u,
+  @builtin(num_workgroups) num_wg: vec3u,
+) {
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    let xi = x[id * params.x_inc];
+    let yi = y[id * params.y_inc];
+    x[id * params.x_inc] =  params.c * xi + params.s * yi;
+    y[id * params.y_inc] = -params.s * xi + params.c * yi;
+  }
+}

package/src/shaders/srotm.wgsl

@@ -0,0 +1,50 @@
+// srotm: applies modified Givens rotation H to vectors x and y.
+// param[0] = flag: -1 (full H), 0 (unit diagonal), 1 (unit off-diagonal)
+// param = [ flag, h11, h21, h12, h22 ]
+// flag == -2 (identity/no-op) is handled in JS before dispatch reaches here.
+
+@group(0) @binding(0) var<storage, read_write> x:     array<f32>;
+@group(0) @binding(1) var<storage, read_write> y:     array<f32>;
+@group(0) @binding(2) var<storage, read>       param: array<f32>;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+  y_inc: u32,
+}
+
+@group(0) @binding(3) var<uniform> params: Params;
+
+const WGS: u32 = 64;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid: vec3u,
+  @builtin(num_workgroups) num_wg: vec3u,
+) {
+  let flag = param[0];
+
+  var h11: f32; var h12: f32;
+  var h21: f32; var h22: f32;
+
+  if (flag == -1.0) {
+    // full 2x2 matrix
+    h11 = param[1]; h21 = param[2];
+    h12 = param[3]; h22 = param[4];
+  } else if (flag == 0.0) {
+    // diagonal fixed at 1
+    h11 = 1.0;      h21 = param[2];
+    h12 = param[3]; h22 = 1.0;
+  } else if (flag == 1.0) {
+    // flag == 1.0: off-diagonal fixed at +1 / -1
+    h11 = param[1]; h21 = -1.0;
+    h12 = 1.0;      h22 = param[4];
+  }
+
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    let xi = x[id * params.x_inc];
+    let yi = y[id * params.y_inc];
+    x[id * params.x_inc] = h11 * xi + h12 * yi;
+    y[id * params.y_inc] = h21 * xi + h22 * yi;
+  }
+}

package/src/shaders/sscal.wgsl

@@ -0,0 +1,23 @@
+// sscal: x = alpha * x
+
+@group(0) @binding(0) var<storage, read_write> x: array<f32>;
+
+struct Params {
+  n:     u32,
+  alpha: f32,
+  x_inc: u32,
+}
+
+@group(0) @binding(1) var<uniform> params: Params;
+
+const WGS: u32 = 64;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid: vec3u,
+  @builtin(num_workgroups) num_wg: vec3u,
+) {
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    x[id * params.x_inc] = params.alpha * x[id * params.x_inc];
+  }
+}

package/src/shaders/sswap.wgsl

@@ -0,0 +1,26 @@
+// sswap: x <-> y
+
+@group(0) @binding(0) var<storage, read_write> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+
+struct Params {
+  n:     u32,
+  x_inc: u32,
+  y_inc: u32,
+}
+
+@group(0) @binding(2) var<uniform> params: Params;
+
+const WGS: u32 = 64;
+
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) gid: vec3u,
+  @builtin(num_workgroups) num_wg: vec3u,
+) {
+  for (var id = gid.x; id < params.n; id += num_wg.x * WGS) {
+    let temp = x[id * params.x_inc];
+    x[id * params.x_inc] = y[id * params.y_inc];
+    y[id * params.y_inc] = temp;
+  }
+}