@genai-fi/nanogpt 0.19.0 → 0.20.0

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

@@ -1,139 +0,0 @@
-import { f as r, c as n } from "../../../webgpu_util-DYlGSwOJ.js";
-import { g as p } from "../../../binary_op_util-pKXltfxI.js";
-import { B as k } from "../../../binary_op_util-pKXltfxI.js";
-import { _ as c, ai as s } from "../../../index-CUXkjxiT.js";
-import { e as u } from "../../../webgpu_program-B4HmApL1.js";
-class x {
-  dispatch;
-  dispatchLayout;
-  outputComponent;
-  op;
-  outputShape;
-  shaderKey;
-  size = !0;
-  variableNames = ["A", "B"];
-  workgroupSize;
-  variableComponents;
-  constructor(e, t, o) {
-    this.outputShape = c(t, o), this.dispatchLayout = r(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 && (s(o) || o[o.length - 1] === 1) || a && (s(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)}
-    };
-    `}
-       ${u("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;
-         }
-       }
-       `;
-  }
-}
-class b {
-  dispatch;
-  dispatchLayout;
-  outputComponent;
-  op;
-  outputShape;
-  shaderKey;
-  size = !0;
-  variableNames = ["A", "B"];
-  workgroupSize;
-  variableComponents;
-  constructor(e, t) {
-    if (this.outputShape = t, this.dispatchLayout = r(this.outputShape), this.op = e, t.length > 0 && t[t.length - 1] % 4 === 0)
-      this.outputComponent = 4, this.variableComponents = [4, 1];
-    else
-      throw new Error("16-bit float binary ops require inner dimension to be multiple of 4");
-    this.shaderKey = `binary_scal_${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)}
-    };
-    `}
-       ${u("index")} {
-         if (index < uniforms.size) {
-            let a = A[index];
-            let b = B[0];
-
-            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 v4res1 = binaryOperation(v4a1, vec4<f32>(b));
-            let v4res2 = binaryOperation(v4a2, vec4<f32>(b));
-
-            let res = vec4<i32>(
-                i32(pack2x16float(v4res1.xy)),
-                i32(pack2x16float(v4res1.zw)),
-                i32(pack2x16float(v4res2.xy)),
-                i32(pack2x16float(v4res2.zw))
-            );
-
-            result[index] = res;
-         }
-       }
-       `;
-  }
-}
-export {
-  x as BinaryOpProgram,
-  b as BinaryOpScalarProgram,
-  k as BinaryOpType
-};

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

@@ -1,7 +0,0 @@
-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

@@ -1,11 +0,0 @@
-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,43 +0,0 @@
-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';
-import { DeviceInformation } from './deviceInfo';
-export interface ReduceParams {
-    reductionOp: 'mean' | 'sum';
-    elementwise?: boolean;
-    forceWorkgroupSize?: number;
-}
-export declare function createReduceInfo(inputs: TensorInfo[], axis: number | number[]): backend_util.ReduceInfo;
-export declare class ReduceProgram implements WebGPUProgram {
-    atomic: boolean;
-    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;
-    utilityFunctions?: string;
-    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, extraUniforms?: {
-    type: string;
-    data: number[];
-}[], output?: TensorInfo): TensorInfo;

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

@@ -1,275 +0,0 @@
-import { ah as h, U as S, h as f } from "../../../index-CUXkjxiT.js";
-import { e as d } from "../../../webgpu_program-B4HmApL1.js";
-import { reshape16 as g } from "../../reshape16.js";
-import { f as z } from "../../../webgpu_util-DYlGSwOJ.js";
-import { c as k } from "../../../axis_util-GTVlo58H.js";
-import { z as x } from "../../../zeros-DvZpK8s6.js";
-function c(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 $(e) {
-  const u = `${e.workgroupSizeX}`, t = e.subgroups && !e.variableSubgroups ? "" : `
-             var<workgroup> bestValues : array<f32, ${e.workgroupSizeX}>;
-           `, i = c(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) -> vec2<f32> {
-                ${e.inputReadSnippet ? e.inputReadSnippet : `
-                let packed = u32(x[index]);
-                return unpack2x16float(packed);
-                `}
-            }
-
-           ${t}
-           ${e.utilityFunctions ?? ""}
-    
-           ${d("index")} {
-                let outputIndex = index / ${u};
-                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 + ${u}) {
-                    var candidate = readInput(offset + k);
-                    ${e.inputSnippet}
-                    bestValue = bestValue + candidate.x + candidate.y;
-                }
-
-                ${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 v(e) {
-  const u = `${e.workgroupSizeX}`, t = e.subgroups && !e.variableSubgroups ? "" : `
-             var<workgroup> bestValues : array<vec2<f32>, ${e.workgroupSizeX}>;
-           `, i = c(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}
-            ${e.utilityFunctions ?? ""}
-    
-           ${d("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 V(e) {
-  return e.elementwise ? $(e) : v(e);
-}
-function w(e) {
-  const u = `${e.workgroupSizeX}`, t = e.subgroups && !e.variableSubgroups ? "" : `
-             var<workgroup> bestValues : array<f32, ${e.workgroupSizeX}>;
-           `, i = c(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}
-              ${e.utilityFunctions ?? ""}
-    
-           ${d("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 = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize)" : ""};
-
-                ${e.reducedSnippet}
-
-                ${e.elementwise ? `for (var k = tid; k < Length;
-                    k = k + ${u}) {
-                    ${e.outputSnippet}
-                }` : `${e.outputSnippet}`}
-           }
-         `;
-}
-function X(e, u) {
-  const t = e[0], o = h(u, t.shape), [, n] = k(t.shape, o), s = S(n), r = S(t.shape) / s;
-  return { windowSize: s, inSize: s, batchSize: r, outSize: r };
-}
-class P {
-  atomic = !1;
-  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;
-  utilityFunctions;
-  constructor(u, t, i, o) {
-    this.params = i, this.inputShape = [t.batchSize, t.inSize], this.deviceInfo = u, this.packed = o;
-    const n = i.forceWorkgroupSize ? i.forceWorkgroupSize : t.inSize % 64 === 0 ? 64 : 32;
-    u.subgroupsSupported && !i.forceWorkgroupSize ? (this.workgroupSize = [Math.min(n, u.subgroupMaxSize), 1, 1], this.subgroups = !0, u.variableSubgroups && (this.subgroupBuiltins = !0)) : this.workgroupSize[0] = n, this.outputShape = i.elementwise ? [t.batchSize, t.inSize] : o ? [t.outSize / 2] : [t.outSize], this.dispatchLayout = z(this.outputShape), this.dispatch = [
-      i.elementwise ? t.batchSize : o ? t.batchSize / 2 : t.batchSize,
-      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 f32(x[index]);";
-  }
-  getUserCode() {
-    const u = this.workgroupSize[0];
-    return this.packed ? V({
-      ...this.params,
-      workgroupSizeX: u,
-      subgroups: this.subgroups,
-      variableSubgroups: this.deviceInfo.variableSubgroups,
-      inputReadSnippet: this.getReadSnippet(),
-      inputSnippet: this.getPreprocessSnippet(),
-      outputSnippet: this.getWriteSnippet(),
-      reducedSnippet: this.getPostprocessSnippet(),
-      utilityFunctions: this.utilityFunctions
-    }) : w({
-      ...this.params,
-      workgroupSizeX: u,
-      subgroups: this.subgroups,
-      variableSubgroups: this.deviceInfo.variableSubgroups,
-      inputReadSnippet: this.getReadSnippet(),
-      inputSnippet: this.getPreprocessSnippet(),
-      outputSnippet: this.getWriteSnippet(),
-      reducedSnippet: this.getPostprocessSnippet(),
-      utilityFunctions: this.utilityFunctions
-    });
-  }
-}
-function A(e, u, t, i, o) {
-  const n = u[0], a = [{ type: "int32", data: [e.inputShape[e.inputShape.length - 1]] }, ...i ?? []];
-  let r = o;
-  !o && e.atomic && (r = x(e.outputShape, "int32"));
-  const l = t.runWebGPUProgram(
-    e,
-    u,
-    e.packed ? "packedF16" : e.atomic ? "int32" : "float32",
-    a,
-    r
-  ), p = f().makeTensorFromTensorInfo(l);
-  if (e.outputShape.length === 1 && e.outputShape[0] <= 2)
-    return p;
-  const b = g(
-    p,
-    e.elementwise ? n.shape : e.packed ? [...n.shape.slice(0, -2), n.shape[n.shape.length - 2] / 2] : [...n.shape.slice(0, -2), n.shape[n.shape.length - 2]]
-  );
-  return p.dispose(), b;
-}
-export {
-  P as ReduceProgram,
-  X as createReduceInfo,
-  A as reduce
-};