import { isPackedTensor as T } from "../../utilities/packed.js";
import { e as p } from "../../webgpu_program-CbjdYLYk.js";
import { f as d, c as u } from "../../webgpu_util-DuofJBMo.js";
import { a as S, a6 as g } from "../../index-DSGwv2Yx.js";
class x {
  variableNames = ["cache", "item"];
  outputShape;
  shaderKey = "AppendCache";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  size = !0;
  uniforms = "cacheT: i32";
  constructor(e, t, a, s, i) {
    const o = Math.min(a + 1, i);
    this.shaderKey = `AppendCache_${o}`, this.outputShape = [e, t, o, s], this.dispatchLayout = d(this.outputShape), this.dispatch = u(this.dispatchLayout, this.outputShape, this.workgroupSize);
  }
  getUserCode() {
    const e = this.outputShape[2];
    return `
        ${p("index")} {
            if (index < uniforms.size) {
                let coords = getCoordsFromIndex(index); // [b, h, t, d]
                let b = coords[0];
                let h = coords[1];
                let t = coords[2];
                let d = coords[3];

                let itemT = 1;
                let maxSize = ${e};
                let totalT = uniforms.cacheT + itemT;
                let start = select(0, 1, totalT >= maxSize);

                let srcT = t + start;
                var val = 0.0;
                if (srcT < uniforms.cacheT) {
                    val = getCache(b, h, srcT, d);
                }
                if (srcT == uniforms.cacheT) {
                    val = getItem(b, h, 0, d);
                }

                setOutputAtIndex(index, val);
            }
        }
        `;
  }
}
class b {
  variableNames = ["cache", "item"];
  outputShape;
  shaderKey = "AppendCache";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  size = !0;
  uniforms = "cacheT: i32";
  constructor(e, t, a, s, i) {
    const o = Math.min(a + 1, i);
    this.shaderKey = `AppendCache_${o}`, this.outputShape = [e, t, o, s], this.dispatchLayout = d(this.outputShape), this.dispatch = u(this.dispatchLayout, this.outputShape, this.workgroupSize);
  }
  getUserCode() {
    const e = this.outputShape[2];
    return `
        ${p("index")} {
            if (index < uniforms.size) {
                let coords = getCoordsFromIndex(index); // [b, h, t, d]
                let b = coords[0];
                let h = coords[1];
                let t = coords[2];
                let d = coords[3];

                let itemT = 1;
                let maxSize = ${e};
                let totalT = uniforms.cacheT + itemT;
                let start = select(0, 1, totalT >= maxSize);

                let srcT = t + start;
                var val: i32 = 0i;
                if (srcT < uniforms.cacheT) {
                    val = cache[getIndexFromCoords4D(vec4<i32>(b, h, srcT, d), uniforms.cacheShape)];
                }
                if (srcT == uniforms.cacheT) {
                    val = item[getIndexFromCoords4D(vec4<i32>(b, h, 0, d), uniforms.itemShape)];
                }

                result[index] = val;
            }
        }
        `;
  }
}
function C(r) {
  const { cache: e, item: t } = r.inputs, { maxSize: a, pastLen: s } = r.attrs, i = r.backend, o = T(e), c = e.shape[0], n = e.shape[2], h = e.shape[1];
  if (g(t.shape, [c, h, 1, t.shape[3]], "Error in AppendCache: "), s < 0 || s > a)
    throw new Error(`Invalid pastLen value: ${s}. Must be in the range [0, ${a}].`);
  const m = o ? new b(c, h, n, t.shape[3], a) : new x(c, h, n, t.shape[3], a), l = [{ type: "int32", data: [s] }], f = o ? "packedF16" : e.dtype;
  return i.runWebGPUProgram(m, [e, t], f, l);
}
const v = {
  kernelName: "AppendCache",
  backendName: "webgpu",
  kernelFunc: C
};
S(v);
