import { e as o } from "../../webgpu_program-CbjdYLYk.js";
import { f as u } from "../../webgpu_util-DuofJBMo.js";
class i {
  variableNames = ["logits"];
  outputShape;
  shaderKey;
  dispatchLayout;
  dispatch;
  workgroupSize;
  minSubgroupSize;
  maxSubgroupSize;
  subgroups = !0;
  subgroupBuiltins = !1;
  constructor(e, a, r) {
    this.minSubgroupSize = a, this.maxSubgroupSize = r, this.outputShape = e, this.dispatchLayout = u(this.outputShape), this.dispatch = [this.outputShape[0], 1, 1], this.workgroupSize = [Math.min(64, r), 1, 1], a !== r && (this.subgroupBuiltins = !0, this.workgroupSize = [64, 1, 1]), this.shaderKey = "softmax16subgroup";
  }
  getUserCode() {
    const e = this.maxSubgroupSize !== this.minSubgroupSize;
    return `
        ${e ? `var<workgroup> bestValues : array<f32, ${this.workgroupSize[0]}>;` : ""}
        const blockSize = ${this.workgroupSize[0]};
        ${o("index")} {
            let row = index / blockSize;
            let tid = i32(localId.x);
            let cols = uniforms.outShape[1];
            let rowIdx = row * cols;

            var threadMax = -3.402823e+38f;
            for (var col = tid; col < cols; col += blockSize) {
                let value = unpack2x16float(u32(logits[rowIdx + col]));
                threadMax = max(threadMax, max(value.x, value.y));
            }

            threadMax = subgroupMax(threadMax);
            ${e ? `
                let lane = localId.x % subgroupSize;
                if (lane == 0) {
                    bestValues[localId.x / subgroupSize] = threadMax;
                }
                workgroupBarrier();
                let numSubgroups = blockSize / subgroupSize;
                threadMax = select(-3.402823e+38f, bestValues[lane], lane < numSubgroups);
                threadMax = subgroupMax(threadMax);
                workgroupBarrier();    
            ` : ""}

            var threadSum = 0.0f;
            for (var col = tid; col < cols; col += blockSize) {
                let value = unpack2x16float(u32(logits[rowIdx + col]));
                let subExp = exp(value - threadMax);
                threadSum += subExp.x + subExp.y;
            }

            threadSum = subgroupAdd(threadSum);
            ${e ? `
                if (lane == 0) {
                    bestValues[localId.x / subgroupSize] = threadSum;
                }
                workgroupBarrier();
                threadSum = select(0.0f, bestValues[lane], lane < numSubgroups);
                threadSum = subgroupAdd(threadSum);    
            ` : ""}

            for (var col = tid; col < cols; col += ${e ? "i32(subgroupSize)" : "blockSize"}) {
                let value = unpack2x16float(u32(logits[rowIdx + col]));
                let valuePair: vec2<f32> = exp(value - threadMax) / threadSum;
                result[rowIdx + col] = i32(pack2x16float(valuePair));
            }
        }
      `;
  }
}
export {
  i as default
};
