@genai-fi/nanogpt 0.9.0 → 0.10.0

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

@@ -0,0 +1,79 @@
+import { f as s, c as n } from "../../../webgpu_util-BBCnKm2X.js";
+import { g as p } from "../../../binary_op_util-pKXltfxI.js";
+import { B as g } from "../../../binary_op_util-pKXltfxI.js";
+import { j as u } from "../../../index-ZyQhjEPo.js";
+import { e as l } from "../../../webgpu_program-Cigz-7RF.js";
+import { l as r } from "../../../tensor-DdQUJZlz.js";
+class y {
+  dispatch;
+  dispatchLayout;
+  outputComponent;
+  op;
+  outputShape;
+  shaderKey;
+  size = !0;
+  variableNames = ["A", "B"];
+  workgroupSize;
+  variableComponents;
+  constructor(e, t, o) {
+    this.outputShape = u(t, o), this.dispatchLayout = s(this.outputShape), this.op = e;
+    const i = t.length > 0 && t[t.length - 1] % 4 === 0, a = o.length > 0 && o[o.length - 1] % 4 === 0;
+    if (i && a)
+      this.outputComponent = 4, this.variableComponents = [4, 4];
+    else throw i && (r(o) || o[o.length - 1] === 1) || a && (r(t) || t[t.length - 1] === 1) ? new Error("Cannot broadcast 16-bit float binary ops with mixed vector sizes") : new Error("16-bit float binary ops require inner dimension to be multiple of 4");
+    this.shaderKey = `binary_${e}_${this.variableComponents}`, this.workgroupSize = [128, 1, 1], this.dispatch = n(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+      this.outputComponent,
+      1,
+      1
+    ]);
+  }
+  getUserCode() {
+    const e = this.outputComponent === 4 ? "vec4<f32>" : "f32";
+    return `
+       ${`
+    fn binaryOperation(a : ${e}, b : ${e}) -> ${e} {
+      ${p(this.op, this.outputComponent === 4)}
+    };
+    `}
+       ${l("index")} {
+         if (index < uniforms.size) {
+            let a = A[index];
+            let b = B[index];
+
+            let v4a1 = vec4<f32>(
+                unpack2x16float(u32(a.x)),
+                unpack2x16float(u32(a.y))
+            );
+            let v4a2 = vec4<f32>(
+                unpack2x16float(u32(a.z)),
+                unpack2x16float(u32(a.w))
+            );
+            let v4b1 = vec4<f32>(
+                unpack2x16float(u32(b.x)),
+                unpack2x16float(u32(b.y))
+            );
+            let v4b2 = vec4<f32>(
+                unpack2x16float(u32(b.z)),
+                unpack2x16float(u32(b.w))
+            );
+
+            let v4res1 = binaryOperation(v4a1, v4b1);
+            let v4res2 = binaryOperation(v4a2, v4b2);
+
+            let res = vec4<i32>(
+                i32(pack2x16float(v4res1.xy)),
+                i32(pack2x16float(v4res1.zw)),
+                i32(pack2x16float(v4res2.xy)),
+                i32(pack2x16float(v4res2.zw))
+            );
+
+            result[index] = res;
+         }
+       }
+       `;
+  }
+}
+export {
+  y as BinaryOpProgram,
+  g as BinaryOpType
+};

package/dist/ops/webgpu/utils/deviceInfo.d.ts

@@ -0,0 +1,7 @@
+import { default as WebGPUBackendPatch } from '../../../patches/webgpu_backend';
+export interface DeviceInformation {
+    subgroupsSupported: boolean;
+    subgroupMaxSize: number;
+    variableSubgroups: boolean;
+}
+export default function createDeviceInformation(backend: WebGPUBackendPatch): DeviceInformation;

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

@@ -0,0 +1,11 @@
+function o(u) {
+  const r = u.device.features.has("subgroups");
+  return {
+    subgroupsSupported: r,
+    subgroupMaxSize: u.subgroupMaxSize,
+    variableSubgroups: u.subgroupMinSize !== u.subgroupMaxSize && r
+  };
+}
+export {
+  o as default
+};

package/dist/ops/webgpu/utils/reductions.d.ts

@@ -1,9 +1,37 @@
 import { backend_util, TensorInfo } from '@tensorflow/tfjs-core';
 import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu/dist/webgpu_program';
 import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu';
-export interface ReduceWebGPUProgram extends WebGPUProgram {
-    inputShape: number[];
+import { DeviceInformation } from './deviceInfo';
+export interface ReduceParams {
+    reductionOp: 'mean' | 'sum';
+    elementwise?: boolean;
 }
-export declare function createReductionShader(workgroupSizeX: number, reductionOp: 'mean' | 'sum', inputSnippet: string, reducedSnippet: string, outputSnippet: string): string;
 export declare function createReduceInfo(inputs: TensorInfo[], axis: number | number[]): backend_util.ReduceInfo;
-export declare function reduce(program: ReduceWebGPUProgram, inputs: TensorInfo[], backend: WebGPUBackend): TensorInfo;
+export declare class ReduceProgram implements WebGPUProgram {
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    variableNames: string[];
+    uniforms: string;
+    inputShape: number[];
+    size: boolean;
+    packed: boolean;
+    outputComponent: number;
+    variableComponents?: number[];
+    elementwise: boolean;
+    subgroups: boolean;
+    subgroupBuiltins: boolean;
+    deviceInfo: DeviceInformation;
+    params: ReduceParams;
+    constructor(deviceInfo: DeviceInformation, reduceInfo: backend_util.ReduceInfo, params: ReduceParams, packed: boolean);
+    protected getWriteSnippet(): string;
+    protected getPreprocessSnippet(): string;
+    protected getPostprocessSnippet(): string;
+    protected getReadSnippet(): string;
+    getUserCode(): string;
+}
+export declare function reduce(program: ReduceProgram, inputs: TensorInfo[], backend: WebGPUBackend): TensorInfo;

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

@@ -1,68 +1,259 @@
-import { p as l, j as d } from "../../../index-BzFyqcy-.js";
-import { g as p } from "../../../webgpu_program-DkQJOJSd.js";
-import { r as f } from "../../../Reshape-DUqYftGC.js";
-import { c as x } from "../../../axis_util-TbGYJ208.js";
-function I(e, r, t, s, u) {
+import { e as l } from "../../../index-ZyQhjEPo.js";
+import { e as a } from "../../../webgpu_program-Cigz-7RF.js";
+import { reshape16 as S } from "../../reshape16.js";
+import { f as b } from "../../../webgpu_util-BBCnKm2X.js";
+import { p as f, s as c } from "../../../tensor-DdQUJZlz.js";
+import { c as h } from "../../../axis_util-BvHEw88j.js";
+function d(e, u, t, i) {
+  return e && !u ? `
+            bestValue = subgroupAdd(bestValue);
+        ` : e ? `
+            bestValue = subgroupAdd(bestValue);
+            let lane = localId.x % subgroupSize;
+            if (lane == 0) {
+                bestValues[localId.x / subgroupSize] = bestValue;
+            }
+            workgroupBarrier();
+            let numSubgroups = ${t} / subgroupSize;
+            bestValue = select(${i ? "vec2<f32>(0.0f)" : "0.0f"}, bestValues[lane], lane < numSubgroups);
+            bestValue = subgroupAdd(bestValue);
+        ` : `
+            bestValues[localId.x] = bestValue;
+            workgroupBarrier();
+
+            var reduceSize = min(u32(Length), ${t}u);
+            for (var currentSize = reduceSize / 2u; reduceSize > 1u;
+                currentSize = reduceSize / 2u) {
+                let interval = DIV_CEIL(reduceSize, 2u);
+                if (localId.x < currentSize) {
+                    let candidate = bestValues[localId.x + interval];
+                    bestValue = bestValue + candidate;
+                    bestValues[localId.x] = bestValue;
+                }
+                reduceSize = interval;
+                workgroupBarrier();
+            }
+
+            bestValue = bestValues[0];
+        `;
+}
+function g(e) {
+  const u = `${e.workgroupSizeX}`, t = e.subgroups && !e.variableSubgroups ? "" : `
+             var<workgroup> bestValues : array<f32, ${e.workgroupSizeX}>;
+           `, i = d(e.subgroups, e.variableSubgroups, e.workgroupSizeX, !1);
   return `
            fn DIV_CEIL(a : u32, b : u32) -> u32 {
             return ((a - 1u) / b + 1u);
            }
 
-           ${`
-             var<workgroup> xBestValues : array<f32, ${e}>;
-           `}
+            fn readInput(index: i32) -> vec2<f32> {
+                ${e.inputReadSnippet ? e.inputReadSnippet : `
+                let packed = u32(x[index]);
+                return unpack2x16float(packed);
+                `}
+            }
+
+           ${t}
     
-           ${p("index")} {
-                let outputIndex = index / ${e};
+           ${a("index")} {
+                let outputIndex = index / ${u};
                 let offset = outputIndex * uniforms.reduceSize;
-                var bestValue = 0.0;
+                var bestValue = 0.0f;
                 let Length = uniforms.reduceSize;
+                let tid = i32(localId.x);
     
-                for (var k = i32(localId.x); k < Length;
-                    k = k + ${e}) {
-                    var candidate = f32(x[offset + k]);
-                    ${t}
-                    bestValue = bestValue + candidate;
+                for (var k = tid; k < Length;
+                    k = k + ${u}) {
+                    var candidate = readInput(offset + k);
+                    ${e.inputSnippet}
+                    bestValue = bestValue + candidate.x + candidate.y;
                 }
-                xBestValues[localId.x] = bestValue;
-                workgroupBarrier();
+
+                ${i}
+                bestValue = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize * 2i)" : ""};
+
+                ${e.reducedSnippet ? e.reducedSnippet : ""}
+
+                for (var k = tid; k < Length;
+                    k = k + ${u}) {
+                    ${e.outputSnippet}
+                }
+           }
+         `;
+}
+function k(e) {
+  const u = `${e.workgroupSizeX}`, t = e.subgroups && !e.variableSubgroups ? "" : `
+             var<workgroup> bestValues : array<vec2<f32>, ${e.workgroupSizeX}>;
+           `, i = d(e.subgroups, e.variableSubgroups, e.workgroupSizeX, !0);
+  return `
+           fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+           }
+
+           fn readInput(index: i32) -> vec2<f32> {
+                ${e.inputReadSnippet ? e.inputReadSnippet : `
+                let packed = u32(x[index]);
+                return unpack2x16float(packed);
+                `}
+            }
+
+           ${t}
     
-                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();
+           ${a("index")} {
+                let outputIndex = index / ${u};
+                let offset1 = outputIndex * 2 * uniforms.reduceSize;
+                let offset2 = offset1 + uniforms.reduceSize;
+                var bestValue = vec2<f32>(0.0f, 0.0f);
+                let Length = uniforms.reduceSize;
+                let tid = i32(localId.x);
+    
+                for (var k = tid; k < Length;
+                    k = k + ${u}) {
+                    var candidate = readInput(offset1 + k);
+                    ${e.inputSnippet}
+                    let bv1 = candidate.x + candidate.y;
+
+                    candidate = readInput(offset2 + k);
+                    ${e.inputSnippet}
+                    let bv2 = candidate.x + candidate.y;
+
+                    bestValue = bestValue + vec2<f32>(bv1, bv2);
+                }
+                ${i}
+                bestValue = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize * 2i)" : ""};
+
+                ${e.reducedSnippet ?? ""}
+                ${e.outputSnippet}
+           }
+         `;
+}
+function z(e) {
+  return e.elementwise ? g(e) : k(e);
+}
+function x(e) {
+  const u = `${e.workgroupSizeX}`, t = e.subgroups && !e.variableSubgroups ? "" : `
+             var<workgroup> bestValues : array<f32, ${e.workgroupSizeX}>;
+           `, i = d(e.subgroups, e.variableSubgroups, e.workgroupSizeX, !1);
+  return `
+           fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+           }
+
+           fn readInput(index: i32) -> f32 {
+                ${e.inputReadSnippet ? e.inputReadSnippet : `
+                return x[index];
+                `}
+            }
+
+           ${t}
+    
+           ${a("index")} {
+                let outputIndex = index / ${e.workgroupSizeX};
+                let offset = outputIndex * uniforms.reduceSize;
+                var bestValue = 0.0f;
+                let Length = uniforms.reduceSize;
+                let tid = i32(localId.x);
+    
+                for (var k = tid; k < Length;
+                    k = k + ${e.workgroupSizeX}) {
+                    var candidate = readInput(offset + k);
+                    ${e.inputSnippet}
+                    bestValue = bestValue + candidate;
                 }
+                ${i}
 
-                bestValue = xBestValues[0] ${r === "mean" ? "/ f32(uniforms.reduceSize)" : ""};
+                bestValue = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize)" : ""};
 
-                ${s}
+                ${e.reducedSnippet}
 
-                for (var k = i32(localId.x); k < Length;
-                    k = k + ${e}) {
-                    ${u}
+                for (var k = tid; k < Length;
+                    k = k + ${u}) {
+                    ${e.outputSnippet}
                 }
            }
          `;
 }
-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 C(e, u) {
+  const t = e[0], n = f(u, t.shape), [, r] = h(t.shape, n), o = c(r), p = c(t.shape) / o;
+  return { windowSize: o, inSize: o, batchSize: p, outSize: 1 };
+}
+class y {
+  outputShape;
+  shaderKey = "reduce16";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  variableNames = ["x"];
+  uniforms = "reduceSize : i32,";
+  inputShape;
+  size = !1;
+  packed = !0;
+  outputComponent;
+  variableComponents;
+  elementwise;
+  subgroups = !1;
+  subgroupBuiltins = !1;
+  deviceInfo;
+  params;
+  constructor(u, t, i, n) {
+    this.params = i, this.inputShape = [t.batchSize, t.inSize], this.deviceInfo = u, this.packed = n;
+    const r = t.inSize % 64 === 0 ? 64 : 32;
+    u.subgroupsSupported ? (this.workgroupSize = [Math.min(r, u.subgroupMaxSize), 1, 1], this.subgroups = !0, u.variableSubgroups && (this.subgroupBuiltins = !0)) : this.workgroupSize[0] = r, this.outputShape = i.elementwise ? [t.batchSize, t.inSize] : [t.batchSize / 2], this.dispatchLayout = b(this.outputShape), this.dispatch = [i.elementwise ? t.batchSize : t.batchSize / 2, 1, 1], this.outputComponent = 1, this.variableComponents = [1], this.elementwise = i.elementwise === !0;
+  }
+  getWriteSnippet() {
+    return this.packed ? "result[outputIndex] = i32(pack2x16float(bestValue));" : "result[outputIndex] = bestValue;";
+  }
+  getPreprocessSnippet() {
+    return "";
+  }
+  getPostprocessSnippet() {
+    return "";
+  }
+  getReadSnippet() {
+    return this.packed ? `
+                let packed = u32(x[index]);
+                return unpack2x16float(packed);
+                ` : "return x[index];";
+  }
+  getUserCode() {
+    const u = this.workgroupSize[0];
+    return this.packed ? z({
+      ...this.params,
+      workgroupSizeX: u,
+      subgroups: this.subgroups,
+      variableSubgroups: this.deviceInfo.variableSubgroups,
+      inputReadSnippet: this.getReadSnippet(),
+      inputSnippet: this.getPreprocessSnippet(),
+      outputSnippet: this.getWriteSnippet(),
+      reducedSnippet: this.getPostprocessSnippet()
+    }) : x({
+      ...this.params,
+      workgroupSizeX: u,
+      subgroups: this.subgroups,
+      variableSubgroups: this.deviceInfo.variableSubgroups,
+      inputReadSnippet: this.getReadSnippet(),
+      inputSnippet: this.getPreprocessSnippet(),
+      outputSnippet: this.getWriteSnippet(),
+      reducedSnippet: this.getPostprocessSnippet()
+    });
+  }
 }
-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;
+function X(e, u, t) {
+  const i = u[0], r = [{ type: "int32", data: [e.inputShape[e.inputShape.length - 1]] }], o = t.runWebGPUProgram(
+    e,
+    u,
+    e.packed ? "int32" : "float32",
+    r
+  );
+  o.packed = e.packed ?? !1;
+  const s = l().makeTensorFromTensorInfo(o), p = S(
+    s,
+    e.elementwise ? i.shape : e.packed ? [...i.shape.slice(0, -2), i.shape[i.shape.length - 2] / 2] : [...i.shape.slice(0, -2), i.shape[i.shape.length - 2]]
+  );
+  return s.dispose(), p;
 }
 export {
-  V as createReduceInfo,
-  I as createReductionShader,
-  $ as reduce
+  y as ReduceProgram,
+  C as createReduceInfo,
+  X as reduce
 };