@genai-fi/nanogpt 0.7.0 → 0.7.2

package/dist/ops/webgpu/utils/reductions.js

@@ -0,0 +1,68 @@
+import { p as l, j as d } from "../../../index-BoWRt-10.js";
+import { g as p } from "../../../webgpu_program-PFzf1hAQ.js";
+import { r as f } from "../../../Reshape-DvudQDvJ.js";
+import { c as x } from "../../../axis_util-BzbKo31C.js";
+function I(e, r, t, s, u) {
+  return `
+           fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+           }
+
+           ${`
+             var<workgroup> xBestValues : array<f32, ${e}>;
+           `}
+    
+           ${p("index")} {
+                let outputIndex = index / ${e};
+                let offset = outputIndex * uniforms.reduceSize;
+                var bestValue = 0.0;
+                let Length = uniforms.reduceSize;
+    
+                for (var k = i32(localId.x); k < Length;
+                    k = k + ${e}) {
+                    var candidate = f32(x[offset + k]);
+                    ${t}
+                    bestValue = bestValue + candidate;
+                }
+                xBestValues[localId.x] = bestValue;
+                workgroupBarrier();
+    
+                var reduceSize = min(u32(Length), ${e}u);
+                for (var currentSize = reduceSize / 2u; reduceSize > 1u;
+                    currentSize = reduceSize / 2u) {
+                    let interval = DIV_CEIL(reduceSize, 2u);
+                    if (localId.x < currentSize) {
+                        let candidate = xBestValues[localId.x + interval];
+                        bestValue = bestValue + candidate;
+                        xBestValues[localId.x] = bestValue;
+                    }
+                    reduceSize = interval;
+                    workgroupBarrier();
+                }
+
+                bestValue = xBestValues[0] ${r === "mean" ? "/ f32(uniforms.reduceSize)" : ""};
+
+                ${s}
+
+                for (var k = i32(localId.x); k < Length;
+                    k = k + ${e}) {
+                    ${u}
+                }
+           }
+         `;
+}
+function V(e, r) {
+  const t = e[0], u = l(r, t.shape), [, n] = x(t.shape, u), a = d(n), i = d(t.shape) / a;
+  return { windowSize: a, inSize: a, batchSize: i, outSize: 1 };
+}
+function $(e, r, t) {
+  const s = [], u = r[0], a = [{ type: "int32", data: [e.inputShape[1]] }], o = t.runWebGPUProgram(e, r, "float32", a);
+  s.push(o);
+  const i = f({ inputs: { x: o }, attrs: { shape: u.shape }, backend: t });
+  return s.forEach((c) => t.disposeData(c.dataId)), i;
+}
+export {
+  V as createReduceInfo,
+  I as createReductionShader,
+  $ as reduce
+};