import { N as f } from "../../index-DSGwv2Yx.js";
import { e as p } from "../../webgpu_program-CbjdYLYk.js";
class B {
  variableNames = ["A", "B"];
  outputShape;
  shaderKey = "MatMul16TB";
  dispatchLayout;
  dispatch;
  workgroupSize = [8, 8, 1];
  // 8x8 threads for 32x32 tile
  dimInner;
  transposeA = !1;
  transposeB = !0;
  broadcastBatch = !0;
  tileInner = 32;
  uniforms;
  scale = !1;
  scaleA = !1;
  scaleB = !1;
  activation;
  causalMask = !1;
  outputComponent;
  variableComponents;
  outputIndexSnippet;
  outputStrideSnippet;
  constructor(e, t, o, a, i, r = !1, s = !1) {
    if (this.transposeA = r, this.transposeB = s, this.variableComponents = [2, 2], this.outputComponent = 2, this.shaderKey = `MatMul16TB_${t}_${o}_${a}_${i}_${r ? "TA" : ""}${s ? "TB" : ""}`, r) {
      if (this.outputShape = [e, a, i / 2], this.dimInner = t, t !== o)
        throw new Error("Inner dimensions of A and B must match for MatMul16 transposeA");
    } else if (s) {
      if (this.outputShape = [e, t, o / 2], this.dimInner = i, i !== a)
        throw new Error("Inner dimensions of A and B must match for MatMul16 transposeB");
    } else if (this.outputShape = [e, t, i / 2], this.dimInner = a, a !== o)
      throw new Error("Inner dimensions of A and B must match for MatMul16");
    if (this.dimInner % this.tileInner !== 0)
      throw new Error(`Inner dimension ${this.dimInner} must be multiple of ${this.tileInner}`);
    if (this.dispatchLayout = { x: [2], y: [1], z: [0] }, this.dispatch = [
      Math.ceil(this.outputShape[2] / (this.workgroupSize[0] * 2)),
      // 4 unpacked cols per thread = 2 packed cols
      Math.ceil(this.outputShape[1] / (this.workgroupSize[1] * 4)),
      // 4 rows per thread
      this.outputShape[0]
    ], i % 32 !== 0)
      throw new Error("Head size must be even for MatMul16 transposeB");
    if (a % 32 !== 0)
      throw new Error("Head size must be even for MatMul16 transposeB");
    if (t % 32 !== 0)
      throw new Error("Sequence length must be multiple of 32 for MatMul16 transposeB");
    if (o % 32 !== 0)
      throw new Error("Sequence length must be multiple of 32 for MatMul16 transposeB");
    this.outputIndexSnippet = "var idx0 = getOutputIndexFromCoords(vec3<i32>(batch, gRow, gColPacked));", this.outputStrideSnippet = "idx0 = idx0 + uniforms.outShapeStrides[1];  // Next row";
  }
  addUniform(e) {
    this.uniforms ? this.uniforms += `, ${e}` : this.uniforms = e;
  }
  /* Note: this is done after constructor because it shouldn't affect dispatch */
  setOutputShape(e, t) {
    const o = f(e), a = f(this.outputShape);
    if (o !== a)
      throw new Error(`New shape size ${o} must match current size ${a}`);
    function i(c, u) {
      return [`${c} / ${u}`, `${c} % ${u}`];
    }
    const r = this.outputShape;
    let s = [];
    if (e.length === r.length + 1)
      if (e[0] * e[1] === r[0])
        s = [
          ...i("batch", e[1]),
          // batch / B2, batch % B2
          "gRow",
          "gColPacked"
        ], this.shaderKey += `_batchSplit_${e[1]}`;
      else if (e[e.length - 2] * e[e.length - 1] === r[r.length - 1])
        s = [
          "batch",
          "gRow",
          ...i("gColPacked", e[e.length - 1])
          // gColPacked / N2, gColPacked % N2
        ], this.shaderKey += `_colSplit_${e[e.length - 1]}`;
      else
        throw new Error("Unsupported output shape split");
    else if (e.length === r.length)
      s = ["batch", "gRow", "gColPacked"];
    else if (e.length === 2 && r[0] === 1)
      s = ["gRow", "gColPacked"], this.shaderKey += "_batchRemoved";
    else
      throw new Error(`Unsupported output shape rank change: ${r.length} -> ${e.length}}`);
    let n = [];
    if (t) {
      if (t.length !== e.length)
        throw new Error("Permutation length must match output rank");
      n = t.map((c) => s[c]), this.shaderKey += `_perm_${t.join("")}`;
    } else
      n = s;
    const l = n.findIndex((c) => c === "gRow"), h = `vec${e.length}<i32>(${n.join(", ")})`;
    this.outputIndexSnippet = `var idx0: i32 = getOutputIndexFromCoords(${h});`, this.outputStrideSnippet = `idx0 = idx0 + uniforms.outShapeStrides${l === 0 ? "" : `[${l}]`}; `, t ? this.outputShape = t.map((c) => e[c]) : this.outputShape = e;
  }
  useScale() {
    this.addUniform("scale: f32"), this.scale = !0, this.shaderKey += "_scaled";
  }
  useScaleA() {
    this.addUniform("scaleA: f32"), this.scaleA = !0, this.shaderKey += "_scaledA";
  }
  useScaleB() {
    this.addUniform("scaleB: f32"), this.scaleB = !0, this.shaderKey += "_scaledB";
  }
  useActivation(e) {
    this.activation = e, this.shaderKey += `_${e}`;
  }
  useCausalMask() {
    this.causalMask = !0, this.addUniform("pastLen: i32"), this.shaderKey += "_causalMask";
  }
  activationSnippet() {
    return this.activation === "gelu" ? `
                // TODO: revisit after https://github.com/gpuweb/gpuweb/issues/4458 is resolved
                fn tanhComplete(x: vec4<f32>) -> vec4<f32> {
                    return vec4<f32>(
                        select(tanh(x.x), sign(x.x), abs(x.x) > 15.0f),
                        select(tanh(x.y), sign(x.y), abs(x.y) > 15.0f),
                        select(tanh(x.z), sign(x.z), abs(x.z) > 15.0f),
                        select(tanh(x.w), sign(x.w), abs(x.w) > 15.0f),
                    );
                }
                fn activation(x : vec4<f32>) -> vec4<f32> {
                    let x3 = x * x * x;
                    var inner = fma(vec4<f32>(${0.044715}f), x3, x);
                    inner = ${0.7978845608028654}f * inner;
                    inner = tanhComplete(inner);
                    inner = 0.5f * (1.0f + inner);
                    return x * inner;
                }
                ` : this.activation === "relu2" ? `
                fn activation(x : vec4<f32>) -> vec4<f32> {
                    let y = max(x, vec4<f32>(0.0));
                    return y * y;
                }
                ` : this.activation === "relu" ? `
                fn activation(x : vec4<f32>) -> vec4<f32> {
                    return max(x, vec4<f32>(0.0));
                }
                ` : "";
  }
  /* Transpose when writing to shared memory */
  readASnippet() {
    const e = `
            var col = i32(localId.x);
            var row = i32(localId.y) * 4;
            var packedA: vec2<i32> = A[offsetA + row * strideA + col];
            var Arow1 = vec4<f32>(
                unpack2x16float(u32(packedA.x)),
                unpack2x16float(u32(packedA.y))
            );
            packedA = A[offsetA + (row + 1) * strideA + col];
            var Arow2 = vec4<f32>(
                unpack2x16float(u32(packedA.x)),
                unpack2x16float(u32(packedA.y))
            );
            packedA = A[offsetA + (row + 2) * strideA + col];
            var Arow3 = vec4<f32>(
                unpack2x16float(u32(packedA.x)),
                unpack2x16float(u32(packedA.y))
            );
            packedA = A[offsetA + (row + 3) * strideA + col];
            var Arow4 = vec4<f32>(
                unpack2x16float(u32(packedA.x)),
                unpack2x16float(u32(packedA.y))
            );
            
            ${this.scaleA ? "Arow1 = Arow1 * uniforms.scaleA;" : ""}
            ${this.scaleA ? "Arow2 = Arow2 * uniforms.scaleA;" : ""}
            ${this.scaleA ? "Arow3 = Arow3 * uniforms.scaleA;" : ""}
            ${this.scaleA ? "Arow4 = Arow4 * uniforms.scaleA;" : ""}
        `;
    return this.transposeA ? `{
                ${e}
                mm_Asub[row][col] = Arow1;
                mm_Asub[row + 1][col] = Arow2;
                mm_Asub[row + 2][col] = Arow3;
                mm_Asub[row + 3][col] = Arow4;
        }` : `{
                ${e}
                
                col = i32(localId.x) * 4;
                row = i32(localId.y);
                
                mm_Asub[col][row] = vec4<f32>(Arow1.x, Arow2.x, Arow3.x, Arow4.x);
                mm_Asub[col + 1][row] = vec4<f32>(Arow1.y, Arow2.y, Arow3.y, Arow4.y);
                mm_Asub[col + 2][row] = vec4<f32>(Arow1.z, Arow2.z, Arow3.z, Arow4.z);
                mm_Asub[col + 3][row] = vec4<f32>(Arow1.w, Arow2.w, Arow3.w, Arow4.w);
        }`;
  }
  /* Transpose when writing to shared memory */
  readBSnippet() {
    const e = `
            var col = i32(localId.x);
            var row = i32(localId.y) * 4;
            var packedB: vec2<i32> = B[offsetB + row * strideB + col];
            var Brow1 = vec4<f32>(
                unpack2x16float(u32(packedB.x)),
                unpack2x16float(u32(packedB.y))
            );
            packedB = B[offsetB + (row + 1) * strideB + col];
            var Brow2 = vec4<f32>(
                unpack2x16float(u32(packedB.x)),
                unpack2x16float(u32(packedB.y))
            );
            packedB = B[offsetB + (row + 2) * strideB + col];
            var Brow3 = vec4<f32>(
                unpack2x16float(u32(packedB.x)),
                unpack2x16float(u32(packedB.y))
            );
            packedB = B[offsetB + (row + 3) * strideB + col];
            var Brow4 = vec4<f32>(
                unpack2x16float(u32(packedB.x)),
                unpack2x16float(u32(packedB.y))
            );
            
            ${this.scaleB ? "Brow1 = Brow1 * uniforms.scaleB;" : ""}
            ${this.scaleB ? "Brow2 = Brow2 * uniforms.scaleB;" : ""}
            ${this.scaleB ? "Brow3 = Brow3 * uniforms.scaleB;" : ""}
            ${this.scaleB ? "Brow4 = Brow4 * uniforms.scaleB;" : ""}
        `;
    return this.transposeB ? `{
                ${e}
                
                col = i32(localId.x) * 4;
                row = i32(localId.y);
                
                mm_Bsub[col][row] = vec4<f32>(Brow1.x, Brow2.x, Brow3.x, Brow4.x);
                mm_Bsub[col + 1][row] = vec4<f32>(Brow1.y, Brow2.y, Brow3.y, Brow4.y);
                mm_Bsub[col + 2][row] = vec4<f32>(Brow1.z, Brow2.z, Brow3.z, Brow4.z);
                mm_Bsub[col + 3][row] = vec4<f32>(Brow1.w, Brow2.w, Brow3.w, Brow4.w);
            }` : `{
                ${e}
                mm_Bsub[row][col] = Brow1;
                mm_Bsub[row + 1][col] = Brow2;
                mm_Bsub[row + 2][col] = Brow3;
                mm_Bsub[row + 3][col] = Brow4;
            }`;
  }
  baseIndexSnippets() {
    const e = `
            let strideA = uniforms.aShape.z / 2;
            let strideB = uniforms.bShape.z / 2;
        `;
    let t = "";
    this.transposeB ? t = "let baseB = getIndexFromCoords3D(vec3<i32>(batchB, globalColStart, 0), vec3<i32>(uniforms.bShape.x, uniforms.bShape.y, strideB));" : t = "let baseB = getIndexFromCoords3D(vec3<i32>(batchB, 0, globalColStart / 4), vec3<i32>(uniforms.bShape.x, uniforms.bShape.y, strideB));";
    let o = "";
    return this.transposeA ? o = "let baseA = getIndexFromCoords3D(vec3<i32>(batchA, 0, globalRowStart / 4), vec3<i32>(uniforms.aShape.x, uniforms.aShape.y, strideA));" : o = "let baseA = getIndexFromCoords3D(vec3<i32>(batchA, globalRowStart, 0), vec3<i32>(uniforms.aShape.x, uniforms.aShape.y, strideA));", `
            ${e}
            ${o}
            ${t}
        `;
  }
  offsetSnippets() {
    let e = "";
    this.transposeA ? e = "let offsetA = baseA + kStart * strideA;" : e = "let offsetA = baseA + kStart / 4;";
    let t = "";
    return this.transposeB ? t = "let offsetB = baseB + kStart / 4;" : t = "let offsetB = baseB + kStart * strideB;", `
            ${e}
            ${t}
        `;
  }
  getUserCode() {
    const e = this.transposeA, t = this.tileInner, o = this.workgroupSize[1] * 4, a = this.workgroupSize[0] * 4, i = e ? o : t, r = e ? t : o, s = this.dimInner, n = Math.ceil(s / t);
    return `
            var<workgroup> mm_Asub : array<array<vec4<f32>, ${i / 4 + (this.transposeA ? 0 : 1)}>, ${r}>;
            var<workgroup> mm_Bsub : array<array<vec4<f32>, ${a / 4 + (this.transposeB ? 1 : 0)}>, ${t}>;

            ${this.activation ? this.activationSnippet() : ""}

            ${p()} {
                let batch = i32(globalId.z);
                let batchA = ${this.broadcastBatch ? "batch % uniforms.aShape[0]" : "batch"};
                let batchB = ${this.broadcastBatch ? "batch % uniforms.bShape[0]" : "batch"};
                var kStart = 0;
                let localRow = i32(localId.y);
                let localCol = i32(localId.x);
                let globalRowStart = i32(workgroupId.y) * ${o};
                let globalColStart = i32(workgroupId.x) * ${a};

                // 4 rows x 4 cols accumulator
                // acc[i] holds row i (4 cols)
                var acc = array<vec4<f32>, 4>(
                    vec4<f32>(0.0), vec4<f32>(0.0), vec4<f32>(0.0), vec4<f32>(0.0)
                );

                ${this.baseIndexSnippets()}

                for (var t = 0; t < ${n}; t++) {
                    ${this.offsetSnippets()}

                    ${this.readASnippet()}
                    ${this.readBSnippet()}

                    kStart = kStart + ${t};
                    workgroupBarrier();

                    for (var k = 0; k < ${t}; k++) {
                        // Load 4 columns of B as a vec4
                        let bVec = mm_Bsub[k][localCol];
                        let aVec = mm_Asub[k][localRow];

                        // Compute 4 rows
                        for (var r = 0; r < 4; r = r + 1) {
                            acc[r] = fma(vec4<f32>(aVec[r]), bVec, acc[r]);
                        }
                    }
                    workgroupBarrier();
                }

                // Write out 4 rows x 2 packed cols (4 unpacked cols)
                let gRow = globalRowStart + localRow * 4;
                let gColPacked = i32(workgroupId.x) * ${this.workgroupSize[0] * 2} + localCol * 2;
                
                ${this.outputIndexSnippet}
                for (var i = 0; i < 4; i = i + 1) {
                    ${this.causalMask ? `
                    // Causal Masking: mask if col > row + pastLen
                    let r = gRow + i;
                    let cBase = gColPacked * 2;
                    let cVec = vec4<i32>(cBase, cBase + 1, cBase + 2, cBase + 3);
                    let mask = cVec > vec4<i32>(r + uniforms.pastLen);
                    acc[i] = select(acc[i], vec4<f32>(-uniforms.INFINITY), mask);
                    ` : ""}

                    ${this.activation ? "acc[i] = activation(acc[i]);" : ""}
                    ${this.scale ? "acc[i] = acc[i] * uniforms.scale;" : ""}
                    result[idx0 / 2] = vec2<i32>(
                        i32(pack2x16float(acc[i].xy)),
                        i32(pack2x16float(acc[i].zw))
                    );

                    ${this.outputStrideSnippet}
                }
            }
        `;
  }
}
export {
  B as default
};
