feat: use uniforms in wgsl

backend/types.ts

@@ -1,4 +1,4 @@
-export type BackendType = "cpu" | "wasm" | "webgpu";
+export type BackendType = "js" | "wasm" | "webgpu";
 export type DataType = "u32" | "i32" | "f32";
 // deno-fmt-ignore
 export type DataArray<T extends DataType> = 
@@ -40,9 +40,3 @@ export interface Backend {
 export interface BackendRequest<T extends DataType = DataType> {
   data: Data<T>[];
 }
-
-export type Operator<T extends DataType = DataType> = (
-  backend: Backend,
-  data: DataArray<T>[],
-  meta: unknown,
-) => Promise<void>;

backend/webgpu/data.ts

@@ -1,26 +1,6 @@
 import { Data, DataArray, DataArrayConstructor, DataType } from "../types.ts";
 import { WebGPUBackend } from "./backend.ts";
 
-export interface WebGPUData<T extends DataType = DataType> extends Data<T> {
-  type: T;
-  backend: WebGPUBackend;
-  buffer: GPUBuffer;
-}
-
-export interface WebGPUDataConstructor<T extends DataType = DataType> {
-  from(
-    backend: WebGPUBackend,
-    source: DataArray<T>,
-  ): Promise<WebGPUData<T>>;
-
-  new (
-    backend: WebGPUBackend,
-    type: DataType,
-    length: number,
-    usage?: number,
-  ): WebGPUData<T>;
-}
-
 export class WebGPUData<T extends DataType = DataType> implements Data<T> {
   type: T;
   backend: WebGPUBackend;
@@ -32,6 +12,7 @@ export class WebGPUData<T extends DataType = DataType> implements Data<T> {
   static async from<T extends DataType>(
     backend: WebGPUBackend,
     source: DataArray<T>,
+    usage?: number,
   ): Promise<WebGPUData<T>> {
     // deno-fmt-ignore
     const type = (
@@ -40,7 +21,7 @@ export class WebGPUData<T extends DataType = DataType> implements Data<T> {
       : source instanceof Float32Array ? "f32"
       : undefined
     )! as T;
-    const data = new this(backend, type, source.length);
+    const data = new this(backend, type, source.length, usage);
     await data.set(source);
     return data;
   }

backend/webgpu/operators/matmul.ts

@@ -13,12 +13,15 @@ export async function matmul<T extends DataType>(
 ) {
   const type = ensureType(a.type, b.type, c.type);
   const pipeline = await backend.register(shader(type));
-
-  const meta = await WebGPUData.from(backend, new Uint32Array([m, n, k]));
+  const uniform = await WebGPUData.from(
+    backend,
+    new Uint32Array([m, n, k]),
+    GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM,
+  );
 
   await backend.execute({
     pipeline,
-    data: [a, b, c, meta],
+    data: [a, b, c, uniform],
     workgroups: [Math.ceil(n / 8), Math.ceil(m / 8), 1],
   });
 }

backend/webgpu/operators/pad.ts

@@ -12,12 +12,15 @@ export async function pad<T extends DataType>(
 ) {
   const type = ensureType(a.type, b.type);
   const pipeline = await backend.register(shader(type));
-
-  const meta = await WebGPUData.from(backend, new Uint32Array([w, h, t]));
+  const uniform = await WebGPUData.from(
+    backend,
+    new Uint32Array([w, h, t]),
+    GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM,
+  );
 
   await backend.execute({
     pipeline,
-    data: [a, b, meta],
+    data: [a, b, uniform],
     workgroups: [Math.ceil(w / 8), Math.ceil(h / 8), 1],
   });
 }

backend/webgpu/operators/transpose.ts

@@ -12,12 +12,15 @@ export async function transpose<T extends DataType>(
 ) {
   const type = ensureType(a.type, b.type);
   const pipeline = await backend.register(shader(type));
-
-  const meta = await WebGPUData.from(backend, new Uint32Array([w, h]));
+  const uniform = await WebGPUData.from(
+    backend,
+    new Uint32Array([w, h]),
+    GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM,
+  );
 
   await backend.execute({
     pipeline,
-    data: [a, b, meta],
+    data: [a, b, uniform],
     workgroups: [Math.ceil(w / 8), Math.ceil(h / 8), 1],
   });
 }

backend/webgpu/shaders/matmul.ts

@@ -2,7 +2,7 @@ import { DataType } from "../../types.ts";
 
 export const matmul = (type: DataType) => `
 [[block]]
-struct Meta {
+struct Uniform {
   m: u32;
   n: u32;
   k: u32;
@@ -20,18 +20,18 @@ var<storage, read> b: Data;
 [[group(0), binding(2)]]
 var<storage, write> c: Data;
 [[group(0), binding(3)]]
-var<storage, read> meta: Meta;
+var<uniform> uniform: Uniform;
 
 [[stage(compute), workgroup_size(8, 8, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-  if (global_id.x >= meta.n || global_id.y >= meta.m) {
+  if (global_id.x >= uniform.n || global_id.y >= uniform.m) {
     return;
   }
 
   var sum = 0${type};
-  for (var k = 0u; k < meta.k; k = k + 1u) {
-    sum = sum + a.values[global_id.y * meta.k + k] * b.values[k * meta.n + global_id.x];
+  for (var k = 0u; k < uniform.k; k = k + 1u) {
+    sum = sum + a.values[global_id.y * uniform.k + k] * b.values[k * uniform.n + global_id.x];
   }
-  c.values[global_id.x + global_id.y * meta.n] = sum;
+  c.values[global_id.x + global_id.y * uniform.n] = sum;
 }
 `;

backend/webgpu/shaders/pad.ts

@@ -2,7 +2,7 @@ import { DataType } from "../../types.ts";
 
 export const pad = (type: DataType) => `
 [[block]]
-struct Meta {
+struct Uniform {
   w: u32;
   h: u32;
   n: u32;
@@ -18,14 +18,14 @@ var<storage, read> a: Data;
 [[group(0), binding(1)]]
 var<storage, write> b: Data;
 [[group(0), binding(2)]]
-var<storage, read> meta: Meta;
+var<uniform> uniform: Uniform;
 
 [[stage(compute), workgroup_size(8, 8, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-  if (global_id.x >= meta.w || global_id.y >= meta.h) {
+  if (global_id.x >= uniform.w || global_id.y >= uniform.h) {
     return;
   }
 
-  b.values[global_id.x + global_id.y * meta.n] = a.values[global_id.x + global_id.y * meta.w];
+  b.values[global_id.x + global_id.y * uniform.n] = a.values[global_id.x + global_id.y * uniform.w];
 }
 `;

backend/webgpu/shaders/transpose.ts

@@ -2,7 +2,7 @@ import { DataType } from "../../types.ts";
 
 export const transpose = (type: DataType) => `
 [[block]]
-struct Meta {
+struct Uniform {
   w: u32;
   h: u32;
 };
@@ -17,14 +17,14 @@ var<storage, read> a: Data;
 [[group(0), binding(1)]]
 var<storage, write> b: Data;
 [[group(0), binding(2)]]
-var<storage, read> meta: Meta;
+var<uniform> uniform: Uniform;
 
 [[stage(compute), workgroup_size(8, 8, 1)]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-  if (global_id.x >= meta.w || global_id.y >= meta.h) {
+  if (global_id.x >= uniform.w || global_id.y >= uniform.h) {
     return;
   }
 
-  b.values[global_id.y + global_id.x * meta.h] = a.values[global_id.x + global_id.y * meta.w];
+  b.values[global_id.y + global_id.x * uniform.h] = a.values[global_id.x + global_id.y * uniform.w];
 }
 `;

backend/webgpu/types.ts

@@ -1,5 +1,4 @@
-import { BackendRequest, DataType, Operator } from "../types.ts";
-import { WebGPUBackend } from "./backend.ts";
+import { BackendRequest, DataType } from "../types.ts";
 import { WebGPUData } from "./data.ts";
 
 export type Workgroups =
@@ -13,8 +12,3 @@ export interface WebGPUBackendRequest<T extends DataType = DataType>
   data: WebGPUData<T>[];
   workgroups: Workgroups;
 }
-
-export type WebGPUOperator<T extends DataType = DataType> = Operator<T> & {
-  backend: WebGPUBackend;
-  data: WebGPUData<T>[];
-};

examples/webgpu_matmul.ts

@@ -5,18 +5,18 @@ import { matmul } from "../backend/webgpu/operators/matmul.ts";
 const backend = new WebGPUBackend();
 await backend.initialize();
 
-const meta = { m: 2, n: 2, k: 2, alpha: 1 };
+const uniform = { m: 2, n: 2, k: 2, alpha: 1 };
 
 const a = await WebGPUData.from(
   backend,
-  new Float32Array(meta.m * meta.k).fill(2),
+  new Float32Array(uniform.m * uniform.k).fill(2),
 );
 const b = await WebGPUData.from(
   backend,
-  new Float32Array(meta.n * meta.k).fill(2),
+  new Float32Array(uniform.n * uniform.k).fill(2),
 );
-const c = new WebGPUData(backend, "f32", meta.m * meta.n);
+const c = new WebGPUData(backend, "f32", uniform.m * uniform.n);
 
-await matmul(backend, a, b, c, meta);
+await matmul(backend, a, b, c, uniform);
 
 console.log(await c.get());

examples/webgpu_reduce.ts

@@ -0,0 +1,19 @@
+// Todo!
+
+// import { WebGPUBackend } from "../backend/webgpu/backend.ts";
+// import { WebGPUData } from "../backend/webgpu/data.ts";
+// import { reduce } from "../backend/webgpu/operators/reduce.ts";
+
+// const backend = new WebGPUBackend();
+// await backend.initialize();
+
+// // deno-fmt-ignore
+// const a = await WebGPUData.from<"f32">(
+//   backend,
+//   new Float32Array(1238).fill(123),
+// );
+
+// console.log(await reduce("return p + c;")(backend, a));
+
+// // what about https://www.w3.org/TR/WGSL/#atomic-type
+// // Probably a good idea, but then again, no idea how to use it lol