wgblas 0.1.0 → 0.1.2

package/README.md

@@ -51,24 +51,21 @@ Restart Firefox after making changes.
 
 ## Requirements
 
-- Node.js 18+
+- Node.js 22+
 
-## Example usage
-
-This package is not yet published to npm. Clone the repo and install it locally in your project:
+## Installation
 
 ```sh
-git clone https://github.com/manit2004/wgblas.git
-cd your-project
-npm install /path/to/wgblas
+npm install wgblas
 ```
 
+## Example usage
+
 ### Example Code Snippet
 
 ```js
-import { init, cleanup } from "wgblas";
+import { init, cleanup, randomFloat32Array } from "wgblas";
 import { sscal } from "wgblas/sscal";
-import { randomFloat32Array } from "wgblas/util/random";
 
 const device = await init();
 
@@ -92,7 +89,7 @@ No bundler needed. Load the pre-built browser bundle from the CDN and use `windo
   <head>
     <meta charset="UTF-8" />
     <title>sscal — wgblas browser example</title>
-    <script src="https://manit2004.github.io/wgblas/wgblas.browser.js"></script>
+    <script src="https://unpkg.com/wgblas/dist/wgblas.browser.js"></script>
   </head>
   <body>
     <pre id="out">Running…</pre>
@@ -126,11 +123,10 @@ No bundler needed. Load the pre-built browser bundle from the CDN and use `windo
 `GpuVector` keeps data resident on the GPU between operations — upload once, chain any number of operations, read back once. This eliminates the redundant uploads and readbacks between steps, which are often more expensive than the compute itself.
 
 ```js
-import { init, cleanup } from "wgblas";
+import { init, cleanup, randomFloat32Array } from "wgblas";
 import { saxpy } from "wgblas/saxpy";
 import { sscal } from "wgblas/sscal";
 import { GpuVector } from "wgblas/classes/GpuVector";
-import { randomFloat32Array } from "wgblas/util/random";
 
 const device = await init();
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "wgblas",
-  "version": "0.1.0",
+  "version": "0.1.2",
   "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;
+  }
+}