webgpu-computed 0.0.11 → 0.0.12

package/package.json

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

package/src/build.d.ts

@@ -1,3 +1,3 @@
-import { GpuComputed } from './utils/GpuComputed';
+import { GpuComputed, AtomicUint32Array } from './utils/GpuComputed';
 import * as WGSL_Fun from "./utils/WGSL_Fun";
-export { GpuComputed, WGSL_Fun };
+export { GpuComputed, WGSL_Fun, AtomicUint32Array };

package/src/index.js

@@ -1,127 +1,140 @@
-async function q(f, e = !0) {
+async function _(f, e = !0) {
   if (typeof global < "u" && typeof require < "u")
     return require(f);
   {
-    let t = await import(
+    let r = await import(
       /* @vite-ignore */
       f
     );
-    return e && (t = t.default), t;
+    return e && (r = r.default), r;
   }
 }
+function q(f) {
+  if (f instanceof Float32Array) return "f32";
+  if (f instanceof Int32Array) return "i32";
+  if (f instanceof Uint32Array) return "u32";
+  throw new Error(
+    `Unsupported ArrayBufferView type: ${f.constructor.name}`
+  );
+}
 function S(f) {
   return f && f[0].toUpperCase() + f.slice(1);
 }
 function x(f) {
   return O.includes(f);
 }
-function w(f) {
+function A(f) {
   return f && Array.isArray(f) && f.length ? f.every((e) => e && typeof e == "object" && "name" in e && "type" in e && x(e.type)) : !1;
 }
 function $(f) {
-  return f && "layout" in f && w(f.layout);
+  return f && "layout" in f && A(f.layout);
 }
-function _(f, e) {
-  const t = B[e];
-  return (t - f % t) % t;
+function U(f, e) {
+  const r = E[e];
+  return (r - f % r) % r;
 }
-function P(f, e) {
+function z(f, e) {
   return f + (e - f % e) % e;
 }
-function T(f) {
-  const e = Object.keys(f), t = e.map((c) => {
-    const n = f[c];
-    if (Array.isArray(n)) {
-      for (const i of ["vec2", "vec3", "vec4", "mat3x3", "mat4x4"])
-        if (E[i] === n.length) return i;
-      throw new Error(`${c} 不支持的数组长度 ${n.length}`);
+function P(f) {
+  const e = Object.keys(f), r = e.map((l) => {
+    const a = f[l];
+    if (Array.isArray(a)) {
+      for (const o of ["vec2", "vec3", "vec4", "mat3x3", "mat4x4"])
+        if (B[o] === a.length) return o;
+      throw new Error(`${l} 不支持的数组长度 ${a.length}`);
     }
-    if (typeof n == "number") return "f32";
-    throw new Error(`${c} 不支持的类型`);
+    if (typeof a == "number") return "f32";
+    throw new Error(`${l} 不支持的类型`);
   });
-  if (e.length !== t.length) throw new Error("keys 与 types 长度不一致");
-  let a = 0;
-  const s = e.map((c, n) => {
-    const i = t[n];
-    a += _(a, i);
-    const u = {
-      name: c,
-      type: i,
-      offset: a,
-      size: E[i]
+  if (e.length !== r.length) throw new Error("keys 与 types 长度不一致");
+  let u = 0;
+  const s = e.map((l, a) => {
+    const o = r[a];
+    u += U(u, o);
+    const i = {
+      name: l,
+      type: o,
+      offset: u,
+      size: B[o]
     };
-    return a += E[i], u;
-  }), r = Math.max(...t.map((c) => B[c]));
+    return u += B[o], i;
+  }), t = Math.max(...r.map((l) => E[l]));
   return {
-    stride: a + (r - a % r) % r,
+    stride: u + (t - u % t) % t,
     layout: s
   };
 }
-const E = {
+const B = {
   f32: 1,
+  u32: 1,
   vec2: 2,
   vec3: 3,
   vec4: 4,
   mat3x3: 12,
   mat4x4: 16
   // array: Infinity
-}, B = {
+}, E = {
   f32: 1,
+  u32: 1,
   vec2: 2,
   vec3: 4,
   vec4: 4,
   mat3x3: 4,
   mat4x4: 4
   // array: 1
-}, O = ["f32", "vec2", "vec3", "vec4", "mat3x3", "mat4x4"];
-let g = null, b = null;
-class z {
+};
+class D extends Uint32Array {
+}
+const O = ["f32", "u32", "vec2", "vec3", "vec4", "mat3x3", "mat4x4"];
+let h = null, b = null;
+class T {
   template;
   option;
   pipeline;
   device;
   groupLayout;
   code;
-  constructor(e, t) {
-    this.template = e, this.option = t, this.improveTemplateOption(this.template);
+  constructor(e, r) {
+    this.template = e, this.option = r, this.improveTemplateOption(this.template);
   }
   /** 完善模版数据
    * @param template 
    */
   improveTemplateOption(e) {
-    const t = (r) => {
-      let o = 0;
-      r.forEach((c) => {
-        o += _(o, c.type), c.offset = o, c.size = E[c.type], o += c.size;
+    const r = (t) => {
+      let n = 0;
+      t.forEach((l) => {
+        n += U(n, l.type), l.offset = n, l.size = B[l.type], n += l.size;
       });
-    }, a = (r) => {
-      t(r.layout);
-      const o = r.layout[r.layout.length - 1], c = o.offset + o.size;
-      let n = 1;
-      for (const i of r.layout) n = Math.max(n, B[i.type]);
-      r.stride = P(c, n);
+    }, u = (t) => {
+      r(t.layout);
+      const n = t.layout[t.layout.length - 1], l = n.offset + n.size;
+      let a = 1;
+      for (const o of t.layout) a = Math.max(a, E[o.type]);
+      t.stride = z(l, a);
     };
-    Object.keys(e).forEach((r) => {
-      const o = e[r];
-      w(o) ? t(o) : $(o) && a(o);
+    Object.keys(e).forEach((t) => {
+      const n = e[t];
+      A(n) ? r(n) : $(n) ? u(n) : Array.isArray(n) && typeof n[0] == "number" && (e[t] = new Float32Array());
     });
   }
   /** 获取Gpu设备
    * @returns 
    */
   async getDevice() {
-    if (!g || !b) throw new Error("webgpu未初始化或不可用");
-    return { adapter: g, device: b };
+    if (!h || !b) throw new Error("webgpu未初始化或不可用");
+    return { adapter: h, device: b };
   }
   /**
    * 初始化计算管线
    */
   async initPipeline() {
     if (!this.template) throw new Error("初始化计算管线错误，未找到可用数据模版");
-    await z.init();
-    const e = this.template, { device: t } = await this.getDevice(), a = [], s = [], r = [];
-    this.device = t, Object.keys(e).forEach((p, d) => {
-      if (r.push({
+    await T.init();
+    const e = this.template, { device: r } = await this.getDevice(), u = [], s = [], t = [];
+    this.device = r, Object.keys(e).forEach((y, d) => {
+      if (t.push({
         binding: d,
         // 绑定到组里的0号位插槽
         visibility: GPUShaderStage.COMPUTE,
@@ -129,45 +142,47 @@ class z {
         buffer: {
           type: "storage"
         }
-      }), w(e[p])) {
-        const v = e[p], A = v.map((h) => `${h.name}:${h.type === "f32" ? "f32" : h.type + "<f32>"}`).join(","), m = `${S(p)}Struct`;
-        a.push(`struct ${S(p)}Struct {${A}};`), s.push(`@group(0) @binding(${d}) var<storage, read_write> ${p}: ${m};`);
-      } else if ($(e[p])) {
-        const v = e[p], A = v.layout.map((h) => `${h.name}:${h.type === "f32" ? "f32" : h.type + "<f32>"}`).join(","), m = `${S(p)}Struct`;
-        a.push(`struct ${m} {${A}};`), s.push(`@group(0) @binding(${d}) var<storage, read_write> ${p}: array<${m}>;`);
-      } else
-        s.push(`@group(0) @binding(${d}) var<storage, read_write> ${p}: array<f32>;`);
+      }), A(e[y])) {
+        const g = e[y], v = g.map((w) => `${w.name}:${w.type === "f32" ? "f32" : w.type + "<f32>"}`).join(","), m = `${S(y)}Struct`;
+        u.push(`struct ${S(y)}Struct {${v}};`), s.push(`@group(0) @binding(${d}) var<storage, read_write> ${y}: ${m};`);
+      } else if ($(e[y])) {
+        const g = e[y], v = g.layout.map((w) => `${w.name}:${w.type === "f32" ? "f32" : w.type + "<f32>"}`).join(","), m = `${S(y)}Struct`;
+        u.push(`struct ${m} {${v}};`), s.push(`@group(0) @binding(${d}) var<storage, read_write> ${y}: array<${m}>;`);
+      } else if (ArrayBuffer.isView(e[y]) && !(e[y] instanceof DataView)) {
+        const g = e[y], v = q(g);
+        g instanceof D ? s.push(`@group(0) @binding(${d}) var<storage, read_write> ${y}: array<atomic<${v}>>;`) : s.push(`@group(0) @binding(${d}) var<storage, read_write> ${y}: array<${v}>;`);
+      }
     });
     const {
-      beforeCodes: o = [],
-      workgroupSize: c = [32, 1, 1],
-      globalInvocationIdName: n = "grid",
-      workgroupIndexName: i = "index",
-      code: u = ""
-    } = this.option ?? {}, l = (
+      beforeCodes: n = [],
+      workgroupSize: l = [32, 1, 1],
+      globalInvocationIdName: a = "grid",
+      workgroupIndexName: o = "index",
+      code: i = ""
+    } = this.option ?? {}, c = (
       /*wgsl*/
       `
-            ${a.join("")}
+            ${u.join("")}
 ${s.join("")}
-            ${o.join(" ") ?? ""}
+            ${n.join(" ") ?? ""}
 
-            @compute @workgroup_size(${c.join(",")})
-            fn main(@builtin(global_invocation_id) ${n}: vec3<u32>) {
-                var ${i} = ${n}.x;
-                ${u}
+            @compute @workgroup_size(${l.join(",")})
+            fn main(@builtin(global_invocation_id) ${a}: vec3<u32>) {
+                var ${o} = ${a}.x;
+                ${i}
             }
         `
     );
-    this.code = l;
-    const y = t.createBindGroupLayout({
-      entries: r
+    this.code = c;
+    const p = r.createBindGroupLayout({
+      entries: t
     });
-    this.groupLayout = y, this.pipeline = t.createComputePipeline({
-      layout: t.createPipelineLayout({
-        bindGroupLayouts: [y]
+    this.groupLayout = p, this.pipeline = r.createComputePipeline({
+      layout: r.createPipelineLayout({
+        bindGroupLayouts: [p]
       }),
       compute: {
-        module: t.createShaderModule({ code: l, label: "" }),
+        module: r.createShaderModule({ code: c, label: "" }),
         entryPoint: "main"
       }
     });
@@ -178,44 +193,46 @@ ${s.join("")}
   createBindGroup(e) {
     if (!this.template) throw new Error("创建buffer组错误，未找到可用数据模版");
     if (!this.device) throw new Error("创建buffer组错误，未找到可用的gpu设备，请确保初始化完计算管线");
-    const t = this.device, a = this.template, s = [];
-    function r(n, i, u = 0, l) {
-      if (!l) {
-        const y = i[i.length - 1], p = y.offset + y.size;
+    const r = this.device, u = this.template, s = [];
+    function t(a, o, i = 0, c) {
+      if (!c) {
+        const p = o[o.length - 1], y = p.offset + p.size;
         let d = 1;
-        for (const v of i) d = Math.max(d, B[v.type]);
-        l = l ?? new Array(P(p, d)).fill(0);
+        for (const g of o) d = Math.max(d, E[g.type]);
+        c = c ?? new Array(z(y, d)).fill(0);
       }
-      return i.forEach((y) => {
-        let p = n[y.name];
-        Array.isArray(p) || (p = [p]);
-        for (let d = 0; d < y.size; d++)
-          l[u + y.offset + d] = Number(p[d] ?? 0);
-      }), l;
+      return o.forEach((p) => {
+        let y = a[p.name];
+        Array.isArray(y) || (y = [y]);
+        for (let d = 0; d < p.size; d++)
+          c[i + p.offset + d] = Number(y[d] ?? 0);
+      }), c;
     }
-    function o(n, i) {
-      const u = new Array(i.stride * n.length).fill(0);
-      return n.forEach((l, y) => {
-        const p = y * i.stride;
-        r(l, i.layout, p, u);
-      }), u;
+    function n(a, o) {
+      const i = new Array(o.stride * a.length).fill(0);
+      return a.forEach((c, p) => {
+        const y = p * o.stride;
+        t(c, o.layout, y, i);
+      }), i;
     }
-    return Object.keys(a).forEach((n) => {
-      if (!(n in e)) throw new Error(`传入的数据中，不存在${n}字段`);
-      const i = a[n], u = e[n];
-      let l = [];
-      w(i) ? l = r(u, i) : $(i) ? l = o(u, i) : Array.isArray(u) && (l = u);
-      const y = new Float32Array(l), p = t.createBuffer({
-        size: y.byteLength,
+    return Object.keys(u).forEach((a) => {
+      if (!(a in e)) throw new Error(`传入的数据中，不存在${a}字段`);
+      const o = u[a], i = e[a];
+      let c = [];
+      A(o) ? c = t(i, o) : $(o) ? c = n(i, o) : Array.isArray(i) && (c = i);
+      let p = null;
+      if (o instanceof Float32Array ? p = new Float32Array(c) : o instanceof Uint32Array ? p = new Uint32Array(c) : o instanceof Int32Array ? p = new Int32Array(c) : p = new Float32Array(c), !p) throw new Error("不支持的数组类型" + o);
+      const y = r.createBuffer({
+        size: p.byteLength,
         usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
       });
-      t.queue.writeBuffer(p, 0, y), s.push({ name: n, buffer: p });
+      r.queue.writeBuffer(y, 0, p), s.push({ name: a, buffer: y });
     }), {
-      group: t.createBindGroup({
+      group: r.createBindGroup({
         layout: this.groupLayout,
-        entries: s.map((n, i) => ({
-          binding: i,
-          resource: { buffer: n.buffer }
+        entries: s.map((a, o) => ({
+          binding: o,
+          resource: { buffer: a.buffer }
         }))
       }),
       buffers: s
@@ -225,24 +242,24 @@ ${s.join("")}
    * @param array 
    * @param key 
    */
-  dataMap(e, t) {
-    if (!(t in this.template)) throw new Error("未找到数据字段：" + t);
-    if ($(this.template[t])) {
-      const a = this.template[t], s = e.length / a.stride, r = [];
-      for (let o = 0; o < s; o++) {
-        const c = o * a.stride, n = {};
-        a.layout.forEach((i) => {
-          const u = e.slice(c + i.offset, c + i.offset + i.size);
-          n[i.name] = u.length === 1 ? u[0] : u;
-        }), r.push(n);
+  dataMap(e, r) {
+    if (!(r in this.template)) throw new Error("未找到数据字段：" + r);
+    if ($(this.template[r])) {
+      const u = this.template[r], s = e.length / u.stride, t = [];
+      for (let n = 0; n < s; n++) {
+        const l = n * u.stride, a = {};
+        u.layout.forEach((o) => {
+          const i = e.slice(l + o.offset, l + o.offset + o.size);
+          a[o.name] = i.length === 1 ? i[0] : i;
+        }), t.push(a);
       }
-      return r;
+      return t;
     }
-    if (w(this.template[t])) {
-      const a = this.template[t], s = {};
-      return a.forEach((r) => {
-        const o = e.slice(r.offset, r.offset + r.size);
-        s[r.name] = o.length === 1 ? o[0] : o;
+    if (A(this.template[r])) {
+      const u = this.template[r], s = {};
+      return u.forEach((t) => {
+        const n = e.slice(t.offset, t.offset + t.size);
+        s[t.name] = n.length === 1 ? n[0] : n;
       }), s;
     }
     return e;
@@ -253,24 +270,25 @@ ${s.join("")}
    * @param synchronize 需要同步的数据字段
    * @returns 
    */
-  async computed(e, t, a) {
+  async computed(e, r, u) {
     if (!this.pipeline) throw new Error("未找到可用计算管线，请确保计算管线已经创建成功");
-    const s = this.device, r = this.pipeline, o = s.createCommandEncoder(), c = o.beginComputePass();
-    c.setPipeline(r), c.setBindGroup(0, e.group), c.dispatchWorkgroups(t[0], t[1], t[2]), c.end();
-    const n = e.buffers?.map((u) => {
-      if (a?.includes(u.name)) {
-        const l = s.createBuffer({
-          size: u.buffer.size,
+    const s = this.device, t = this.pipeline, n = s.createCommandEncoder(), l = n.beginComputePass();
+    l.setPipeline(t), l.setBindGroup(0, e.group), l.dispatchWorkgroups(r[0], r[1], r[2]), l.end();
+    const a = e.buffers?.map((i) => {
+      if (u?.includes(i.name)) {
+        const c = s.createBuffer({
+          size: i.buffer.size,
           usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
         });
-        return o.copyBufferToBuffer(u.buffer, 0, l, 0, l.size), { buffer: l, name: u.name };
+        return n.copyBufferToBuffer(i.buffer, 0, c, 0, c.size), { buffer: c, name: i.name };
       }
-    }).filter((u) => !!u);
-    return s.queue.submit([o.finish()]), await s.queue.onSubmittedWorkDone(), await Promise.all(
-      n.map(async (u) => {
-        await u.buffer.mapAsync(GPUMapMode.READ);
-        const l = u.buffer.getMappedRange();
-        return [...new Float32Array(l)];
+    }).filter((i) => !!i);
+    return s.queue.submit([n.finish()]), await s.queue.onSubmittedWorkDone(), await Promise.all(
+      a.map(async (i) => {
+        await i.buffer.mapAsync(GPUMapMode.READ);
+        const c = i.buffer.getMappedRange();
+        let p = null;
+        return this.template[i.name] instanceof Float32Array ? p = new Float32Array(c) : this.template[i.name] instanceof Uint32Array ? p = new Uint32Array(c) : this.template[i.name] instanceof Int32Array ? p = new Int32Array(c) : p = new Float32Array(c), [...p];
       })
     );
   }
@@ -278,14 +296,14 @@ ${s.join("")}
    * @returns 
    */
   static async init() {
-    if (!(g && b)) {
+    if (!(h && b)) {
       if (typeof globalThis < "u" && typeof window > "u") {
-        const { create: e, globals: t } = await q("webgpu", !1);
-        Object.assign(globalThis, t), globalThis.navigator || (globalThis.navigator = {}), Object.assign(globalThis.navigator, { gpu: e([]) });
+        const { create: e, globals: r } = await _("webgpu", !1);
+        Object.assign(globalThis, r), globalThis.navigator || (globalThis.navigator = {}), Object.assign(globalThis.navigator, { gpu: e([]) });
       }
       if (!navigator.gpu) throw new Error("该环境不支持webgpu");
-      if (g || (g = await navigator.gpu.requestAdapter({})), !g) throw new Error("获取适配器失败");
-      if (b = await g.requestDevice(), !g) throw new Error("获取设备失败");
+      if (h || (h = await navigator.gpu.requestAdapter({})), !h) throw new Error("获取适配器失败");
+      if (b = await h.requestDevice(), !h) throw new Error("获取设备失败");
     }
   }
   /** 注销gpu设备
@@ -297,19 +315,20 @@ ${s.join("")}
    * @param data 
    */
   static buildBufferTypeByData(e) {
-    return Object.keys(e).reduce((a, s) => {
-      const r = e[s];
-      if (Array.isArray(r))
-        if (typeof r[0] == "number") a[s] = [];
-        else if (typeof r[0] == "object" || r.length) {
-          const o = T(r[0]);
-          a[s] = o;
+    return Object.keys(e).reduce((u, s) => {
+      let t = e[s];
+      if (Array.isArray(t) && typeof t[0] == "number" && (t = new Float32Array()), Array.isArray(t))
+        if (typeof t[0] == "object" || t.length) {
+          const n = P(t[0]);
+          u[s] = n;
         } else console.log(`字段：${s}， 不支持该值对应数据类型或数组为空`);
-      else if (typeof r == "object") {
-        const o = T(r);
-        a[s] = o.layout;
+      else if (ArrayBuffer.isView(t) && !(t instanceof DataView))
+        u[s] = t;
+      else if (typeof t == "object") {
+        const n = P(t);
+        u[s] = n.layout;
       } else console.log(`字段：${s}， 不支持的数据类型`);
-      return a;
+      return u;
     }, {});
   }
   /** 通过数据创建
@@ -317,18 +336,18 @@ ${s.join("")}
    * @returns 
    */
   static async fromByData(e) {
-    let { data: t, ...a } = e;
-    const s = this.buildBufferTypeByData(t), r = new z(s, a);
-    return await r.initPipeline(), r;
+    let { data: r, ...u } = e;
+    const s = this.buildBufferTypeByData(r), t = new T(s, u);
+    return await t.initPipeline(), t;
   }
   /** 快捷计算方法
    * @param opt 
    * @returns 
    */
   static async computed(e) {
-    let { data: t, map: a = !1, workgroupCount: s, synchronize: r, onSuccess: o, ...c } = e;
-    const n = await this.fromByData({ data: t, ...c }), i = n.createBindGroup(t), u = await n.computed(i, s, r);
-    return o && o({ gpuComputed: n, group: i, results: u }), a ? u.map((l, y) => n.dataMap(l, r[y])) : u;
+    let { data: r, map: u = !1, workgroupCount: s, synchronize: t, onSuccess: n, ...l } = e;
+    const a = await this.fromByData({ data: r, ...l }), o = a.createBindGroup(r), i = await a.computed(o, s, t);
+    return n && n({ gpuComputed: a, group: o, results: i }), u ? i.map((c, p) => a.dataMap(c, t[p])) : i;
   }
 }
 const j = (
@@ -340,7 +359,7 @@ const j = (
         return v + q.w * t + cross(q.xyz, t);
     }
 `
-), D = (
+), G = (
   /* wgsl */
   `
     fn point_in_obb(
@@ -360,12 +379,13 @@ const j = (
         return all(abs(pLocal) <= halfSize);
     }
 `
-), G = /* @__PURE__ */ Object.freeze(/* @__PURE__ */ Object.defineProperty({
+), L = /* @__PURE__ */ Object.freeze(/* @__PURE__ */ Object.defineProperty({
   __proto__: null,
-  point_in_obb: D,
+  point_in_obb: G,
   quat_rotate: j
 }, Symbol.toStringTag, { value: "Module" }));
 export {
-  z as GpuComputed,
-  G as WGSL_Fun
+  D as AtomicUint32Array,
+  T as GpuComputed,
+  L as WGSL_Fun
 };

package/src/utils/GpuComputed.d.ts

@@ -1,4 +1,4 @@
-type WGSl_TYPE = "f32" | "vec2" | "vec3" | "vec4" | "mat3x3" | "mat4x4";
+type WGSl_TYPE = "f32" | "u32" | "vec2" | "vec3" | "vec4" | "mat3x3" | "mat4x4";
 type GpuComputedOption = {
     workgroupSize?: [number, number, number];
     globalInvocationIdName?: string;
@@ -25,7 +25,9 @@ interface IStructArray {
     stride?: number;
     layout: IStruct;
 }
-type BufferDataType = Record<string, (number[]) | IStruct | IStructArray>;
+type BufferDataType = Record<string, (number[]) | ArrayBufferView | IStruct | IStructArray>;
+export declare class AtomicUint32Array extends Uint32Array {
+}
 export declare class GpuComputed {
     template?: BufferDataType;
     option?: GpuComputedOption;
@@ -81,7 +83,7 @@ export declare class GpuComputed {
     /**
      * @param data
      */
-    static buildBufferTypeByData(data: Record<string, any>): Record<string, IStructArray | number[] | IStruct>;
+    static buildBufferTypeByData(data: Record<string, any>): Record<string, IStructArray | IStruct | ArrayBufferView<ArrayBufferLike>>;
     /** 通过数据创建
      * @param opt
      * @returns