@genai-fi/nanogpt 0.9.1 → 0.10.1

package/dist/patches/webgpu_program.js

@@ -0,0 +1,401 @@
+import { h as z } from "../index-ZyQhjEPo.js";
+import { g as F, d as b, a as x, b as g, s as O, t as l, c as y } from "../webgpu_program-Cigz-7RF.js";
+import { e as _, a as A, b as L } from "../tensor-DdQUJZlz.js";
+var N = /* @__PURE__ */ ((s) => (s[s.FROM_PIXELS = 0] = "FROM_PIXELS", s[s.DRAW = 1] = "DRAW", s))(N || {});
+const K = (s, t, e, o, u) => {
+  const a = { dtype: o.dtype, shape: o.shape }, n = D(e, a, t), r = s.createShaderModule({ code: n, label: t.constructor.name });
+  let d = _().get("WEBGPU_PRINT_SHADER");
+  if (d !== "") {
+    d = d.toLowerCase();
+    const f = d.split(",");
+    (d === "all" || f.some((v) => t.shaderKey.toLowerCase().includes(v))) && (console.group(t.shaderKey), console.debug(n), console.groupEnd());
+  }
+  return u ? s.createComputePipelineAsync({
+    compute: { module: r, entryPoint: "_start" },
+    label: t.constructor.name,
+    layout: "auto"
+  }) : s.createComputePipeline({
+    compute: { module: r, entryPoint: "_start" },
+    label: t.constructor.name,
+    layout: "auto"
+  });
+};
+function D(s, t, e) {
+  const o = [], u = e.workgroupSize[0] * e.workgroupSize[1] * e.workgroupSize[2];
+  if (e.outputComponent = e.outputComponent ? e.outputComponent : 1, o.push(`
+
+      var<private> localId: vec3<u32>;
+      var<private> localIndex: u32;
+      var<private> globalId: vec3<u32>;
+      var<private> numWorkgroups: vec3<u32>;
+      var<private> workgroupId: vec3<u32>;
+        ${e.subgroupBuiltins ? "var<private> subgroupInvocationId: u32;" : ""}
+        ${e.subgroupBuiltins ? "var<private> subgroupSize: u32;" : ""}
+
+      // Only used when the y/z dimension of workgroup size is 1.
+      fn getGlobalIndex() -> i32 {
+        ${R(e) ? "  return i32(globalId.x);" : `  return i32((workgroupId.z * numWorkgroups.x * numWorkgroups.y +
+                workgroupId.y * numWorkgroups.x + workgroupId.x) * ${u}u +
+                localIndex);
+        `}
+      }
+    `), e.pixelsOpType != null) {
+    const i = e.pixelsOpType === 0 ? `@group(0) @binding(0) var<storage, read_write> result: array<${b(
+      t.dtype,
+      e.outputComponent
+    )}>;` : `@group(0) @binding(1) var<storage, read> inBuf : array<${b(
+      s[0].dtype,
+      e.outputComponent
+    )}>;`, h = t.shape.length === 3 ? "vec2<i32>" : "i32";
+    o.push(`
+        struct Uniform {
+          outShapeStrides : ${h},
+          size            : i32,
+          numChannels     : i32,
+          alpha           : f32,
+        };
+
+        ${i}
+        @group(0) @binding(2) var<uniform> uniforms: Uniform;
+      `);
+    const I = w(e);
+    return [
+      C,
+      o.join(`
+`),
+      x(t.shape),
+      e.getUserCode(),
+      m(I, e)
+    ].join(`
+`);
+  }
+  let a, n, r = "struct Uniforms { NAN : f32, INFINITY : f32, ";
+  e.variableNames.forEach((i, h) => {
+    const I = g(s[h].shape.length);
+    r += `${i.charAt(0).toLowerCase() + i.slice(1)}Shape : ${I}, `, a = s[h].shape.length - 1, n = g(a), r += `${i.charAt(0).toLowerCase() + i.slice(1)}ShapeStrides: ${n}, `;
+  });
+  const d = g(t.shape.length);
+  r += `outShape : ${d}, `, a = t.shape.length - 1, n = g(a), r += `
+         outShapeStrides: ${n}, `, e.size && (r += "size : i32, "), e.uniforms && (r += e.uniforms), r += "};", r = P(r), o.push(r), e.atomic ? o.push(`
+      @group(0) @binding(0) var<storage, read_write> result: array<atomic<i32>>;
+    `) : o.push(`
+      @group(0) @binding(0) var<storage, read_write> result: array<${b(
+    t.dtype,
+    e.outputComponent
+  )}>;
+    `), e.variableNames.forEach((i, h) => {
+    o.push(`
+      @group(0) @binding(${1 + h}) var<storage, read> ${i}: array<${e.variableComponents ? b(s[h].dtype, e.variableComponents[h]) : b(s[h].dtype, e.outputComponent)}>;
+        `);
+  }), r !== "" && o.push(`
+      @group(0) @binding(${1 + e.variableNames.length}) var<uniform> uniforms: Uniforms;
+      `);
+  const f = B(t.shape, e.dispatchLayout), $ = [
+    e.subgroups ? "enable subgroups;" : "",
+    C,
+    o.join(`
+`) + T,
+    x(t.shape),
+    f,
+    W(t.shape.length)
+  ];
+  e.atomic || $.push(G(t.shape, t.dtype, e.outputComponent)), e.variableNames.forEach((i, h) => {
+    $.push(`${x(s[h].shape, i)}`);
+  });
+  const c = s.map(
+    (i, h) => U(
+      i,
+      t.shape,
+      e.variableComponents ? e.variableComponents[h] : e.outputComponent,
+      e.dispatchLayout.x.length === t.shape.length
+    )
+  ).join(`
+`);
+  $.push(c), $.push(e.getUserCode());
+  const p = w(e);
+  return $.push(m(p, e)), $.join(`
+`);
+}
+const C = `
+  struct vec5 {x: i32, y: i32, z: i32, w: i32, u: i32};
+  struct vec6 {x: i32, y: i32, z: i32, w: i32, u: i32, v: i32};
+
+  // Checks whether coordinates lie within the bounds of the shape.
+  fn coordsInBounds2D(coord : vec2<i32>, shape : vec2<i32>) -> bool {
+    return all(coord >= vec2<i32>(0)) && all(coord < shape);
+  }
+  fn coordsInBounds3D(coord : vec3<i32>, shape : vec3<i32>) -> bool {
+    return all(coord >= vec3<i32>(0)) && all(coord < shape);
+  }
+  fn coordsInBounds4D(coord : vec4<i32>, shape : vec4<i32>) -> bool {
+    return all(coord >= vec4<i32>(0)) && all(coord < shape);
+  }
+
+  fn getIndexFromCoords1D(coord : i32, shape : i32) -> i32 {
+    return coord;
+  }
+  fn getIndexFromCoords2D(coords : vec2<i32>, shape : vec2<i32>) -> i32 {
+    return dot(coords, vec2<i32>(shape.y, 1));
+  }
+  fn getIndexFromCoords3D(coords : vec3<i32>, shape : vec3<i32>) -> i32 {
+    return dot(coords, vec3<i32>(shape.y * shape.z, shape.z, 1));
+  }
+  fn getIndexFromCoords4D(coords : vec4<i32>, shape : vec4<i32>) -> i32 {
+    return dot(coords, vec4<i32>(
+        shape.y * shape.z * shape.w, shape.z * shape.w, shape.w, 1));
+  }
+  fn getIndexFromCoords5D(coords : vec5, shape : vec5) -> i32 {
+    let shapeStrides: vec5 = vec5(shape.y * shape.z * shape.w * shape.u, shape.z * shape.w * shape.u, shape.w * shape.u, shape.u, 1);
+    return coords.x*shapeStrides.x + coords.y*shapeStrides.y + coords.z*shapeStrides.z + coords.w*shapeStrides.w + coords.u*shapeStrides.u;
+  }
+  fn getIndexFromCoords6D(coords : vec6, shape : vec6) -> i32 {
+    let shapeStrides: vec6 = vec6(shape.y * shape.z * shape.w * shape.u * shape.v, shape.z * shape.w * shape.u * shape.v, shape.w * shape.u * shape.v, shape.u * shape.v, shape.v, 1);
+    return coords.x*shapeStrides.x + coords.y*shapeStrides.y + coords.z*shapeStrides.z + coords.w*shapeStrides.w + coords.u*shapeStrides.u + coords.v*shapeStrides.v;
+  }
+
+  // NaN defination in IEEE 754-1985 is :
+  //   - sign = either 0 or 1.
+  //   - biased exponent = all 1 bits.
+  //   - fraction = anything except all 0 bits (since all 0 bits represents infinity).
+  // https://en.wikipedia.org/wiki/IEEE_754-1985#Representation_of_non-numbers
+  fn isnan(val: f32) -> bool {
+    let floatToUint: u32 = bitcast<u32>(val);
+    return (floatToUint & 0x7fffffffu) > 0x7f800000u;
+  }
+  fn isnanVec4(val : vec4<f32>) -> vec4<bool> {
+    let floatToUint: vec4<u32> = bitcast<vec4<u32>>(val);
+    return (floatToUint & vec4<u32>(0x7fffffffu)) > vec4<u32>(0x7f800000u);
+  }
+`, T = `
+  fn isinf(val: f32) -> bool {
+    return abs(val) == uniforms.INFINITY;
+  }
+`;
+function j(s, t) {
+  const e = s.name, o = s.shape.length, u = g(o), a = "get" + e.charAt(0).toUpperCase() + e.slice(1), n = ["d0", "d1", "d2", "d3", "d4", "d5"].slice(0, o), r = n.map((v) => `${v} : i32`).join(", ");
+  if (o < 1)
+    return `
+      fn ${a}() -> ${l(t)} {
+        return ${l(t)}(${e}[0]);
+      }
+    `;
+  const d = `uniforms.${e.charAt(0).toLowerCase() + e.slice(1)}Shape`;
+  let f = `${o}D`;
+  return o === 0 && (f = "1D"), `
+    fn ${a}(${r}) -> ${l(t)} {
+      return ${l(t)}(${e}[getIndexFromCoords${f}(${u}(${n.join(",")}),
+        ${d})${t === 1 ? "" : ` / ${t}`}]);
+    }
+   `;
+}
+function E(s, t, e, o) {
+  const u = s.name, a = u.charAt(0).toUpperCase() + u.slice(1), n = "get" + a + "ByOutput", r = s.shape.length, d = t.length, f = g(d);
+  if (L(s.shape, t) && o)
+    return `
+    fn ${n}Index(globalIndex : i32) -> ${l(e)} {
+      return ${l(e)}(${u}[globalIndex]);
+    }
+
+    fn ${n}Coords(coords : ${f}) -> ${l(e)} {
+      return ${l(e)}(${u}[${d > 1 ? "getOutputIndexFromCoords(coords)" : "coords"}${e === 1 ? "" : ` / ${e}`}]);
+    }
+    `;
+  const v = z(s.shape, t), $ = d - r;
+  let c = "";
+  if (r === 0)
+    return `
+    fn ${n}Index(globalIndex : i32) -> ${l(e)}{
+      return get${a}();
+    }
+
+    fn ${n}Coords(coords : ${f}) -> ${l(e)}{
+      return get${a}();
+    }
+  `;
+  d < 2 && v.length >= 1 ? c = "coords = 0;" : c = v.map((h) => `coords.${y(h + $)} = 0;`).join(`
+`);
+  let p = "";
+  if (d < 2 && r > 0)
+    p = "coords";
+  else if (d > 1) {
+    const h = g(r), I = s.shape.map((M, k) => `coords.${y(k + $)}`).join(", ");
+    p = `${h}(${I})`;
+  } else
+    p = "coords";
+  const S = `uniforms.${u.charAt(0).toLowerCase() + u.slice(1)}Shape`, i = `${r}D`;
+  return `
+  fn ${n}Index(globalIndex : i32) -> ${l(e)} {
+    var coords = getCoordsFromIndex(globalIndex);
+    ${c}
+    return ${l(e)}(${u}[getIndexFromCoords${i}(${p}, ${S})${e === 1 ? "" : ` / ${e}`}]);
+  }
+
+  fn ${n}Coords(coordsIn : ${f}) -> ${l(e)} {
+    var coords = coordsIn;
+    ${c}
+    return ${l(e)}(${u}[getIndexFromCoords${i}(${p}, ${S})${e === 1 ? "" : ` / ${e}`}]);
+  }
+`;
+}
+function U(s, t, e, o) {
+  let u = j(s, e);
+  return s.shape.length <= t.length && (u += E(s, t, e, o)), u;
+}
+function B(s, t) {
+  const { x: e, y: o = [], z: u = [] } = t, a = s.length, n = e.length + o.length + u.length;
+  if (n !== a)
+    return "";
+  if (e.length === a)
+    return `fn getOutputCoords() -> ${g(a)}{
+    let globalIndex = getGlobalIndex();
+    return getCoordsFromIndex(globalIndex);
+  }
+  `;
+  let r = "";
+  const d = [e, o, u];
+  for (let c = 0; c < d.length; c++) {
+    const p = d[c];
+    if (p.length !== 0)
+      if (p.length === 1)
+        r += `let d${p[0]} = i32(globalId[${c}]);`;
+      else {
+        const S = O(p, "uniforms.outShape");
+        r += `var index${c} = i32(globalId[${c}]);`;
+        for (let i = 0; i < S.length; i++)
+          r += `let d${p[i]} = index${c} / ${S[i]};`, i === S.length - 1 ? r += `let d${p[i + 1]} = index${c} - d${p[i]} * ${S[i]};` : r += `index${c} = index${c} - d${p[i]} * ${S[i]};`;
+      }
+  }
+  const f = [];
+  for (let c = 0; c < n; c++)
+    f.push(`d${c}`);
+  const v = g(n);
+  let $ = `fn getOutputCoords() -> ${v} {
+  ${r}
+`;
+  return f.length === 0 ? $ += `return ${v}(0); }` : $ += `return ${v}(${f.join(",")}); }`, $;
+}
+function W(s) {
+  let t = "";
+  switch (s) {
+    case 0:
+    case 1:
+      t += `
+        fn getOutputIndexFromCoords(coords : i32) -> i32 {
+          return coords;
+        }
+        `;
+      break;
+    case 2:
+      t += `
+        fn getOutputIndexFromCoords(coords : vec2<i32>) -> i32 {
+          return dot(coords, vec2<i32>(uniforms.outShapeStrides, 1));
+        }
+        `;
+      break;
+    case 3:
+      t += `
+        fn getOutputIndexFromCoords(coords : vec3<i32>) -> i32 {
+          return dot(coords, vec3<i32>(uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, 1));
+        }
+        `;
+      break;
+    case 4:
+      t += `
+        fn getOutputIndexFromCoords(coords : vec4<i32>) -> i32 {
+          return dot(coords, vec4<i32>(
+            uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, uniforms.outShapeStrides.z, 1));
+        }
+        `;
+      break;
+    case 5:
+      t += `
+        fn getOutputIndexFromCoords(coords : vec5) -> i32 {
+          return coords.x * uniforms.outShapeStrides.x +
+              coords.y * uniforms.outShapeStrides.y +
+              coords.z * uniforms.outShapeStrides.z +
+              coords.w * uniforms.outShapeStrides.w +
+              coords.u;
+        }
+        `;
+      break;
+    case 6:
+      t += `
+        fn getOutputIndexFromCoords(coords : vec6) -> i32 {
+          return coords.x * uniforms.outShapeStrides.x +
+              coords.y * uniforms.outShapeStrides.y +
+              coords.z * uniforms.outShapeStrides.z +
+              coords.w * uniforms.outShapeStrides.w +
+              coords.u * uniforms.outShapeStrides.u +
+              coords.v;
+        }
+        `;
+      break;
+    default:
+      A(!1, () => `Unsupported ${s}D shape`);
+      break;
+  }
+  return t;
+}
+function R(s) {
+  return s.dispatch[1] === 1 && s.dispatch[2] === 1;
+}
+function G(s, t, e) {
+  const o = s.length, u = b(t, e);
+  let a = `fn setOutputAtIndex(flatIndex : i32, value : ${l(e)}) {
+      result[flatIndex] = ${u}(value);
+    }
+
+    fn setOutputAtIndexI32(flatIndex : i32, value : ${l(e, "i32")}) {
+      result[flatIndex] = ${u}(value);
+    }
+    `;
+  if (o >= 2) {
+    const n = ["d0", "d1", "d2", "d3", "d4", "d5"].slice(0, o), r = g(o);
+    a += `
+      fn setOutputAtCoords(${n.map((d) => `${d} : i32`).join(", ")}, value : ${l(e)}) {
+        let flatIndex = getOutputIndexFromCoords(${r}(${n.join(", ")}));
+        setOutputAtIndex(flatIndex${e === 1 ? "" : ` / ${e}`}, value);
+      }
+      fn setOutputAtCoordsI32(${n.map((d) => `${d} : i32`).join(", ")}, value : ${l(e, "i32")}) {
+        let flatIndex = getOutputIndexFromCoords(${r}(${n.join(", ")}));
+        setOutputAtIndexI32(flatIndex${e === 1 ? "" : ` / ${e}`}, value);
+      }
+    `;
+  }
+  return a;
+}
+function P(s) {
+  const t = /(\w+)\s*:\s*vec(5|6)/g;
+  s = s.replace(t, (o) => "@align(16) " + o);
+  const e = /vec(5|6)\s*,\s*(\w+)/g;
+  return s = s.replace(e, (o, u, a) => `vec${u}, @align(16) ${a}`), s;
+}
+function w(s) {
+  return !(s.dispatchLayout.hasOwnProperty("y") && s.dispatchLayout.y?.length !== 0 || s.dispatchLayout.hasOwnProperty("z") && s.dispatchLayout.z?.length !== 0);
+}
+function m(s, t) {
+  return `
+     ${F(t)}
+      fn _start(@builtin(local_invocation_id) LocalId : vec3<u32>,
+                @builtin(global_invocation_id) GlobalId : vec3<u32>,
+                @builtin(local_invocation_index) LocalIndex: u32,
+                @builtin(workgroup_id) WorkgroupId : vec3<u32>,
+                ${t.subgroupBuiltins ? "@builtin(subgroup_invocation_id) SubgroupInvocationId : u32," : ""}
+                ${t.subgroupBuiltins ? "@builtin(subgroup_size) SubgroupSize : u32," : ""}
+                @builtin(num_workgroups) NumWorkgroups : vec3<u32>) {
+        localId = LocalId;
+        localIndex = LocalIndex;
+        globalId = GlobalId;
+        numWorkgroups = NumWorkgroups;
+        workgroupId = WorkgroupId;
+        ${t.subgroupBuiltins ? "subgroupInvocationId = SubgroupInvocationId" : ""};
+        ${t.subgroupBuiltins ? "subgroupSize = SubgroupSize" : ""};
+        ${s ? "main(getGlobalIndex());" : "main();"};
+      }
+    `;
+}
+export {
+  N as PixelsOpType,
+  K as compileProgram,
+  m as getStartHeaderString
+};