import { a as $, a6 as y, e as L } from "../../index-DSGwv2Yx.js";
import { createReduceInfo as M } from "./utils/reductions.js";
import { f as w } from "../../webgpu_util-DuofJBMo.js";
import { e as _ } from "../../webgpu_program-CbjdYLYk.js";
import { p as x, u as R } from "../../pack16-WlOSOuZA.js";
import { isPackedTensor as h } from "../../utilities/packed.js";
import { reshape16 as z } from "../reshape16.js";
import { sum16 as N } from "../sum16.js";
import { slice16 as v } from "../slice16.js";
class Y {
  outputShape;
  shaderKey = "RMSNormGrad";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  variableNames = ["x", "gamma", "dy"];
  uniforms = "reduceSize : i32, batchSize: i32";
  inputShape;
  size = !1;
  rowsPerWorkgroup;
  packed = !1;
  outputComponent;
  constructor(e, a = 4, t = !1) {
    if (this.packed = t, this.shaderKey = `RMSNormGrad_${a}`, this.rowsPerWorkgroup = a, this.inputShape = [e.batchSize, e.inSize], this.outputShape = [e.batchSize + e.batchSize / this.rowsPerWorkgroup, e.inSize], this.dispatchLayout = w(this.outputShape), this.dispatch = [e.batchSize / this.rowsPerWorkgroup, 1, 1], e.batchSize % this.rowsPerWorkgroup !== 0)
      throw new Error(
        `RMSNormGradProgram: batch size ${e.batchSize} must be divisible by rowsPerWorkgroup ${this.rowsPerWorkgroup}`
      );
    if (e.inSize > 1024)
      throw new Error(`RMSNormGradProgram: inSize ${e.inSize} exceeds max of 1024`);
  }
  getUserCode() {
    const e = this.workgroupSize[0], a = this.rowsPerWorkgroup, t = `
          var<workgroup> partials : array<vec2<f32>, ${e}>;
          var<workgroup> accumulation: array<${this.packed ? "vec2<f32>" : "f32"}, 1024>;
        `, u = this.packed ? `
                let X = unpack2x16float(u32(x[offset + k]));
                let DY = unpack2x16float(u32(dy[offset + k]));
                let G  = unpack2x16float(u32(gamma[k]));
                sum_x2 = fma(X.x, X.x, sum_x2);
                sum_x2 = fma(X.y, X.y, sum_x2);
                sum_dygx = fma(DY.x * G.x, X.x, sum_dygx);
                sum_dygx = fma(DY.y * G.y, X.y, sum_dygx);
        ` : `
                let X = f32(x[offset + k]);
                let DY = f32(dy[offset + k]);
                let G  = f32(gamma[k]);
                sum_x2 = fma(X, X, sum_x2);
                sum_dygx = fma(DY * G, X, sum_dygx);
            `, m = this.packed ? `
                let X  = unpack2x16float(u32(x[offset + k]));
                let DY = unpack2x16float(u32(dy[offset + k]));
                let G  = unpack2x16float(u32(gamma[k]));

                let dyGamma = DY * G;
                let dx = vec2<f32>(
                    fma(dyGamma.x, invRMS, -X.x * scale),
                    fma(dyGamma.y, invRMS, -X.y * scale)
                );

                result[offset + k] = i32(pack2x16float(dx));

                // dGamma
                accumulation[k] = fma(DY, X * invRMS, accumulation[k]);
        ` : `
                let X  = f32(x[offset + k]);
                let DY = f32(dy[offset + k]);
                let G  = f32(gamma[k]);

                let dyGamma = DY * G;
                let dx = fma(dyGamma, invRMS, -X * scale);

                result[offset + k] = dx;

                // dGamma
                accumulation[k] = fma(DY, X * invRMS, accumulation[k]);
        `, n = this.packed ? `
                result[outDgBase + k] = i32(pack2x16float(accumulation[k]));
        ` : `
                result[outDgBase + k] = accumulation[k];
        `;
    return `
          fn DIV_CEIL(a : u32, b : u32) -> u32 {
            return ((a - 1u) / b + 1u);
          }

          ${t}

          ${_("index")} {
            // One workgroup per row (batch).
            let Length = uniforms.reduceSize;
            let BatchSize = uniforms.batchSize;
            for (var k = i32(localId.x); k < Length; k = k + ${e}) {
                accumulation[k] = ${this.packed ? "vec2<f32>(0.0f)" : "0.0f"};
            }

            for (var rowOff = 0; rowOff < ${a}; rowOff = rowOff + 1) {
                let row = i32(workgroupId.x) * ${a} + rowOff;
                let offset = row * Length;
                
                var sum_x2 = 0.0f;
                var sum_dygx = 0.0f;

                for (var k = i32(localId.x); k < Length; k = k + ${e}) {
                    ${u}
                }

                partials[localId.x] = vec2<f32>(sum_x2, sum_dygx);
                workgroupBarrier();

                var reduceSize = min(u32(Length), ${e}u);
                for (var currentSize = reduceSize / 2u; reduceSize > 1u; currentSize = reduceSize / 2u) {
                    let interval = DIV_CEIL(reduceSize, 2u);
                    if (localId.x < currentSize) {
                        partials[localId.x] = partials[localId.x] + partials[localId.x + interval];
                    }
                    reduceSize = interval;
                    workgroupBarrier();
                }

                let invN = 1.0f / f32(${this.packed ? "Length * 2" : "Length"});
                let mean_x2   = fma(partials[0].x, invN, 1e-8);
                let mean_dygx = partials[0].y * invN;

                let invRMS = inverseSqrt(mean_x2);
                let scale = (mean_dygx / (mean_x2)) * invRMS;

                // write dx and dGamma.
                for (var k = i32(localId.x); k < Length; k = k + ${e}) {
                    ${m}
                }

                workgroupBarrier();
            }

            // Write out the partially accumulated dGamma
            let outDgBase = BatchSize * Length + i32(workgroupId.x) * Length;
            for (var k = i32(localId.x); k < Length; k = k + ${e}) {
                ${n}
            }
          }
        `;
  }
}
class I {
  outputShape;
  shaderKey = "RMSNormGrad";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  variableNames = ["x", "dy"];
  uniforms = "reduceSize : i32, batchSize: i32";
  inputShape;
  size = !1;
  packed = !1;
  outputComponent;
  constructor(e, a = !1) {
    this.packed = a, this.shaderKey = "RMSNormGrad_NoGamma", this.inputShape = [e.batchSize, e.inSize], this.outputShape = [e.batchSize, e.inSize], this.dispatchLayout = w(this.outputShape), this.dispatch = [e.batchSize, 1, 1];
  }
  getUserCode() {
    const e = this.workgroupSize[0], a = `
          var<workgroup> partials : array<vec2<f32>, ${e}>;
        `, t = this.packed ? `
                let X = unpack2x16float(u32(x[offset + k]));
                let DY = unpack2x16float(u32(dy[offset + k]));
                sum_x2 = fma(X.x, X.x, sum_x2);
                sum_x2 = fma(X.y, X.y, sum_x2);
                sum_dygx = fma(DY.x, X.x, sum_dygx);
                sum_dygx = fma(DY.y, X.y, sum_dygx);
        ` : `
                let X = f32(x[offset + k]);
                let DY = f32(dy[offset + k]);
                sum_x2 = fma(X, X, sum_x2);
                sum_dygx = fma(DY, X, sum_dygx);
            `, u = this.packed ? `
                let X  = unpack2x16float(u32(x[offset + k]));
                let DY = unpack2x16float(u32(dy[offset + k]));

                let dx = vec2<f32>(
                    fma(DY.x, invRMS, -X.x * scale),
                    fma(DY.y, invRMS, -X.y * scale)
                );

                result[offset + k] = i32(pack2x16float(dx));
        ` : `
                let X  = f32(x[offset + k]);
                let DY = f32(dy[offset + k]);

                let dx = fma(DY, invRMS, -X * scale);

                result[offset + k] = dx;
        `;
    return `
          fn DIV_CEIL(a : u32, b : u32) -> u32 {
            return ((a - 1u) / b + 1u);
          }

          ${a}

          ${_("index")} {
            // One workgroup per row (batch).
            let Length = uniforms.reduceSize;
            let BatchSize = uniforms.batchSize;

                let row = i32(workgroupId.x);
                let offset = row * Length;
                
                var sum_x2 = 0.0f;
                var sum_dygx = 0.0f;

                for (var k = i32(localId.x); k < Length; k = k + ${e}) {
                    ${t}
                }

                partials[localId.x] = vec2<f32>(sum_x2, sum_dygx);
                workgroupBarrier();

                var reduceSize = min(u32(Length), ${e}u);
                for (var currentSize = reduceSize / 2u; reduceSize > 1u; currentSize = reduceSize / 2u) {
                    let interval = DIV_CEIL(reduceSize, 2u);
                    if (localId.x < currentSize) {
                        partials[localId.x] = partials[localId.x] + partials[localId.x + interval];
                    }
                    reduceSize = interval;
                    workgroupBarrier();
                }

                let invN = 1.0f / f32(${this.packed ? "Length * 2" : "Length"});
                let mean_x2   = fma(partials[0].x, invN, 1e-8);
                let mean_dygx = partials[0].y * invN;

                let invRMS = inverseSqrt(mean_x2);
                let scale = (mean_dygx / (mean_x2)) * invRMS;

                // write dx and dGamma.
                for (var k = i32(localId.x); k < Length; k = k + ${e}) {
                    ${u}
                }
          }
        `;
  }
}
function P(d) {
  const { dy: e, x: a, gamma: t } = d.inputs, u = 4;
  y(a.shape, e.shape, "Error in RMSNormGrad dy: ");
  const m = h(a), n = t ? h(t) : !1, f = h(e), r = m || n || f, i = !r || m ? a : x(a), s = !r || n || !t ? t : x(t), c = !r || f ? e : x(e);
  s && y(s.shape, [i.shape[i.shape.length - 1]], "Error in RMSNormGrad gamma: ");
  const G = d.backend, o = M(s ? [i, s, c] : [i, c], -1), k = t ? new Y(o, u, r) : new I(o, r), X = [
    { type: "int32", data: [k.inputShape[1]] },
    // Reduce size
    { type: "int32", data: [k.inputShape[0]] }
    // Batch size
  ];
  if (o.inSize > 1024)
    throw new Error(`rmsNormGradGPU: inSize ${o.inSize} exceeds max of 1024`);
  const D = G.runWebGPUProgram(
    k,
    s ? [i, s, c] : [i, c],
    r ? "packedF16" : "float32",
    X
  );
  r && !m && i.dispose(), r && !n && s && s.dispose(), r && !f && c.dispose();
  const l = L().makeTensorFromTensorInfo(D);
  if (t) {
    const p = v(l, [0, 0], [o.batchSize, o.inSize]), S = v(
      l,
      [o.batchSize, 0],
      [o.batchSize / u, o.inSize]
    );
    l.dispose();
    const b = z(p, a.shape);
    p.dispose();
    const g = N(S, [0]);
    return S.dispose(), [b, !r || n ? g : R(g)];
  } else {
    const p = z(l, a.shape);
    return l.dispose(), [p];
  }
}
const C = {
  kernelName: "RMSNormGrad",
  backendName: "webgpu",
  kernelFunc: P
};
$(C);
