import { a as x, ah as I, N as c } from "../../index-DSGwv2Yx.js";
import { e as D } from "../../webgpu_program-CbjdYLYk.js";
import { f as $, c as F } from "../../webgpu_util-DuofJBMo.js";
import { r as g } from "../../Reshape-pwprEaej.js";
import { a as L, c as d } from "../../concat_util-C1Mxe27t.js";
class T {
  outputShape;
  shaderKey;
  dispatchLayout;
  dispatch;
  variableNames;
  uniforms = "";
  workPerThread = 1;
  workgroupSize = [64, 1, 1];
  size = !0;
  offsetLength;
  constructor(t) {
    this.outputShape = d(
      t,
      1
      /* axis */
    ), this.variableNames = t.map((e, a) => `T${a}`), this.dispatchLayout = $(this.outputShape), this.dispatch = F(this.dispatchLayout, this.outputShape, this.workgroupSize, [
      this.workPerThread,
      1,
      1
    ]), this.offsetLength = t.length - 1;
    for (let e = 0; e < this.offsetLength; e++)
      this.uniforms += `offset${e} : i32,`;
    this.shaderKey = "concat16";
  }
  getUserCode() {
    const t = [];
    if (this.offsetLength > 0) {
      t.push(
        "if (yC < uniforms.offset0){ result[getIndexFromCoords2D(coords, uniforms.outShape)] = T0[getIndexFromCoords2D(vec2<i32>(yR, yC), uniforms.t0Shape)]; }"
      );
      for (let s = 1; s < this.offsetLength; s++)
        t.push(
          `else if (yC < uniforms.offset${[s]}){ result[getIndexFromCoords2D(coords, uniforms.outShape)] = T${s}[getIndexFromCoords2D(vec2<i32>(yR, yC - uniforms.offset${s - 1}), uniforms.t${s}Shape)]; }`
        );
      const a = this.offsetLength, i = this.offsetLength - 1;
      t.push(
        `else { result[getIndexFromCoords2D(coords, uniforms.outShape)] = T${a}[getIndexFromCoords2D(vec2<i32>(yR, yC - uniforms.offset${i}), uniforms.t${a}Shape)]; }`
      );
    } else
      t.push(
        "result[getIndexFromCoords2D(coords, uniforms.outShape)] = T0[getIndexFromCoords2D(vec2<i32>(yR, yC), uniforms.t0Shape)];"
      );
    return `
      ${D("index")} {
        for(var i = 0; i < ${this.workPerThread}; i = i + 1) {
          let flatIndex = index * ${this.workPerThread} + i;
          if(flatIndex < uniforms.size) {
            let coords = getCoordsFromIndex(flatIndex);
            let yR = coords.x;
            let yC = coords.y;

            ${t.join(`
        `)}
          }
        }
      }
    `;
  }
}
function m(n, t, e) {
  const a = e.device.limits.maxStorageBuffersPerShaderStage - 1;
  if (n.length > a) {
    const o = [];
    for (let p = 0; p < n.length; p += a) {
      const C = n.slice(p, p + a);
      o.push(m(C, t, e));
    }
    const S = m(o, t, e);
    for (const p of o)
      e.disposeData(p.dataId);
    return S;
  }
  const { tensors2D: i, outShape: s } = P(n, t, e), h = i.map((o) => o.shape), u = new T(h), f = [], r = new Array(h.length - 1);
  if (r.length > 0) {
    r[0] = h[0][1], f.push({ type: "int32", data: [r[0]] });
    for (let o = 1; o < r.length; o++)
      r[o] = r[o - 1] + h[o][1], f.push({ type: "int32", data: [r[o]] });
  }
  const l = e.runWebGPUProgram(u, i, i[0].dtype, f);
  i.forEach((o) => e.disposeData(o.dataId));
  const y = g({ inputs: { x: l }, backend: e, attrs: { shape: s } });
  return e.disposeData(l.dataId), y;
}
function P(n, t, e) {
  const a = d(
    n.map((s) => s.shape),
    t
  );
  return { tensors2D: n.map(
    (s) => g({
      inputs: { x: s },
      backend: e,
      attrs: {
        shape: [c(s.shape.slice(0, t)), c(s.shape.slice(t))]
      }
    })
  ), outShape: a };
}
function w(n) {
  const { inputs: t, backend: e, attrs: a } = n, { axis: i } = a, s = I(i, t[0].shape)[0], h = t.map((r) => r.shape);
  L(h, s);
  const u = d(
    t.map((r) => r.shape),
    s
  );
  if (c(u) === 0)
    return e.makeTensorInfo(u, t[0].dtype, []);
  const f = t.filter((r) => c(r.shape) > 0);
  return m(f, s, e);
}
const v = {
  kernelName: "Concat16",
  backendName: "webgpu",
  kernelFunc: w
};
x(v);
export {
  T as ConcatProgram,
  v as concatConfig
};
