diff --git a/backend/types.ts b/backend/types.ts index 2afa20d..6b62ffb 100644 --- a/backend/types.ts +++ b/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 = @@ -40,9 +40,3 @@ export interface Backend { export interface BackendRequest { data: Data[]; } - -export type Operator = ( - backend: Backend, - data: DataArray[], - meta: unknown, -) => Promise; diff --git a/backend/webgpu/data.ts b/backend/webgpu/data.ts index 0d6be57..3acb671 100644 --- a/backend/webgpu/data.ts +++ b/backend/webgpu/data.ts @@ -1,26 +1,6 @@ import { Data, DataArray, DataArrayConstructor, DataType } from "../types.ts"; import { WebGPUBackend } from "./backend.ts"; -export interface WebGPUData extends Data { - type: T; - backend: WebGPUBackend; - buffer: GPUBuffer; -} - -export interface WebGPUDataConstructor { - from( - backend: WebGPUBackend, - source: DataArray, - ): Promise>; - - new ( - backend: WebGPUBackend, - type: DataType, - length: number, - usage?: number, - ): WebGPUData; -} - export class WebGPUData implements Data { type: T; backend: WebGPUBackend; @@ -32,6 +12,7 @@ export class WebGPUData implements Data { static async from( backend: WebGPUBackend, source: DataArray, + usage?: number, ): Promise> { // deno-fmt-ignore const type = ( @@ -40,7 +21,7 @@ export class WebGPUData implements Data { : 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; } diff --git a/backend/webgpu/operators/matmul.ts b/backend/webgpu/operators/matmul.ts index a4346c5..a77402d 100644 --- a/backend/webgpu/operators/matmul.ts +++ b/backend/webgpu/operators/matmul.ts @@ -13,12 +13,15 @@ export async function matmul( ) { 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], }); } diff --git a/backend/webgpu/operators/pad.ts b/backend/webgpu/operators/pad.ts index d3a5054..2d9a275 100644 --- a/backend/webgpu/operators/pad.ts +++ b/backend/webgpu/operators/pad.ts @@ -12,12 +12,15 @@ export async function pad( ) { 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], }); } diff --git a/backend/webgpu/operators/transpose.ts b/backend/webgpu/operators/transpose.ts index 9dfe0a8..a341ad1 100644 --- a/backend/webgpu/operators/transpose.ts +++ b/backend/webgpu/operators/transpose.ts @@ -12,12 +12,15 @@ export async function transpose( ) { 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], }); } diff --git a/backend/webgpu/shaders/matmul.ts b/backend/webgpu/shaders/matmul.ts index 03e5fb0..b48a337 100644 --- a/backend/webgpu/shaders/matmul.ts +++ b/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 b: Data; [[group(0), binding(2)]] var c: Data; [[group(0), binding(3)]] -var meta: Meta; +var uniform: Uniform; [[stage(compute), workgroup_size(8, 8, 1)]] fn main([[builtin(global_invocation_id)]] global_id: vec3) { - 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; } `; diff --git a/backend/webgpu/shaders/pad.ts b/backend/webgpu/shaders/pad.ts index bafe39c..b129a46 100644 --- a/backend/webgpu/shaders/pad.ts +++ b/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 a: Data; [[group(0), binding(1)]] var b: Data; [[group(0), binding(2)]] -var meta: Meta; +var uniform: Uniform; [[stage(compute), workgroup_size(8, 8, 1)]] fn main([[builtin(global_invocation_id)]] global_id: vec3) { - 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]; } `; diff --git a/backend/webgpu/shaders/transpose.ts b/backend/webgpu/shaders/transpose.ts index 930920b..0f12d4c 100644 --- a/backend/webgpu/shaders/transpose.ts +++ b/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 a: Data; [[group(0), binding(1)]] var b: Data; [[group(0), binding(2)]] -var meta: Meta; +var uniform: Uniform; [[stage(compute), workgroup_size(8, 8, 1)]] fn main([[builtin(global_invocation_id)]] global_id: vec3) { - 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]; } `; diff --git a/backend/webgpu/types.ts b/backend/webgpu/types.ts index 42da5c2..0cb979c 100644 --- a/backend/webgpu/types.ts +++ b/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 data: WebGPUData[]; workgroups: Workgroups; } - -export type WebGPUOperator = Operator & { - backend: WebGPUBackend; - data: WebGPUData[]; -}; diff --git a/examples/webgpu_matmul.ts b/examples/webgpu_matmul.ts index dacc4e6..fc1f3d0 100644 --- a/examples/webgpu_matmul.ts +++ b/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()); diff --git a/examples/webgpu_reduce.ts b/examples/webgpu_reduce.ts new file mode 100644 index 0000000..425abba --- /dev/null +++ b/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