webgpu-computed 0.0.17 → 0.0.19

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "webgpu-computed",
-  "version": "0.0.17",
+  "version": "0.0.19",
   "description": "对webgpu的封装，处理了繁琐的前置工作，只关注wgsl本身逻辑",
   "main": "./src/index.js",
   "scripts": {

package/src/index.js

@@ -40,28 +40,28 @@ function U(f) {
   const e = Object.keys(f), r = e.map((c) => {
     const l = f[c];
     if (Array.isArray(l)) {
-      for (const d of ["vec2", "vec3", "vec4", "mat3x3", "mat4x4"])
-        if ($[d] === l.length) return d;
+      for (const y of ["vec2", "vec3", "vec4", "mat3x3", "mat4x4"])
+        if ($[y] === l.length) return y;
       throw new Error(`${c} 不支持的数组长度 ${l.length}`);
     }
     if (typeof l == "number") return "f32";
     throw new Error(`${c} 不支持的类型`);
   });
   if (e.length !== r.length) throw new Error("keys 与 types 长度不一致");
-  let s = 0;
+  let i = 0;
   const u = e.map((c, l) => {
-    const d = r[l];
-    s += _(s, d);
+    const y = r[l];
+    i += _(i, y);
     const a = {
       name: c,
-      type: d,
-      offset: s,
-      size: $[d]
+      type: y,
+      offset: i,
+      size: $[y]
     };
-    return s += $[d], a;
+    return i += $[y], a;
   }), t = Math.max(...r.map((c) => E[c]));
   return {
-    stride: s + (t - s % t) % t,
+    stride: i + (t - i % t) % t,
     layout: u
   };
 }
@@ -107,16 +107,16 @@ class T {
       t.forEach((c) => {
         n += _(n, c.type), c.offset = n, c.size = $[c.type], n += c.size;
       });
-    }, s = (t) => {
+    }, i = (t) => {
       r(t.layout);
       const n = t.layout[t.layout.length - 1], c = n.offset + n.size;
       let l = 1;
-      for (const d of t.layout) l = Math.max(l, E[d.type]);
+      for (const y of t.layout) l = Math.max(l, E[y.type]);
       t.stride = P(c, l);
     };
     Object.keys(e).forEach((t) => {
       const n = e[t];
-      m(n) ? r(n) : B(n) ? s(n) : Array.isArray(n) && typeof n[0] == "number" && (e[t] = new Float32Array());
+      m(n) ? r(n) : B(n) ? i(n) : Array.isArray(n) && typeof n[0] == "number" && (e[t] = new Float32Array());
     });
   }
   /** 获取Gpu设备
@@ -132,43 +132,43 @@ class T {
   async initPipeline() {
     if (!this.template) throw new Error("初始化计算管线错误，未找到可用数据模版");
     await T.init();
-    const e = this.template, { device: r } = await this.getDevice(), s = [], u = [], t = [];
-    this.device = r, Object.keys(e).forEach((i, y) => {
+    const e = this.template, { device: r } = await this.getDevice(), i = [], u = [], t = [];
+    this.device = r, Object.keys(e).forEach((s, d) => {
       if (t.push({
-        binding: y,
+        binding: d,
         // 绑定到组里的0号位插槽
         visibility: GPUShaderStage.COMPUTE,
         // 数据在哪些阶段可以使用, 计算着色器、片元着色器、顶点着色器
         buffer: {
           type: "storage"
         }
-      }), m(e[i])) {
-        const g = e[i], h = g.map((b) => `${b.name}:${b.type === "f32" ? "f32" : b.type + "<f32>"}`).join(","), v = `${S(i)}Struct`;
-        s.push(`struct ${S(i)}Struct {${h}};`), u.push(`@group(0) @binding(${y}) var<storage, read_write> ${i}: ${v};`);
-      } else if (B(e[i])) {
-        const g = e[i], h = g.layout.map((b) => `${b.name}:${b.type === "f32" ? "f32" : b.type + "<f32>"}`).join(","), v = `${S(i)}Struct`;
-        s.push(`struct ${v} {${h}};`), u.push(`@group(0) @binding(${y}) var<storage, read_write> ${i}: array<${v}>;`);
-      } else if (ArrayBuffer.isView(e[i]) && !(e[i] instanceof DataView)) {
-        const g = e[i], h = q(g);
-        g instanceof D ? u.push(`@group(0) @binding(${y}) var<storage, read_write> ${i}: array<atomic<${h}>>;`) : u.push(`@group(0) @binding(${y}) var<storage, read_write> ${i}: array<${h}>;`);
+      }), m(e[s])) {
+        const g = e[s], h = g.map((b) => `${b.name}:${b.type === "f32" ? "f32" : b.type + "<f32>"}`).join(","), v = `${S(s)}Struct`;
+        i.push(`struct ${S(s)}Struct {${h}};`), u.push(`@group(0) @binding(${d}) var<storage, read_write> ${s}: ${v};`);
+      } else if (B(e[s])) {
+        const g = e[s], h = g.layout.map((b) => `${b.name}:${b.type === "f32" ? "f32" : b.type + "<f32>"}`).join(","), v = `${S(s)}Struct`;
+        i.push(`struct ${v} {${h}};`), u.push(`@group(0) @binding(${d}) var<storage, read_write> ${s}: array<${v}>;`);
+      } else if (ArrayBuffer.isView(e[s]) && !(e[s] instanceof DataView)) {
+        const g = e[s], h = q(g);
+        g instanceof D ? u.push(`@group(0) @binding(${d}) var<storage, read_write> ${s}: array<atomic<${h}>>;`) : u.push(`@group(0) @binding(${d}) var<storage, read_write> ${s}: array<${h}>;`);
       }
     });
     const {
       beforeCodes: n = [],
       workgroupSize: c = [32, 1, 1],
       globalInvocationIdName: l = "grid",
-      workgroupIndexName: d = "index",
+      workgroupIndexName: y = "index",
       code: a = ""
     } = this.option ?? {}, o = (
       /*wgsl*/
       `
-            ${s.join("")}
+            ${i.join("")}
 ${u.join("")}
             ${n.join(" ") ?? ""}
 
             @compute @workgroup_size(${c.join(",")})
             fn main(@builtin(global_invocation_id) ${l}: vec3<u32>) {
-                var ${d} = ${l}.x;
+                var ${y} = ${l}.x;
                 ${a}
             }
         `
@@ -196,28 +196,28 @@ ${u.join("")}
   createBindGroup(e, r) {
     if (!this.template) throw new Error("创建buffer组错误，未找到可用数据模版");
     if (!this.device) throw new Error("创建buffer组错误，未找到可用的gpu设备，请确保初始化完计算管线");
-    const s = this.device, u = this.template, t = [], n = r?.buffers?.reduce((a, o) => (a.set(o.name, o.buffer), a), /* @__PURE__ */ new Map());
-    function c(a, o, p = 0, i) {
+    const i = this.device, u = this.template, t = [], n = r?.buffers?.reduce((a, o) => (a.set(o.name, o.buffer), a), /* @__PURE__ */ new Map());
+    function c(a, o, p = 0, s) {
       if (ArrayBuffer.isView(a) || Array.isArray(a)) return a;
-      if (!i) {
-        const y = o[o.length - 1], g = y.offset + y.size;
+      if (!s) {
+        const d = o[o.length - 1], g = d.offset + d.size;
         let h = 1;
         for (const v of o) h = Math.max(h, E[v.type]);
-        i = i ?? new Float32Array(P(g, h)).fill(0);
+        s = s ?? new Float32Array(P(g, h)).fill(0);
       }
-      return o.forEach((y) => {
-        let g = a[y.name];
+      return o.forEach((d) => {
+        let g = a[d.name];
         Array.isArray(g) || (g = [g]);
-        for (let h = 0; h < y.size; h++)
-          i[p + y.offset + h] = Number(g[h] ?? 0);
-      }), i;
+        for (let h = 0; h < d.size; h++)
+          s[p + d.offset + h] = Number(g[h] ?? 0);
+      }), s;
     }
     function l(a, o) {
       if (ArrayBuffer.isView(a) || typeof a[0] == "number") return a;
       const p = new Float32Array(o.stride * a.length).fill(0);
-      return a.forEach((i, y) => {
-        const g = y * o.stride;
-        c(i, o.layout, g, p);
+      return a.forEach((s, d) => {
+        const g = d * o.stride;
+        c(s, o.layout, g, p);
       }), p;
     }
     return Object.keys(u).forEach((a) => {
@@ -227,17 +227,17 @@ ${u.join("")}
         throw new Error(`传入的数据中，不存在${a}字段`);
       }
       const o = u[a], p = e[a];
-      let i = [];
-      m(o) ? i = c(p, o) : B(o) ? i = l(p, o) : (Array.isArray(p) || ArrayBuffer.isView(p)) && (i = p);
-      let y = null;
-      if (o instanceof Float32Array ? y = new Float32Array(i) : o instanceof Uint32Array ? y = new Uint32Array(i) : o instanceof Int32Array ? y = new Int32Array(i) : y = new Float32Array(i), !y) throw new Error("不支持的数组类型" + o);
-      const g = s.createBuffer({
-        size: y.byteLength,
+      let s = [];
+      m(o) ? s = c(p, o) : B(o) ? s = l(p, o) : (Array.isArray(p) || ArrayBuffer.isView(p)) && (s = p);
+      let d = null;
+      if (o instanceof Float32Array ? d = new Float32Array(s) : o instanceof Uint32Array ? d = new Uint32Array(s) : o instanceof Int32Array ? d = new Int32Array(s) : d = new Float32Array(s), !d) throw new Error("不支持的数组类型" + o);
+      const g = i.createBuffer({
+        size: d.byteLength,
         usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
       });
-      s.queue.writeBuffer(g, 0, y), t.push({ name: a, buffer: g });
+      i.queue.writeBuffer(g, 0, d), t.push({ name: a, buffer: g });
     }), {
-      group: s.createBindGroup({
+      group: i.createBindGroup({
         layout: this.groupLayout,
         entries: t.map((a, o) => ({
           binding: o,
@@ -254,21 +254,21 @@ ${u.join("")}
   dataMap(e, r) {
     if (!(r in this.template)) throw new Error("未找到数据字段：" + r);
     if (B(this.template[r])) {
-      const s = this.template[r], u = e.length / s.stride, t = [];
+      const i = this.template[r], u = e.length / i.stride, t = [];
       for (let n = 0; n < u; n++) {
-        const c = n * s.stride, l = {};
-        s.layout.forEach((d) => {
-          const a = e.slice(c + d.offset, c + d.offset + d.size);
-          l[d.name] = a.length === 1 ? a[0] : a;
+        const c = n * i.stride, l = {};
+        i.layout.forEach((y) => {
+          const a = e.slice(c + y.offset, c + y.offset + y.size);
+          l[y.name] = a.length === 1 ? a[0] : [...a];
         }), t.push(l);
       }
       return t;
     }
     if (m(this.template[r])) {
-      const s = this.template[r], u = {};
-      return s.forEach((t) => {
+      const i = this.template[r], u = {};
+      return i.forEach((t) => {
         const n = e.slice(t.offset, t.offset + t.size);
-        u[t.name] = n.length === 1 ? n[0] : n;
+        u[t.name] = n.length === 1 ? n[0] : [...n];
       }), u;
     }
     return e;
@@ -279,12 +279,12 @@ ${u.join("")}
    * @param synchronize 需要同步的数据字段
    * @returns 
    */
-  async computed(e, r, s = []) {
+  async computed(e, r, i = []) {
     if (!this.pipeline) throw new Error("未找到可用计算管线，请确保计算管线已经创建成功");
     const u = this.device, t = this.pipeline, n = u.createCommandEncoder(), c = n.beginComputePass();
     c.setPipeline(t), c.setBindGroup(0, e.group), c.dispatchWorkgroups(r[0], r[1], r[2]), c.end();
     const l = e.buffers?.map((o) => {
-      if (s?.includes(o.name)) {
+      if (i?.includes(o.name)) {
         const p = u.createBuffer({
           size: o.buffer.size,
           usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
@@ -293,30 +293,28 @@ ${u.join("")}
       }
     }).filter((o) => !!o);
     u.queue.submit([n.finish()]), await u.queue.onSubmittedWorkDone();
-    const d = /* @__PURE__ */ new Map();
+    const y = /* @__PURE__ */ new Map();
     return await Promise.all(
       l.map(async (o) => {
         await o.buffer.mapAsync(GPUMapMode.READ);
         const p = o.buffer.getMappedRange();
-        let i = null;
-        this.template[o.name] instanceof Float32Array ? i = new Float32Array(p) : this.template[o.name] instanceof Uint32Array ? i = new Uint32Array(p) : this.template[o.name] instanceof Int32Array ? i = new Int32Array(p) : i = new Float32Array(p);
-        const y = [...i];
-        d.set(o.name, y);
+        let s = null;
+        this.template[o.name] instanceof Float32Array ? s = new Float32Array(p) : this.template[o.name] instanceof Uint32Array ? s = new Uint32Array(p) : this.template[o.name] instanceof Int32Array ? s = new Int32Array(p) : s = new Float32Array(p), y.set(o.name, s);
       })
-    ), s.map((o) => d.get(o));
+    ), i.map((o) => y.get(o));
   }
   /** 初始化gpu设备
    * @returns 
    */
-  static async init() {
+  static async init(e = {}) {
     if (!(w && A)) {
       if (typeof globalThis < "u" && typeof window > "u") {
-        const { create: e, globals: r } = await z("webgpu", !1);
-        Object.assign(globalThis, r), globalThis.navigator || (globalThis.navigator = {}), Object.assign(globalThis.navigator, { gpu: e([]) });
+        const { create: r, globals: i } = await z("webgpu", !1);
+        Object.assign(globalThis, i), globalThis.navigator || (globalThis.navigator = {}), Object.assign(globalThis.navigator, { gpu: r([]) });
       }
       if (!navigator.gpu) throw new Error("该环境不支持webgpu");
-      if (w || (w = await navigator.gpu.requestAdapter({})), !w) throw new Error("获取适配器失败");
-      if (A = await w.requestDevice(), !w) throw new Error("获取设备失败");
+      if (w || (w = await navigator.gpu.requestAdapter(e)), !w) throw new Error("获取适配器失败");
+      if (A = await w.requestDevice(e), !w) throw new Error("获取设备失败");
     }
   }
   /** 注销gpu设备
@@ -328,20 +326,20 @@ ${u.join("")}
    * @param data 
    */
   static buildBufferTypeByData(e) {
-    return Object.keys(e).reduce((s, u) => {
+    return Object.keys(e).reduce((i, u) => {
       let t = e[u];
       if (Array.isArray(t) && typeof t[0] == "number" && (t = new Float32Array()), Array.isArray(t))
         if (typeof t[0] == "object" || t.length) {
           const n = U(t[0]);
-          s[u] = n;
+          i[u] = n;
         } else console.log(`字段：${u}， 不支持该值对应数据类型或数组为空`);
       else if (ArrayBuffer.isView(t) && !(t instanceof DataView))
-        s[u] = t;
+        i[u] = t;
       else if (typeof t == "object") {
         const n = U(t);
-        s[u] = n.layout;
+        i[u] = n.layout;
       } else console.log(`字段：${u}， 不支持的数据类型`);
-      return s;
+      return i;
     }, {});
   }
   /** 通过数据创建
@@ -349,8 +347,8 @@ ${u.join("")}
    * @returns 
    */
   static async fromByData(e) {
-    let { data: r, ...s } = e;
-    const u = this.buildBufferTypeByData(r), t = new T(u, s);
+    let { data: r, ...i } = e;
+    const u = this.buildBufferTypeByData(r), t = new T(u, i);
     return await t.initPipeline(), t;
   }
   /** 快捷计算方法
@@ -358,9 +356,9 @@ ${u.join("")}
    * @returns 
    */
   static async computed(e) {
-    let { data: r, map: s = !1, workgroupCount: u, synchronize: t, onSuccess: n, ...c } = e;
-    const l = await this.fromByData({ data: r, ...c }), d = l.createBindGroup(r), a = await l.computed(d, u, t);
-    return n && n({ gpuComputed: l, group: d, results: a }), s ? a.map((o, p) => l.dataMap(o, t[p])) : a;
+    let { data: r, map: i = !1, workgroupCount: u, synchronize: t, onSuccess: n, ...c } = e;
+    const l = await this.fromByData({ data: r, ...c }), y = l.createBindGroup(r), a = await l.computed(y, u, t);
+    return n && n({ gpuComputed: l, group: y, results: a }), i ? a.map((o, p) => l.dataMap(o, t[p])) : a;
   }
 }
 const M = (

package/src/utils/GpuComputed.d.ts

@@ -26,6 +26,15 @@ interface IStructArray {
     layout: IStruct;
 }
 type BufferDataType = Record<string, (number[]) | ArrayBufferView | IStruct | IStructArray>;
+type InitOption = {
+    featureLevel?: string;
+    powerPreference?: GPUPowerPreference;
+    forceFallbackAdapter?: boolean;
+    xrCompatible?: boolean;
+    requiredFeatures?: Iterable<GPUFeatureName>;
+    requiredLimits?: Record<string, GPUSize64 | undefined>;
+    defaultQueue?: GPUQueueDescriptor;
+};
 export declare class AtomicUint32Array extends Uint32Array {
 }
 export declare class GpuComputed {
@@ -74,18 +83,18 @@ export declare class GpuComputed {
      * @param array
      * @param key
      */
-    dataMap(array: number[], key: string): number[] | Record<string, number | number[]> | Record<string, number | number[]>[];
+    dataMap(array: Float32Array | Uint32Array | Int32Array, key: string): Float32Array<ArrayBufferLike> | Int32Array<ArrayBufferLike> | Uint32Array<ArrayBufferLike> | Record<string, number | number[]> | Record<string, number | number[]>[];
     /** 开始计算
      * @param group 数据组
      * @param workgroupCount 工作组大小
      * @param synchronize 需要同步的数据字段
      * @returns
      */
-    computed(group: BufferGroup, workgroupCount: [number, number?, number?], synchronize?: string[]): Promise<number[][]>;
+    computed(group: BufferGroup, workgroupCount: [number, number?, number?], synchronize?: string[]): Promise<(Float32Array<ArrayBufferLike> | Int32Array<ArrayBufferLike> | Uint32Array<ArrayBufferLike>)[]>;
     /** 初始化gpu设备
      * @returns
      */
-    static init(): Promise<void>;
+    static init(opt?: InitOption): Promise<void>;
     /** 注销gpu设备
      */
     static destroy(): void;
@@ -112,8 +121,8 @@ export declare class GpuComputed {
         onSuccess?: (opt: {
             gpuComputed: GpuComputed;
             group: BufferGroup;
-            results: number[][];
+            results: (Float32Array | Uint32Array | Int32Array)[];
         }) => void;
-    } & GpuComputedOption): Promise<(number[] | Record<string, number | number[]> | Record<string, number | number[]>[])[]>;
+    } & GpuComputedOption): Promise<(Float32Array<ArrayBufferLike> | Int32Array<ArrayBufferLike> | Uint32Array<ArrayBufferLike> | Record<string, number | number[]> | Record<string, number | number[]>[])[]>;
 }
 export {};