import { ad as z, ac as F, ab as E, a9 as j, v as A } from "./index-DSGwv2Yx.js";
function L(t, s) {
  if (Math.max(...t) > 5)
    throw new Error("Cannot symbolically compute strides for rank > 6 tensor.");
  const e = t.length, o = "xyzwuv", i = t.map((n) => `${s}.${o[n]}`), u = new Array(e - 1);
  u[e - 2] = i[e - 1];
  for (let n = e - 3; n >= 0; --n)
    u[n] = `(${u[n + 1]} * ${i[n + 1]})`;
  return u;
}
const X = (t, s, e) => e === "int32" ? `atomicAdd(${t}, bitcast<i32>(${s}));` : `
          {
            var oldValue = 0;
            loop {
              let newValueF32 = bitcast<f32>(oldValue) + (${s});
              let newValue = bitcast<i32>(newValueF32);
              let res = atomicCompareExchangeWeak(${t}, oldValue, newValue);
              if res.exchanged {
                break;
              }
              oldValue = res.old_value;
            }
          }`;
var w;
(function(t) {
  t[t.FROM_PIXELS = 0] = "FROM_PIXELS", t[t.DRAW = 1] = "DRAW";
})(w || (w = {}));
const H = (t, s, e, o, i) => {
  const u = { dtype: o.dtype, shape: o.shape }, n = D(e, u, s), r = t.createShaderModule({ code: n, label: s.constructor.name });
  let d = E().get("WEBGPU_PRINT_SHADER");
  if (d !== "") {
    d = d.toLowerCase();
    const p = d.split(",");
    (d === "all" || p.some((a) => s.shaderKey.toLowerCase().includes(a))) && (console.group(s.shaderKey), console.debug(n), console.groupEnd());
  }
  return i ? t.createComputePipelineAsync({
    compute: { module: r, entryPoint: "_start" },
    label: s.constructor.name,
    layout: "auto"
  }) : t.createComputePipeline({
    compute: { module: r, entryPoint: "_start" },
    label: s.constructor.name,
    layout: "auto"
  });
}, h = (t, s = "f32") => {
  switch (t) {
    case 1:
      return `${s}`;
    case 2:
      return `vec2<${s}>`;
    case 3:
      return `vec3<${s}>`;
    case 4:
      return `vec4<${s}>`;
    default:
      throw new Error(`${t}-component ${s} is not supported.`);
  }
};
function g(t) {
  if (t <= 1)
    return "i32";
  if (t === 2)
    return "vec2<i32>";
  if (t === 3)
    return "vec3<i32>";
  if (t === 4)
    return "vec4<i32>";
  if (t === 5)
    return "vec5";
  if (t === 6)
    return "vec6";
  throw Error(`GPU for rank ${t} is not yet supported`);
}
function I(t) {
  if (t === 0)
    return "x";
  if (t === 1)
    return "y";
  if (t === 2)
    return "z";
  if (t === 3)
    return "w";
  if (t === 4)
    return "u";
  if (t === 5)
    return "v";
  throw Error(`Index ${t} is not yet supported`);
}
function Y(...t) {
  let s;
  switch (t.length) {
    case 0:
      s = `
        fn main()
      `;
      break;
    case 1:
      s = `
        fn main(${t[0]} : i32)
      `;
      break;
    default:
      throw Error("Unreachable");
  }
  return s;
}
function C(t, s) {
  let e;
  return e = `
     ${N(s)}
      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>,
                ${s.subgroupBuiltins ? "@builtin(subgroup_invocation_id) SubgroupInvocationId : u32," : ""}
                ${s.subgroupBuiltins ? "@builtin(subgroup_size) SubgroupSize : u32," : ""}
                @builtin(num_workgroups) NumWorkgroups : vec3<u32>) {
        localId = LocalId;
        localIndex = LocalIndex;
        globalId = GlobalId;
        numWorkgroups = NumWorkgroups;
        workgroupId = WorkgroupId;
        ${s.subgroupBuiltins ? "subgroupInvocationId = SubgroupInvocationId" : ""};
        ${s.subgroupBuiltins ? "subgroupSize = SubgroupSize" : ""};
        ${t ? "main(getGlobalIndex());" : "main();"};
      }
    `, e;
}
function N(t) {
  return `
  @compute @workgroup_size(${t.workgroupSize[0]}, ${t.workgroupSize[1]}, ${t.workgroupSize[2]})
`;
}
function D(t, s, e) {
  const o = [], i = 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 {
        ${_(e) ? "  return i32(globalId.x);" : `  return i32((workgroupId.z * numWorkgroups.x * numWorkgroups.y +
                workgroupId.y * numWorkgroups.x + workgroupId.x) * ${i}u +
                localIndex);
        `}
      }
    `), e.pixelsOpType != null) {
    const c = e.pixelsOpType === w.FROM_PIXELS ? `@group(0) @binding(0) var<storage, read_write> result: array<${b(s.dtype, e.outputComponent)}>;` : `@group(0) @binding(1) var<storage, read> inBuf : array<${b(t[0].dtype, e.outputComponent)}>;`, $ = s.shape.length === 3 ? "vec2<i32>" : "i32";
    o.push(`
        struct Uniform {
          outShapeStrides : ${$},
          size            : i32,
          numChannels     : i32,
          alpha           : f32,
        };

        ${c}
        @group(0) @binding(2) var<uniform> uniforms: Uniform;
      `);
    const x = k(e);
    return [
      m,
      o.join(`
`),
      y(s.shape),
      e.getUserCode(),
      C(x, e)
    ].join(`
`);
  }
  let u, n, r = "struct Uniforms { NAN : f32, INFINITY : f32, ";
  e.variableNames.forEach((c, $) => {
    const x = g(t[$].shape.length);
    r += `${c.charAt(0).toLowerCase() + c.slice(1)}Shape : ${x}, `, u = t[$].shape.length - 1, n = g(u), r += `${c.charAt(0).toLowerCase() + c.slice(1)}ShapeStrides: ${n}, `;
  });
  const d = g(s.shape.length);
  r += `outShape : ${d}, `, u = s.shape.length - 1, n = g(u), r += `
         outShapeStrides: ${n}, `, e.size && (r += "size : i32, "), e.uniforms && (r += e.uniforms), r += "};", r = V(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(s.dtype, e.outputComponent)}>;
    `), e.variableNames.forEach((c, $) => {
    o.push(`
      @group(0) @binding(${1 + $}) var<storage, read> ${c}: array<${e.variableComponents ? b(t[$].dtype, e.variableComponents[$]) : b(t[$].dtype, e.outputComponent)}>;
        `);
  }), r !== "" && o.push(`
      @group(0) @binding(${1 + e.variableNames.length}) var<uniform> uniforms: Uniforms;
      `);
  const p = P(s.shape, e.dispatchLayout), v = [
    e.subgroups ? "enable subgroups;" : "",
    m,
    o.join(`
`) + T,
    y(s.shape),
    p,
    R(s.shape.length)
  ];
  e.atomic || v.push(G(s.shape, s.dtype, e.outputComponent)), e.variableNames.forEach((c, $) => {
    v.push(`${y(t[$].shape, c)}`);
  });
  const l = t.map((c, $) => B(c, s.shape, e.variableComponents ? e.variableComponents[$] : e.outputComponent, e.dispatchLayout.x.length === s.shape.length)).join(`
`);
  v.push(l), v.push(e.getUserCode());
  const f = k(e);
  return v.push(C(f, e)), v.join(`
`);
}
function q(t, s, e) {
  let o = t.shaderKey;
  if (t.pixelsOpType != null)
    return o;
  const i = [], u = [];
  s.forEach((a) => {
    i.push(a.shape), u.push(a.dtype);
  }), i.push(e.shape), u.push(e.dtype);
  const n = s.map((a) => z(a.shape, e.shape)), r = s.map((a) => F(a.shape, e.shape)).join("_"), d = n.map((a) => a.join("_")).join(";"), p = _(t) ? "flatDispatch" : "";
  return o += "_" + (t.workgroupSize ? t.workgroupSize.join(",") : "") + i.map((a) => a.length).join(",") + u.join(",") + t.variableNames.join(",") + d + r + p, o;
}
const m = `
  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 y(t, s = "") {
  const e = t.length, o = s !== "" ? `get${s.charAt(0).toUpperCase() + s.slice(1)}CoordsFromIndex` : "getCoordsFromIndex", i = s !== "" ? `${s.charAt(0).toLowerCase() + s.slice(1)}ShapeStrides` : "outShapeStrides";
  if (e <= 1)
    return `fn ${o}(index : i32) -> i32 { return index; }`;
  const u = j(t), n = g(e), r = [];
  for (let p = 0; p < e; p++)
    r.push(`d${p}`);
  if (u.length === 1)
    return `    fn ${o}(index : i32) -> vec2<i32> {
      let d0 = index / uniforms.${i}; let d1 = index - d0 * uniforms.${i};
      return vec2<i32>(d0, d1);
    }`;
  let d;
  return d = "var index2 = index;" + u.map((p, a) => {
    const v = `let ${r[a]} = index2 / uniforms.${i}.${I(a)}`, l = a === u.length - 1 ? `let ${r[a + 1]} = index2 - ${r[a]} * uniforms.${i}.${I(a)}` : `index2 = index2 - ${r[a]} * uniforms.${i}.${I(a)}`;
    return `${v}; ${l};`;
  }).join(""), `
    fn ${o}(index : i32) -> ${n} {
      ${d}
      return ${n}(${r.join(",")});
    }
  `;
}
function U(t, s) {
  const e = t.name, o = t.shape.length, i = g(o), u = "get" + e.charAt(0).toUpperCase() + e.slice(1), n = ["d0", "d1", "d2", "d3", "d4", "d5"].slice(0, o), r = n.map((a) => `${a} : i32`).join(", ");
  if (o < 1)
    return `
      fn ${u}() -> ${h(s)} {
        return ${h(s)}(${e}[0]);
      }
    `;
  const d = `uniforms.${e.charAt(0).toLowerCase() + e.slice(1)}Shape`;
  let p = `${o}D`;
  return o === 0 && (p = "1D"), `
    fn ${u}(${r}) -> ${h(s)} {
      return ${h(s)}(${e}[getIndexFromCoords${p}(${i}(${n.join(",")}),
        ${d})${s === 1 ? "" : ` / ${s}`}]);
    }
   `;
}
function W(t, s, e, o) {
  const i = t.name, u = i.charAt(0).toUpperCase() + i.slice(1), n = "get" + u + "ByOutput", r = t.shape.length, d = s.length, p = g(d);
  if (F(t.shape, s) && o)
    return `
    fn ${n}Index(globalIndex : i32) -> ${h(e)} {
      return ${h(e)}(${i}[globalIndex]);
    }

    fn ${n}Coords(coords : ${p}) -> ${h(e)} {
      return ${h(e)}(${i}[${d > 1 ? "getOutputIndexFromCoords(coords)" : "coords"}${e === 1 ? "" : ` / ${e}`}]);
    }
    `;
  const a = z(t.shape, s), v = d - r;
  let l = "";
  if (r === 0)
    return `
    fn ${n}Index(globalIndex : i32) -> ${h(e)}{
      return get${u}();
    }

    fn ${n}Coords(coords : ${p}) -> ${h(e)}{
      return get${u}();
    }
  `;
  d < 2 && a.length >= 1 ? l = "coords = 0;" : l = a.map(($) => `coords.${I($ + v)} = 0;`).join(`
`);
  let f = "";
  if (d < 2 && r > 0)
    f = "coords";
  else if (d > 1) {
    const $ = g(r), x = t.shape.map((M, O) => `coords.${I(O + v)}`).join(", ");
    f = `${$}(${x})`;
  } else
    f = "coords";
  const S = `uniforms.${i.charAt(0).toLowerCase() + i.slice(1)}Shape`, c = `${r}D`;
  return `
  fn ${n}Index(globalIndex : i32) -> ${h(e)} {
    var coords = getCoordsFromIndex(globalIndex);
    ${l}
    return ${h(e)}(${i}[getIndexFromCoords${c}(${f}, ${S})${e === 1 ? "" : ` / ${e}`}]);
  }

  fn ${n}Coords(coordsIn : ${p}) -> ${h(e)} {
    var coords = coordsIn;
    ${l}
    return ${h(e)}(${i}[getIndexFromCoords${c}(${f}, ${S})${e === 1 ? "" : ` / ${e}`}]);
  }
`;
}
function B(t, s, e, o) {
  let i = U(t, e);
  return t.shape.length <= s.length && (i += W(t, s, e, o)), i;
}
function P(t, s) {
  const { x: e, y: o = [], z: i = [] } = s, u = t.length, n = e.length + o.length + i.length;
  if (n !== u)
    return "";
  if (e.length === u)
    return `fn getOutputCoords() -> ${g(u)}{
    let globalIndex = getGlobalIndex();
    return getCoordsFromIndex(globalIndex);
  }
  `;
  let r = "";
  const d = [e, o, i];
  for (let l = 0; l < d.length; l++) {
    const f = d[l];
    if (f.length !== 0)
      if (f.length === 1)
        r += `let d${f[0]} = i32(globalId[${l}]);`;
      else {
        const S = L(f, "uniforms.outShape");
        r += `var index${l} = i32(globalId[${l}]);`;
        for (let c = 0; c < S.length; c++)
          r += `let d${f[c]} = index${l} / ${S[c]};`, c === S.length - 1 ? r += `let d${f[c + 1]} = index${l} - d${f[c]} * ${S[c]};` : r += `index${l} = index${l} - d${f[c]} * ${S[c]};`;
      }
  }
  const p = [];
  for (let l = 0; l < n; l++)
    p.push(`d${l}`);
  const a = g(n);
  let v = `fn getOutputCoords() -> ${a} {
  ${r}
`;
  return p.length === 0 ? v += `return ${a}(0); }` : v += `return ${a}(${p.join(",")}); }`, v;
}
function R(t) {
  let s = "";
  switch (t) {
    case 0:
    case 1:
      s += `
        fn getOutputIndexFromCoords(coords : i32) -> i32 {
          return coords;
        }
        `;
      break;
    case 2:
      s += `
        fn getOutputIndexFromCoords(coords : vec2<i32>) -> i32 {
          return dot(coords, vec2<i32>(uniforms.outShapeStrides, 1));
        }
        `;
      break;
    case 3:
      s += `
        fn getOutputIndexFromCoords(coords : vec3<i32>) -> i32 {
          return dot(coords, vec3<i32>(uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, 1));
        }
        `;
      break;
    case 4:
      s += `
        fn getOutputIndexFromCoords(coords : vec4<i32>) -> i32 {
          return dot(coords, vec4<i32>(
            uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, uniforms.outShapeStrides.z, 1));
        }
        `;
      break;
    case 5:
      s += `
        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:
      s += `
        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 ${t}D shape`);
      break;
  }
  return s;
}
function _(t) {
  return t.dispatch[1] === 1 && t.dispatch[2] === 1;
}
function b(t, s = 1) {
  if (t === "float32")
    return h(s, "f32");
  if (t === "int32" || t === "bool" || t === "packedF16")
    return h(s, "i32");
  throw new Error(`type ${t} is not supported.`);
}
function G(t, s, e) {
  const o = t.length, i = b(s, e);
  let u = `fn setOutputAtIndex(flatIndex : i32, value : ${h(e)}) {
      result[flatIndex] = ${i}(value);
    }

    fn setOutputAtIndexI32(flatIndex : i32, value : ${h(e, "i32")}) {
      result[flatIndex] = ${i}(value);
    }
    `;
  if (o >= 2) {
    const n = ["d0", "d1", "d2", "d3", "d4", "d5"].slice(0, o), r = g(o);
    u += `
      fn setOutputAtCoords(${n.map((d) => `${d} : i32`).join(", ")}, value : ${h(e)}) {
        let flatIndex = getOutputIndexFromCoords(${r}(${n.join(", ")}));
        setOutputAtIndex(flatIndex${e === 1 ? "" : ` / ${e}`}, value);
      }
      fn setOutputAtCoordsI32(${n.map((d) => `${d} : i32`).join(", ")}, value : ${h(e, "i32")}) {
        let flatIndex = getOutputIndexFromCoords(${r}(${n.join(", ")}));
        setOutputAtIndexI32(flatIndex${e === 1 ? "" : ` / ${e}`}, value);
      }
    `;
  }
  return u;
}
function V(t) {
  const s = /(\w+)\s*:\s*vec(5|6)/g;
  t = t.replace(s, (o) => "@align(16) " + o);
  const e = /vec(5|6)\s*,\s*(\w+)/g;
  return t = t.replace(e, (o, i, u) => `vec${i}, @align(16) ${u}`), t;
}
function k(t) {
  return !(t.dispatchLayout.hasOwnProperty("y") && t.dispatchLayout.y.length !== 0 || t.dispatchLayout.hasOwnProperty("z") && t.dispatchLayout.z.length !== 0);
}
export {
  w as P,
  y as a,
  g as b,
  I as c,
  b as d,
  Y as e,
  H as f,
  N as g,
  X as h,
  q as m,
  L as s,
  h as t
};
