webgpu-computed 0.0.1

package/README.md

@@ -0,0 +1,380 @@
+# webgpu-computed
+
+一个简化的 WebGPU 计算库，封装了繁琐的初始化和缓冲区管理，让开发者专注于 WGSL 着色器逻辑。
+
+## 特性
+
+- 🚀 简化 WebGPU 初始化
+- 📦 自动缓冲区管理和布局计算
+- 🔧 支持复杂数据结构（向量、矩阵）
+- ⚡ 高性能 GPU 计算
+- 📚 内置常用 WGSL 函数
+
+## 安装
+
+```bash
+npm install webgpu-computed
+```
+
+## 快速开始
+
+### 1. 初始化 WebGPU
+
+在使用任何计算功能前，需要先初始化 WebGPU 环境：
+
+```javascript
+import { GpuComputed } from 'webgpu-computed';
+
+// 初始化 WebGPU
+await GpuComputed.init();
+```
+
+### 2. 执行简单计算
+
+以下是一个简单的向量加法示例：
+
+```javascript
+import { GpuComputed } from 'webgpu-computed';
+
+// 准备数据
+const data = {
+  inputA: [1.0, 2.0, 3.0, 4.0],
+  inputB: [0.5, 1.5, 2.5, 3.5],
+  output: new Array(4).fill(0) // 输出缓冲区
+};
+
+// WGSL 计算代码
+const code = `
+  output[index] = inputA[index] + inputB[index];
+`;
+
+// 执行计算
+const results = await GpuComputed.computed({
+  code,
+  data,
+  workgroupCount: [1] // 工作组数量
+});
+
+console.log(results); // [[1.5, 3.5, 5.5, 7.5]]
+```
+
+### 3. 使用复杂数据结构
+
+库支持向量和矩阵类型：
+
+```javascript
+const data = {
+  positions: [
+    { pos: [1.0, 2.0, 3.0], vel: [0.1, 0.2, 0.3] },
+    { pos: [4.0, 5.0, 6.0], vel: [0.4, 0.5, 0.6] }
+  ],
+  output: new Array(2).fill({ pos: [0,0,0], vel: [0,0,0] })
+};
+
+const code = `
+  output[index].pos = positions[index].pos + positions[index].vel;
+  output[index].vel = positions[index].vel * 2.0;
+`;
+
+const results = await GpuComputed.computed({
+  code,
+  data,
+  workgroupCount: [1]
+});
+```
+
+## API 参考
+
+### GpuComputed 类
+
+#### 静态方法
+
+##### `GpuComputed.init()`
+
+初始化 WebGPU 环境。必须在使用其他功能前调用。
+
+**返回值**: `Promise<void>`
+
+**抛出**: 如果浏览器不支持 WebGPU 或获取适配器/设备失败
+
+##### `GpuComputed.computed(options)`
+
+执行 GPU 计算任务。
+
+**参数**:
+
+- `code` (string): WGSL 计算代码
+- `data` (object): 输入/输出数据对象
+- `workgroupCount` (array): 工作组数量 [x, y?, z?]
+- `workgroupSize` (array, 可选): 工作组大小，默认 [32, 1, 1]
+- `globalInvocationIdName` (string, 可选): 全局调用 ID 变量名，默认 "grid"
+- `workgroupIndexName` (string, 可选): 工作组索引变量名，默认 "index"
+- `synchronize` (array, 可选): 需要同步回 CPU 的缓冲区名称数组
+- `beforeCodes` (array, 可选): 计算函数前的 WGSL 代码片段
+- `onSuccess` (function, 可选): 计算成功回调函数
+
+**返回值**: `Promise<Array<Float32Array>>` - 同步缓冲区的数据
+
+### 数据类型
+
+支持以下 WGSL 类型：
+
+- `f32`: 单精度浮点数
+- `vec2`: 二维向量
+- `vec3`: 三维向量
+- `vec4`: 四维向量
+- `mat3x3`: 3x3 矩阵
+- `mat4x4`: 4x4 矩阵
+
+### 内置 WGSL 函数
+
+库提供了一些常用的 WGSL 辅助函数：
+
+#### 四元数旋转
+
+```wgsl
+fn quat_rotate(q: vec4<f32>, v: vec3<f32>) -> vec3<f32>
+```
+
+使用示例：
+
+```javascript
+import { WGSL_Fun } from 'webgpu-computed';
+
+await GpuComputed.computed({
+    code: "",
+    data: {....},
+    beforeCodes:[WGSL_Fun.quat_rotate]
+})
+
+```
+
+#### 点在 OBB 中的检测
+
+```wgsl
+fn point_in_obb(point: vec3<f32>, center: vec3<f32>, halfSize: vec3<f32>, quat: vec4<f32>) -> bool
+```
+
+## 高级用法
+
+### 自定义工作组配置
+
+```javascript
+await GpuComputed.computed({
+  code: '...',
+  data: {...},
+  workgroupCount: [4, 4], // 16 个工作组
+  workgroupSize: [16, 16], // 每个工作组 256 个线程
+});
+```
+
+### 同步数据回 CPU
+
+```javascript
+const results = await GpuComputed.computed({
+  code: '...',
+  data: {...},
+  synchronize: ['output'], // 指定需要同步的缓冲区
+  workgroupCount: [1]
+});
+
+// results 包含同步回的数据
+```
+
+### 回调函数
+
+```javascript
+await GpuComputed.computed({
+  code: '...',
+  data: {...},
+  workgroupCount: [1],
+  onSuccess: ({ code, bufferInfoList, results }) => {
+    console.log('计算完成', results);
+  }
+});
+```
+
+## 浏览器支持
+
+- Chrome 113+
+- Edge 113+
+- Firefox (部分支持)
+- Safari (部分支持)
+
+确保浏览器支持 WebGPU API。
+
+## 示例项目
+
+```js
+    import { GpuComputed } from "webgpu-computed"
+
+    // 1. 初始化 WebGPU
+    console.log('初始化 WebGPU...');
+    await GpuComputed.init();
+    console.log('WebGPU 初始化成功');
+
+    // 2. 简单数组计算示例
+    console.log('\n=== 简单数组计算 ===');
+    const simpleData = {
+        inputA: [1.0, 2.0, 3.0, 4.0],
+        inputB: [0.5, 1.5, 2.5, 3.5],
+        output: new Array(4).fill(0)
+    };
+
+    const simpleCode = `
+        output[index] = inputA[index] + inputB[index];
+    `;
+
+    const simpleResults = await GpuComputed.computed({
+        code: simpleCode,
+        data: simpleData,
+        workgroupCount: [1],
+        synchronize: ['output']
+    });
+
+    console.log('简单计算结果:', simpleResults[0]); // [1.5, 3.5, 5.5, 7.5]
+
+    // 3. 复杂数据结构示例（结构体）
+    console.log('\n=== 复杂数据结构计算 ===');
+    const complexData = {
+        particles: [
+            { position: [1.0, 2.0, 3.0], velocity: [0.1, 0.2, 0.3], mass: 1.0 },
+            { position: [4.0, 5.0, 6.0], velocity: [0.4, 0.5, 0.6], mass: 2.0 }
+        ],
+        output: [
+            { position: [0, 0, 0], velocity: [0, 0, 0], mass: 0 },
+            { position: [0, 0, 0], velocity: [0, 0, 0], mass: 0 }
+        ]
+    };
+
+    const complexCode = `
+        output[index].position = particles[index].position + particles[index].velocity;
+        output[index].velocity = particles[index].velocity * 2.0;
+        output[index].mass = particles[index].mass * 1.5;
+    `;
+
+    const complexResults = await GpuComputed.computed({
+        code: complexCode,
+        data: complexData,
+        workgroupCount: [1],
+        synchronize: ['output']
+    });
+
+    console.log('复杂计算结果:', complexResults[0]);
+
+    // 4. 使用内置 WGSL 函数示例
+    console.log('\n=== 使用内置 WGSL 函数 ===');
+    const wgslFunData = {
+        points: [
+            {
+                x: 1.0, y: 0.0,  z: 0.0
+            },
+            {
+                x: 0.0, y: 1.0,  z: 0.0
+            },
+            {
+                x: -1.0, y: 0.0,  z: 0.0
+            }
+        ],
+        obbCenter: [0.0, 0.0, 0.0],
+        obbHalfSize: [2.0, 2.0, 2.0],
+        obbRotation: [0.0, 0.0, 0.0, 1.0], // 单位四元数，无旋转
+        results: new Array(3).fill(0)
+    };
+
+    const wgslFunCode = `
+        let point = vec3(points[index].x, points[index].y, points[index].z);
+        let center = vec3<f32>(obbCenter[0], obbCenter[1], obbCenter[2]);
+        let halfSize = vec3<f32>(obbHalfSize[0], obbHalfSize[1], obbHalfSize[2]);
+        let quat = vec4<f32>(obbRotation[0], obbRotation[1], obbRotation[2], obbRotation[3]);
+
+        if (point_in_obb(point, center, halfSize, quat)) {
+            results[index] = 1.0;
+        } else {
+            results[index] = 0.0;
+        }
+    `;
+
+    const wgslFunResults = await GpuComputed.computed({
+        code: wgslFunCode,
+        data: wgslFunData,
+        workgroupCount: [1],
+        beforeCodes: [WGSL_Fun.quat_rotate, WGSL_Fun.point_in_obb, /** 可添加自己的函数代码 */],
+        synchronize: ['results']
+    });
+
+    console.log('OBB 检测结果:', wgslFunResults[0]); // [1, 1, 1] 所有点都在 OBB 内
+
+    // 5. 自定义工作组配置示例
+    console.log('\n=== 自定义工作组配置 ===');
+    const largeData = {
+        largeArray: new Array(1024).fill(0).map((_, i) => i * 1.0),
+        output: new Array(1024).fill(0)
+    };
+
+    const largeCode = `
+        output[index] = largeArray[index] * 2.0;
+    `;
+
+    const largeResults = await GpuComputed.computed({
+        code: largeCode,
+        data: largeData,
+        workgroupCount: [32], // 32 个工作组
+        workgroupSize: [32, 1, 1], // 每个工作组 32 个线程，总共 1024 个线程
+        synchronize: ['output']
+    });
+
+    console.log('大数组计算结果 (前10个):', largeResults[0].slice(0, 10));
+
+    // 6. 使用回调函数示例
+    console.log('\n=== 使用回调函数 ===');
+    const callbackData = {
+        values: [10.0, 20.0, 30.0],
+        squares: new Array(3).fill(0)
+    };
+
+    const callbackCode = `
+        squares[index] = values[index] * values[index];
+    `;
+
+    await GpuComputed.computed({
+        code: callbackCode,
+        data: callbackData,
+        workgroupCount: [1],
+        synchronize: ['squares'],
+        onSuccess: ({ code, bufferInfoList, results }) => {
+            console.log('回调函数触发，计算平方结果:', results[0]); // [100, 400, 900]
+        }
+    });
+
+    // 7. 多维工作组示例
+    console.log('\n=== 多维工作组 ===');
+    const matrixData = {
+        matrixA: new Array(16).fill(0).map((_, i) => i * 1.0),
+        matrixB: new Array(16).fill(0).map((_, i) => (i + 1) * 1.0),
+        result: new Array(16).fill(0)
+    };
+
+    const matrixCode = `
+        let x = index % 4u;
+        let y = index / 4u;
+        let idx = y * 4u + x;
+        result[idx] = matrixA[idx] + matrixB[idx];
+    `;
+
+    const matrixResults = await GpuComputed.computed({
+        code: matrixCode,
+        data: matrixData,
+        workgroupCount: [4, 4], // 4x4 工作组网格
+        workgroupSize: [1, 1, 1], // 每个工作组 1 个线程
+        synchronize: ['result']
+    });
+
+    console.log('矩阵计算结果:', matrixResults[0]);
+
+    console.log('\n所有功能示例运行完成！');
+```
+
+## 许可证
+
+ISC License

package/package.json

@@ -0,0 +1,14 @@
+{
+  "name": "webgpu-computed",
+  "version": "0.0.1",
+  "description": "对webgpu的封装，处理了繁琐的前置工作，只关注wgsl本身逻辑",
+  "main": "./src/index.js",
+  "scripts": {
+    "test": "echo \"Error: no test specified\" && exit 1"
+  },
+  "keywords": [
+    "webgpu"
+  ],
+  "author": "夏过初秋",
+  "license": "ISC"
+}

package/src/index.d.ts

@@ -0,0 +1,76 @@
+declare type BufferType = {
+    buffer: number[];
+    stride: number;
+    layout: {
+        name: string;
+        type: ValueType;
+        offset: number;
+        size: number;
+    }[];
+    count: number;
+};
+
+export declare class GpuComputed {
+    constructor();
+    static init(): Promise<void>;
+    getDevice(): Promise<{
+        adapter: GPUAdapter;
+        device: GPUDevice;
+    }>;
+    createPipeline(code: string, buffers: Record<string, BufferType | number[]>): Promise<{
+        pipeline: GPUComputePipeline;
+        group: GPUBindGroup;
+        device: GPUDevice;
+        bufferInfoList: {
+            name: string;
+            buffer: GPUBuffer;
+            float32Array: Float32Array<ArrayBuffer>;
+            groupLayoutItem: GPUBindGroupLayoutEntry;
+            groupItem: {
+                binding: number;
+                resource: {
+                    buffer: GPUBuffer;
+                };
+            };
+        }[];
+    }>;
+    private buildCode;
+    private buildBuffer;
+    computed(opt: {
+        code: string;
+        data: Record<string, any[]>;
+        onSuccess?: (opt: {
+            code: string;
+            bufferInfoList: any[];
+            results: any;
+        }) => void;
+    } & OptionType): Promise<number[][]>;
+    static instance: GpuComputed;
+    static get computed(): (opt: {
+        code: string;
+        data: Record<string, any[]>;
+        onSuccess?: (opt: {
+            code: string;
+            bufferInfoList: any[];
+            results: any;
+        }) => void;
+    } & OptionType) => Promise<number[][]>;
+}
+
+declare type OptionType = {
+    workgroupSize?: [number, number, number];
+    workgroupCount: [number, number?, number?];
+    globalInvocationIdName?: string;
+    workgroupIndexName?: string;
+    synchronize?: string[];
+    beforeCodes?: string[];
+};
+
+declare type ValueType = "f32" | "vec2" | "vec3" | "vec4" | "mat3x3" | "mat4x4";
+
+export declare const WGSL_Fun: Readonly<{
+    quat_rotate: "\n        fn quat_rotate(q: vec4<f32>, v: vec3<f32>) -> vec3<f32> {\n            // q.xyz = vector part, q.w = scalar part\n            let t = 2.0 * cross(q.xyz, v);\n            return v + q.w * t + cross(q.xyz, t);\n        }\n    ";
+    point_in_obb: "\n        fn point_in_obb(\n            point: vec3<f32>,\n            center: vec3<f32>,\n            halfSize: vec3<f32>,\n            quat: vec4<f32>\n        ) -> bool {\n            // 世界空间 → OBB 局部空间\n            let local = point - center;\n\n            // 逆旋转（共轭四元数）\n            let invQuat = vec4<f32>(-quat.xyz, quat.w);\n            let pLocal = quat_rotate(invQuat, local);\n\n            // AABB 判断\n            return all(abs(pLocal) <= halfSize);\n        }\n    ";
+}>;
+
+export { }

package/src/index.js

@@ -0,0 +1,217 @@
+const m = {
+  f32: 1,
+  vec2: 2,
+  vec3: 3,
+  vec4: 4,
+  mat3x3: 12,
+  mat4x4: 16
+}, w = {
+  f32: 1,
+  vec2: 2,
+  vec3: 4,
+  vec4: 4,
+  mat3x3: 4,
+  mat4x4: 4
+};
+function P(A, c) {
+  const o = w[c];
+  return (o - A % o) % o;
+}
+let b = null, v = null;
+class h {
+  constructor() {
+  }
+  static async init() {
+    if (!navigator.gpu) throw new Error("该环境不支持webgpu");
+    if (b = await navigator.gpu.requestAdapter({}), !b) throw new Error("获取适配器失败");
+    if (v = await b.requestDevice(), !b) throw new Error("获取设备失败");
+  }
+  async getDevice() {
+    if (!b || !v) throw new Error("webgpu未初始化或不可用");
+    return { adapter: b, device: v };
+  }
+  async createPipeline(c, o) {
+    const { device: r } = await this.getDevice(), n = Object.keys(o).map((f, s) => {
+      const a = o[f], i = new Float32Array(Array.isArray(a) ? a : a.buffer), e = r.createBuffer({
+        size: i.byteLength,
+        usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
+      });
+      return r.queue.writeBuffer(e, 0, i), {
+        name: f,
+        buffer: e,
+        float32Array: i,
+        groupLayoutItem: {
+          binding: s,
+          // 绑定到组里的0号位插槽
+          visibility: GPUShaderStage.COMPUTE,
+          // 数据在哪些阶段可以使用, 计算着色器、片元着色器、顶点着色器
+          buffer: {
+            type: "storage"
+          }
+        },
+        groupItem: {
+          binding: s,
+          resource: { buffer: e }
+        }
+      };
+    }), l = r.createBindGroupLayout({
+      entries: n.map((f) => f.groupLayoutItem)
+    }), g = r.createBindGroup({
+      layout: l,
+      entries: n.map((f) => f.groupItem)
+    });
+    return { pipeline: r.createComputePipeline({
+      layout: r.createPipelineLayout({
+        bindGroupLayouts: [l]
+      }),
+      compute: {
+        module: r.createShaderModule({ code: c, label: "" }),
+        entryPoint: "main"
+      }
+    }), group: g, device: r, bufferInfoList: n };
+  }
+  buildCode(c, o, r) {
+    const {
+      workgroupSize: n = [32, 1, 1],
+      globalInvocationIdName: l = "grid",
+      workgroupIndexName: g = "index"
+    } = r ?? {}, y = (e) => e && e[0].toUpperCase() + e.slice(1), f = Object.keys(c).map((e) => {
+      const p = c[e];
+      return Array.isArray(p) ? void 0 : `struct ${y(e)}Struct { ${p.layout.map((t) => `${t.name}: ${t.type === "f32" ? "f32" : `${t.type}<f32>`}`).join(",")} };`;
+    }).filter((e) => !!e).join(`
+`), s = Object.keys(c).map((e, p) => {
+      const u = c[e], t = "read_write";
+      return Array.isArray(u) ? `@group(0) @binding(${p}) var<storage, ${t}> ${e}: array<f32>;` : `@group(0) @binding(${p}) var<storage, ${t}> ${e}: array<${y(e)}Struct>;`;
+    }).join(`
+`), a = l;
+    return (
+      /*wgsl*/
+      `
+            ${f}
+${s}
+
+            ${r?.beforeCodes?.join(" ") ?? ""}
+
+            @compute @workgroup_size(${n.join(",")})
+            fn main(@builtin(global_invocation_id) ${a}: vec3<u32>) {
+                var ${g} = ${a}.x;
+                ${o}
+            }
+        `
+    );
+  }
+  buildBuffer(c, o, r) {
+    if (!Array.isArray(c) || c.length === 0)
+      throw new Error("数据必须是非空数组");
+    if (o || (o = Object.keys(c[0])), r || (r = o.map((s) => {
+      const a = c[0][s];
+      if (Array.isArray(a)) {
+        for (const i of ["vec2", "vec3", "vec4", "mat3x3", "mat4x4"])
+          if (m[i] === a.length) return i;
+        throw new Error(`${s} 不支持的数组长度 ${a.length}`);
+      }
+      if (typeof a == "number") return "f32";
+      throw new Error(`${s} 不支持的类型`);
+    })), o.length !== r.length) throw new Error("keys 与 types 长度不一致");
+    let n = 0;
+    const l = o.map((s, a) => {
+      const i = r[a];
+      n += P(n, i);
+      const e = {
+        name: s,
+        type: i,
+        offset: n,
+        size: m[i]
+      };
+      return n += m[i], e;
+    }), g = Math.max(...r.map((s) => w[s])), y = n + (g - n % g) % g, f = new Array(y * c.length).fill(0);
+    return c.forEach((s, a) => {
+      const i = a * y;
+      l.forEach(({ name: e, offset: p, size: u }) => {
+        let t = s[e];
+        Array.isArray(t) || (t = [t]);
+        for (let d = 0; d < u; d++)
+          f[i + p + d] = Number(t[d] ?? 0);
+      });
+    }), {
+      buffer: f,
+      stride: y,
+      layout: l,
+      count: c.length
+    };
+  }
+  async computed(c) {
+    let { code: o, data: r, ...n } = c;
+    const l = Object.keys(r).reduce((u, t) => {
+      const d = r[t];
+      if (typeof d[0] == "number") u[t] = d;
+      else {
+        const $ = this.buildBuffer(d);
+        u[t] = $;
+      }
+      return u;
+    }, {});
+    o = this.buildCode(l, o, n);
+    const { pipeline: g, group: y, device: f, bufferInfoList: s } = await this.createPipeline(o, l), a = f.createCommandEncoder(), i = a.beginComputePass();
+    i.setPipeline(g), i.setBindGroup(0, y), i.dispatchWorkgroups(n.workgroupCount[0], n.workgroupCount[1], n.workgroupCount[2]), i.end();
+    const e = s?.map((u) => {
+      if (n.synchronize?.includes(u.name)) {
+        const t = f.createBuffer({
+          size: u.float32Array.byteLength,
+          usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
+        });
+        return a.copyBufferToBuffer(u.buffer, 0, t, 0, t.size), { buffer: t, name: u.name };
+      }
+    }).filter((u) => !!u);
+    f.queue.submit([a.finish()]), await f.queue.onSubmittedWorkDone();
+    const p = await Promise.all(
+      e.map(async (u) => {
+        await u.buffer.mapAsync(GPUMapMode.READ);
+        const t = u.buffer.getMappedRange();
+        return [...new Float32Array(t)];
+      })
+    );
+    return n?.onSuccess && n.onSuccess({ code: o, bufferInfoList: s, results: p }), p;
+  }
+  static instance;
+  static get computed() {
+    return this.instance || (this.instance = new h()), this.instance.computed.bind(this.instance);
+  }
+}
+const B = Object.freeze({
+  quat_rotate: (
+    /* wgsl */
+    `
+        fn quat_rotate(q: vec4<f32>, v: vec3<f32>) -> vec3<f32> {
+            // q.xyz = vector part, q.w = scalar part
+            let t = 2.0 * cross(q.xyz, v);
+            return v + q.w * t + cross(q.xyz, t);
+        }
+    `
+  ),
+  point_in_obb: (
+    /* wgsl */
+    `
+        fn point_in_obb(
+            point: vec3<f32>,
+            center: vec3<f32>,
+            halfSize: vec3<f32>,
+            quat: vec4<f32>
+        ) -> bool {
+            // 世界空间 → OBB 局部空间
+            let local = point - center;
+
+            // 逆旋转（共轭四元数）
+            let invQuat = vec4<f32>(-quat.xyz, quat.w);
+            let pLocal = quat_rotate(invQuat, local);
+
+            // AABB 判断
+            return all(abs(pLocal) <= halfSize);
+        }
+    `
+  )
+});
+export {
+  h as GpuComputed,
+  B as WGSL_Fun
+};