import { ah as h, N as S, e as f } from "../../../index-DSGwv2Yx.js";
import { e as d } from "../../../webgpu_program-CbjdYLYk.js";
import { reshape16 as g } from "../../reshape16.js";
import { f as z } from "../../../webgpu_util-DuofJBMo.js";
import { c as k } from "../../../axis_util-QWWgLjut.js";
import { z as x } from "../../../zeros-Bw0puq_w.js";
function c(e, u, t, i) {
  return e && !u ? `
            bestValue = subgroupAdd(bestValue);
        ` : e ? `
            bestValue = subgroupAdd(bestValue);
            let lane = localId.x % subgroupSize;
            if (lane == 0) {
                bestValues[localId.x / subgroupSize] = bestValue;
            }
            workgroupBarrier();
            let numSubgroups = ${t} / subgroupSize;
            bestValue = select(${i ? "vec2<f32>(0.0f)" : "0.0f"}, bestValues[lane], lane < numSubgroups);
            bestValue = subgroupAdd(bestValue);
        ` : `
            bestValues[localId.x] = bestValue;
            workgroupBarrier();

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

            bestValue = bestValues[0];
        `;
}
function $(e) {
  const u = `${e.workgroupSizeX}`, t = e.subgroups && !e.variableSubgroups ? "" : `
             var<workgroup> bestValues : array<f32, ${e.workgroupSizeX}>;
           `, i = c(e.subgroups, e.variableSubgroups, e.workgroupSizeX, !1);
  return `
           fn DIV_CEIL(a : u32, b : u32) -> u32 {
            return ((a - 1u) / b + 1u);
           }

            fn readInput(index: i32) -> vec2<f32> {
                ${e.inputReadSnippet ? e.inputReadSnippet : `
                let packed = u32(x[index]);
                return unpack2x16float(packed);
                `}
            }

           ${t}
           ${e.utilityFunctions ?? ""}
    
           ${d("index")} {
                let outputIndex = index / ${u};
                let offset = outputIndex * uniforms.reduceSize;
                var bestValue = 0.0f;
                let Length = uniforms.reduceSize;
                let tid = i32(localId.x);
    
                for (var k = tid; k < Length;
                    k = k + ${u}) {
                    var candidate = readInput(offset + k);
                    ${e.inputSnippet}
                    bestValue = bestValue + candidate.x + candidate.y;
                }

                ${i}
                bestValue = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize * 2i)" : ""};

                ${e.reducedSnippet ? e.reducedSnippet : ""}

                for (var k = tid; k < Length;
                    k = k + ${u}) {
                    ${e.outputSnippet}
                }
           }
         `;
}
function v(e) {
  const u = `${e.workgroupSizeX}`, t = e.subgroups && !e.variableSubgroups ? "" : `
             var<workgroup> bestValues : array<vec2<f32>, ${e.workgroupSizeX}>;
           `, i = c(e.subgroups, e.variableSubgroups, e.workgroupSizeX, !0);
  return `
           fn DIV_CEIL(a : u32, b : u32) -> u32 {
            return ((a - 1u) / b + 1u);
           }

           fn readInput(index: i32) -> vec2<f32> {
                ${e.inputReadSnippet ? e.inputReadSnippet : `
                let packed = u32(x[index]);
                return unpack2x16float(packed);
                `}
            }

            ${t}
            ${e.utilityFunctions ?? ""}
    
           ${d("index")} {
                let outputIndex = index / ${u};
                let offset1 = outputIndex * 2 * uniforms.reduceSize;
                let offset2 = offset1 + uniforms.reduceSize;
                var bestValue = vec2<f32>(0.0f, 0.0f);
                let Length = uniforms.reduceSize;
                let tid = i32(localId.x);
    
                for (var k = tid; k < Length;
                    k = k + ${u}) {
                    var candidate = readInput(offset1 + k);
                    ${e.inputSnippet}
                    let bv1 = candidate.x + candidate.y;

                    candidate = readInput(offset2 + k);
                    ${e.inputSnippet}
                    let bv2 = candidate.x + candidate.y;

                    bestValue = bestValue + vec2<f32>(bv1, bv2);
                }
                ${i}
                bestValue = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize * 2i)" : ""};

                ${e.reducedSnippet ?? ""}
                ${e.outputSnippet}
           }
         `;
}
function V(e) {
  return e.elementwise ? $(e) : v(e);
}
function w(e) {
  const u = `${e.workgroupSizeX}`, t = e.subgroups && !e.variableSubgroups ? "" : `
             var<workgroup> bestValues : array<f32, ${e.workgroupSizeX}>;
           `, i = c(e.subgroups, e.variableSubgroups, e.workgroupSizeX, !1);
  return `
           fn DIV_CEIL(a : u32, b : u32) -> u32 {
            return ((a - 1u) / b + 1u);
           }

           fn readInput(index: i32) -> f32 {
                ${e.inputReadSnippet ? e.inputReadSnippet : `
                return x[index];
                `}
            }

           ${t}
              ${e.utilityFunctions ?? ""}
    
           ${d("index")} {
                let outputIndex = index / ${e.workgroupSizeX};
                let offset = outputIndex * uniforms.reduceSize;
                var bestValue = 0.0f;
                let Length = uniforms.reduceSize;
                let tid = i32(localId.x);
    
                for (var k = tid; k < Length;
                    k = k + ${e.workgroupSizeX}) {
                    var candidate = readInput(offset + k);
                    ${e.inputSnippet}
                    bestValue = bestValue + candidate;
                }
                ${i}

                bestValue = bestValue ${e.reductionOp === "mean" ? "/ f32(uniforms.reduceSize)" : ""};

                ${e.reducedSnippet}

                ${e.elementwise ? `for (var k = tid; k < Length;
                    k = k + ${u}) {
                    ${e.outputSnippet}
                }` : `${e.outputSnippet}`}
           }
         `;
}
function X(e, u) {
  const t = e[0], o = h(u, t.shape), [, n] = k(t.shape, o), s = S(n), r = S(t.shape) / s;
  return { windowSize: s, inSize: s, batchSize: r, outSize: r };
}
class P {
  atomic = !1;
  outputShape;
  shaderKey = "reduce16";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  variableNames = ["x"];
  uniforms = "reduceSize : i32,";
  inputShape;
  size = !1;
  packed = !0;
  outputComponent;
  variableComponents;
  elementwise;
  subgroups = !1;
  subgroupBuiltins = !1;
  deviceInfo;
  params;
  utilityFunctions;
  constructor(u, t, i, o) {
    this.params = i, this.inputShape = [t.batchSize, t.inSize], this.deviceInfo = u, this.packed = o;
    const n = i.forceWorkgroupSize ? i.forceWorkgroupSize : t.inSize % 64 === 0 ? 64 : 32;
    u.subgroupsSupported && !i.forceWorkgroupSize ? (this.workgroupSize = [Math.min(n, u.subgroupMaxSize), 1, 1], this.subgroups = !0, u.variableSubgroups && (this.subgroupBuiltins = !0)) : this.workgroupSize[0] = n, this.outputShape = i.elementwise ? [t.batchSize, t.inSize] : o ? [t.outSize / 2] : [t.outSize], this.dispatchLayout = z(this.outputShape), this.dispatch = [
      i.elementwise ? t.batchSize : o ? t.batchSize / 2 : t.batchSize,
      1,
      1
    ], this.outputComponent = 1, this.variableComponents = [1], this.elementwise = i.elementwise === !0;
  }
  getWriteSnippet() {
    return this.packed ? "result[outputIndex] = i32(pack2x16float(bestValue));" : "result[outputIndex] = bestValue;";
  }
  getPreprocessSnippet() {
    return "";
  }
  getPostprocessSnippet() {
    return "";
  }
  getReadSnippet() {
    return this.packed ? `
                let packed = u32(x[index]);
                return unpack2x16float(packed);
                ` : "return f32(x[index]);";
  }
  getUserCode() {
    const u = this.workgroupSize[0];
    return this.packed ? V({
      ...this.params,
      workgroupSizeX: u,
      subgroups: this.subgroups,
      variableSubgroups: this.deviceInfo.variableSubgroups,
      inputReadSnippet: this.getReadSnippet(),
      inputSnippet: this.getPreprocessSnippet(),
      outputSnippet: this.getWriteSnippet(),
      reducedSnippet: this.getPostprocessSnippet(),
      utilityFunctions: this.utilityFunctions
    }) : w({
      ...this.params,
      workgroupSizeX: u,
      subgroups: this.subgroups,
      variableSubgroups: this.deviceInfo.variableSubgroups,
      inputReadSnippet: this.getReadSnippet(),
      inputSnippet: this.getPreprocessSnippet(),
      outputSnippet: this.getWriteSnippet(),
      reducedSnippet: this.getPostprocessSnippet(),
      utilityFunctions: this.utilityFunctions
    });
  }
}
function A(e, u, t, i, o) {
  const n = u[0], a = [{ type: "int32", data: [e.inputShape[e.inputShape.length - 1]] }, ...i ?? []];
  let r = o;
  !o && e.atomic && (r = x(e.outputShape, "int32"));
  const l = t.runWebGPUProgram(
    e,
    u,
    e.packed ? "packedF16" : e.atomic ? "int32" : "float32",
    a,
    r
  ), p = f().makeTensorFromTensorInfo(l);
  if (e.outputShape.length === 1 && e.outputShape[0] <= 2)
    return p;
  const b = g(
    p,
    e.elementwise ? n.shape : e.packed ? [...n.shape.slice(0, -2), n.shape[n.shape.length - 2] / 2] : [...n.shape.slice(0, -2), n.shape[n.shape.length - 2]]
  );
  return p.dispose(), b;
}
export {
  P as ReduceProgram,
  X as createReduceInfo,
  A as reduce
};
