import { isPackedTensor as y } from "../../utilities/packed.js";
import { e as c } from "../../webgpu_program-CbjdYLYk.js";
import { f as x, c as l } from "../../webgpu_util-DuofJBMo.js";
import { a as w, a6 as b } from "../../index-DSGwv2Yx.js";
class v {
  variableNames = ["x", "sin", "cos"];
  outputShape;
  shaderKey = "Rope";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  size = !0;
  uniforms = "pastLen: i32";
  constructor(e, o, a, t) {
    this.shaderKey = `Rope_${t}`, this.outputShape = [e, o, a, t], this.dispatchLayout = x(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize);
  }
  getUserCode() {
    const e = this.outputShape[3];
    return `
        ${c("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 rotaryDim = ${e};

                var outVal = 0.0;

                let xIdx = b * uniforms.outShapeStrides[0] +
                    h * uniforms.outShapeStrides[1] +
                    t * uniforms.outShapeStrides[2] +
                    d;

                if (d < rotaryDim) {
                    let idx = (t + uniforms.pastLen) * uniforms.cosShape[1] + d / 2;
                    let cos = cos[idx];
                    let sin = sin[idx];

                    let ownX = x[xIdx] * cos;
                    var evenOdd = 0.0;

                    if (d % 2 == 0) {
                        // even index
                        evenOdd = -x[xIdx + 1];
                    } else {
                        // odd index
                        evenOdd = x[xIdx - 1];
                    }

                    outVal = fma(evenOdd, sin, ownX);
                } else {
                    // pass through for non-rotary dims
                    outVal = x[xIdx];
                }

                setOutputAtIndex(index, outVal);
            }
        }
        `;
  }
}
class k {
  variableNames = ["x", "sin", "cos"];
  outputShape;
  shaderKey = "Rope";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  size = !0;
  uniforms = "pastLen: i32";
  constructor(e, o, a, t) {
    this.shaderKey = `Rope_${t}`, this.outputShape = [e, o, a, t / 2], this.dispatchLayout = x(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize);
  }
  getUserCode() {
    return `
        ${c("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];

                var outVal = vec2<f32>(0.0, 0.0);

                let xIdx = b * uniforms.outShapeStrides[0] +
                    h * uniforms.outShapeStrides[1] +
                    t * uniforms.outShapeStrides[2] +
                    d;

                let idx = (t + uniforms.pastLen) * uniforms.cosShape[1] + d;
                let cos = cos[idx];
                let sin = sin[idx];

                let xPair = unpack2x16float(u32(x[xIdx]));
                let ownX = vec2<f32>(xPair.x * cos, xPair.y * cos);

                let evenOdd = vec2<f32>(
                    -xPair.y,
                    xPair.x
                );

                outVal = vec2<f32>(
                    fma(evenOdd.x, sin, ownX.x),
                    fma(evenOdd.y, sin, ownX.y)
                );

                result[index] = i32(pack2x16float(outVal));
            }
        }
        `;
  }
}
function L(i) {
  const { x: e } = i.inputs, { pastLen: o, negSin: a, ropeCache: t } = i.attrs, m = i.backend, d = y(e), p = e.shape[0], h = e.shape[1], r = e.shape[2], n = d ? e.shape[3] * 2 : e.shape[3], s = a ? t.getNegSin() : t.getSin(), u = t.getCos();
  if (b(s.shape, u.shape, "Error in Rope: "), s.shape[0] < r + o)
    throw new Error(
      `Sin tensor shape ${s.shape} is not compatible with seqLength ${r} and pastLen ${o}.`
    );
  if (s.shape[1] * 2 < n)
    throw new Error(`Sin tensor shape ${s.shape} is not compatible with feature dimension ${n}.`);
  if (s.shape.length !== 3)
    throw new Error(`Sin tensor must be 3-dimensional, but got shape ${s.shape}.`);
  const f = d ? new k(p, h, r, n) : new v(p, h, r, n), S = [{ type: "int32", data: [o] }], g = d ? "packedF16" : e.dtype;
  return m.runWebGPUProgram(f, [e, s, u], g, S);
}
const P = {
  kernelName: "Rope",
  backendName: "webgpu",
  kernelFunc: L
};
w(P);
