import { a as d } from "../../index-DSGwv2Yx.js";
import { e as s } from "../../webgpu_program-CbjdYLYk.js";
import { f as n, c as o } from "../../webgpu_util-DuofJBMo.js";
import { isPackedTensor as l } from "../../utilities/packed.js";
const u = 0.7978845608028654, r = 0.044715;
class c {
  outputShape;
  shaderKey;
  dispatchLayout;
  dispatch;
  variableNames = ["A"];
  workgroupSize;
  size = !0;
  constructor(e) {
    this.workgroupSize = [128, 1, 1], this.outputShape = e, this.dispatchLayout = n(this.outputShape), this.dispatch = o(this.dispatchLayout, this.outputShape, this.workgroupSize), this.shaderKey = "unary_gelu";
  }
  getUserCode() {
    return `
            // TODO: revisit after https://github.com/gpuweb/gpuweb/issues/4458 is resolved
            fn tanhComplete(x: f32) -> f32 {
                return select(tanh(x), sign(x), abs(x) > 15.0);
            }
            fn unaryOperation(x : f32) -> f32 {
                let x3 = x * x * x;
                var inner = fma(${r}, x3, x);
                inner = ${u} * inner;
                inner = tanhComplete(inner);
                inner = 0.5 * (1.0 + inner);
                return x * inner;
            }
            ${s("index")} {
                if (index < uniforms.size) {
                let a = getAByOutputIndex(index);
                setOutputAtIndex(index, unaryOperation(a));
                }
            }
      `;
  }
}
function x(t) {
  const { x: e } = t.inputs, a = t.backend, i = new c(e.shape);
  return a.runWebGPUProgram(i, [e], "float32");
}
const g = {
  kernelName: "Gelu",
  backendName: "webgpu",
  kernelFunc: x
};
d(g);
class f {
  // Inputs: dy, x
  variableNames = ["dy", "x"];
  outputShape;
  shaderKey = "GeluGrad";
  dispatchLayout;
  dispatch;
  workgroupSize = [128, 1, 1];
  size = !0;
  constructor(e) {
    this.outputShape = e, this.dispatchLayout = n(this.outputShape), this.dispatch = o(this.dispatchLayout, this.outputShape, this.workgroupSize);
  }
  getUserCode() {
    return `
            // TODO: revisit after https://github.com/gpuweb/gpuweb/issues/4458 is resolved
            fn tanhComplete(x: f32) -> f32 {
                return select(tanh(x), sign(x), abs(x) > 15.0);
            }
            fn activationGrad(dy: f32, X: f32) -> f32 {
                let x2 = X * X;
                let x3 = x2 * X;
                let u  = ${u} * (X + ${r} * x3);
                let t  = tanhComplete(u);
                let sech2 = 1.0 - t * t;
                let du_dx = ${u} * (1.0 + 3.0 * ${r} * x2);
                let dgelu = 0.5 * (1.0 + t) + 0.5 * X * sech2 * du_dx;
                return dy *dgelu;
            }
            ${s("index")} {
                if (index < uniforms.size) {
                    let X  = unpack2x16float(u32(x[index]));
                    let DY = unpack2x16float(u32(dy[index]));
                    let dgelu = vec2<f32>(
                        activationGrad(DY.x, X.x),
                        activationGrad(DY.y, X.y)
                    );
                    result[index] = i32(pack2x16float(dgelu));
                }
            }`;
  }
}
class m {
  // Inputs: dy, x
  variableNames = ["dy", "x"];
  outputShape;
  shaderKey = "GeluGrad";
  dispatchLayout;
  dispatch;
  workgroupSize = [128, 1, 1];
  size = !0;
  constructor(e) {
    this.outputShape = e, this.dispatchLayout = n(this.outputShape), this.dispatch = o(this.dispatchLayout, this.outputShape, this.workgroupSize);
  }
  getUserCode() {
    return `
            // TODO: revisit after https://github.com/gpuweb/gpuweb/issues/4458 is resolved
            fn tanhComplete(x: f32) -> f32 {
                return select(tanh(x), sign(x), abs(x) > 15.0);
            }
            fn activationGrad(dy: f32, X: f32) -> f32 {
                let x2 = X * X;
                let x3 = x2 * X;
                let u  = ${u} * (X + ${r} * x3);
                let t  = tanhComplete(u);
                let sech2 = 1.0 - t * t;
                let du_dx = ${u} * (1.0 + 3.0 * ${r} * x2);
                let dgelu = 0.5 * (1.0 + t) + 0.5 * X * sech2 * du_dx;
                return dy *dgelu;
            }
            ${s("index")} {
                if (index < uniforms.size) {
                    let X  = getXByOutputIndex(index);
                    let DY = getDyByOutputIndex(index);
                    let dgelu = activationGrad(DY, X);
                    setOutputAtIndex(index, dgelu);
                }
            }`;
  }
}
function y(t) {
  const { dy: e, x: a } = t.inputs, i = t.backend, p = l(e), h = p ? new f(a.shape) : new m(a.shape);
  return i.runWebGPUProgram(h, [e, a], p ? "packedF16" : "float32");
}
const b = {
  kernelName: "GeluGrad",
  backendName: "webgpu",
  kernelFunc: y
};
d(b);
export {
  c as GeluProgram
};
