import { reduce as g, ReduceProgram as h } from "./utils/reductions.js";
import { a as w, N as S } from "../../index-DSGwv2Yx.js";
import k from "./utils/deviceInfo.js";
class v extends h {
  shaderKey = "norm2";
  atomic = !0;
  utilityFunctions = `
        fn atomicAddF32(sum: ptr<storage, atomic<i32>, read_write>, value: f32) -> f32 {
            var old = atomicLoad(sum);
            loop {
                let new_value = value + bitcast<f32>(old);
                let exchange_result = atomicCompareExchangeWeak(sum, old, bitcast<i32>(new_value));
                if (exchange_result.exchanged) {
                    return new_value;
                }
                old = exchange_result.old_value;
            }
        }
    `;
  constructor(o, r, i) {
    super(
      o,
      r,
      {
        reductionOp: "sum",
        elementwise: !1,
        forceWorkgroupSize: i
      },
      !1
    ), this.uniforms += "lossScaling: f32, index: i32";
  }
  getPreprocessSnippet() {
    return `
            candidate = candidate * uniforms.lossScaling;
            candidate = candidate * candidate;
        `;
  }
  getWriteSnippet() {
    return `
            if (tid == 0) {
                atomicAddF32(&result[uniforms.index], bestValue);
            }
        `;
  }
}
function x(t) {
  const { x: o, output: r } = t.inputs, { invLossScaling: i, index: c } = t.attrs, n = t.backend, u = [], d = k(n);
  let e = Math.min(512, n.device.limits.maxComputeWorkgroupSizeX);
  const a = 4, s = S(o.shape);
  for (; s % (e * a) !== 0 && e > 1; )
    e /= 2;
  if (e === 1)
    throw new Error(`Cannot find suitable workgroup size for Norm2Program with reduce size ${s}`);
  const l = {
    inSize: e * a,
    outSize: 1,
    batchSize: s / (e * a),
    windowSize: e
  }, m = new v(d, l, e), p = g(
    m,
    [o],
    n,
    [
      { type: "float32", data: [i] },
      { type: "int32", data: [c] }
    ],
    r
  );
  return u.forEach((f) => f.dispose()), p;
}
const z = {
  kernelName: "Norm2",
  backendName: "webgpu",
  kernelFunc: x
};
w(z);
