import { ab as z, v as F, ac as O, ad as _ } from "../index-DSGwv2Yx.js";
import { g as A, d as b, a as x, b as g, s as L, t as l, c as y } from "../webgpu_program-CbjdYLYk.js";
var N = /* @__PURE__ */ ((s) => (s[s.FROM_PIXELS = 0] = "FROM_PIXELS", s[s.DRAW = 1] = "DRAW", s))(N || {});
const H = (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 = z().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(),
      k(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(k(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 (O(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 = _(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, m) => `coords.${y(m + $)}`).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 = L(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:
      F(!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 k(s, t) {
  return `
     ${A(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,
  H as compileProgram,
  k as getStartHeaderString
};
