import { a as i } from "../../index-DSGwv2Yx.js";
class l {
  variableNames = ["x", "kernel"];
  outputShape;
  userCode;
  // enableShapeUniforms = true;
  customUniforms = [{ name: "mode", type: "int" }];
  constructor(e, t, s, o) {
    const n = o / t;
    this.outputShape = [e, t, s, n], this.userCode = `
        void main() {
            ivec4 coords = getOutputCoords(); // [b, h, t, d]
            int b = coords.x;
            int h = coords.y;
            int t = coords.z;
            int d = coords.w;

            // Compute output channel index in fused kernel
            int out_offset = mode * ${t} * ${n} + h * ${n} + d;

            float sum = 0.0;
            for (int c = 0; c < ${o}; ++c) {
                float xval = getX(b, t, c); // fetch from x
                float kval = getKernel(c, out_offset); // fetch from kernel
                sum += xval * kval;
            }

            setOutput(sum);
        }
        `;
  }
}
function m(r) {
  const { x: e, kernel: t } = r.inputs, { heads: s } = r.attrs, o = r.backend, n = e.shape[0], c = e.shape[1], u = e.shape[2], a = new l(n, s, c, u);
  return [
    o.runWebGLProgram(a, [e, t], "float32", [[0]]),
    o.runWebGLProgram(a, [e, t], "float32", [[1]]),
    o.runWebGLProgram(a, [e, t], "float32", [[2]])
  ];
}
const d = {
  kernelName: "QKV",
  backendName: "webgl",
  kernelFunc: m
};
i(d);
