import { e as r } from "../../webgpu_program-CbjdYLYk.js";
import { f as a, c as u } from "../../webgpu_util-DuofJBMo.js";
class p {
  variableNames = ["q", "k"];
  outputShape;
  shaderKey = "AttentionMask";
  dispatchLayout;
  dispatch;
  uniforms = "divisor: f32, pastLen: i32, inf: f32";
  workgroupSize = [64, 1, 1];
  size = !0;
  hs;
  nh;
  T1;
  T2;
  constructor(e, i, s, o, t) {
    if (this.shaderKey = `AttentionMask_${t}`, this.outputShape = [e, i, s, o], this.hs = t, this.nh = i, this.T1 = s, this.T2 = o, this.dispatchLayout = a(this.outputShape), this.dispatch = u(this.dispatchLayout, this.outputShape, this.workgroupSize), t % 4 !== 0)
      throw new Error("Head size must be a multiple of 4 for AttentionMaskProgram");
  }
  getUserCode() {
    return `
            ${r("index")} {
                
                let coords = getCoordsFromIndex(index);
                let b = coords[0];
                let h = coords[1];
                let t1 = coords[2];
                let t2 = coords[3];

                if (index < uniforms.size) {
                    if (t2 > t1 + uniforms.pastLen) {
                        setOutputAtIndex(index, uniforms.inf);
                        return;
                    }

                    let q0 = getIndexFromCoords4D(vec4<i32>(b, h, t1, 0), uniforms.qShape);
                    let k0 = getIndexFromCoords4D(vec4<i32>(b, h, t2, 0), uniforms.kShape);
                    
                    var sum: f32 = 0.0;
                    for (var i: i32 = 0; i < ${this.hs}; i = i + 4) {
                        let qv = vec4<f32>(q[q0 + i], q[q0 + i + 1], q[q0 + i + 2], q[q0 + i + 3]);
                        let kv = vec4<f32>(k[k0 + i], k[k0 + i + 1], k[k0 + i + 2], k[k0 + i + 3]);
                        sum = sum + dot(qv, kv);
                    }
                    let scaled = sum * uniforms.divisor;
                    setOutputAtIndex(index, scaled);
                }
            }
        `;
  }
}
export {
  p as default
};
