#version 450

#include "types.glsl"
#if defined(UNARY_MUL_FUSION)
#include "generic_binary_head.glsl"
#else
#include "generic_unary_head.glsl"
#endif

#if defined(UNARY_MUL_FUSION)
// OP on src1
layout(constant_id = 1) const bool op_on_b = false;
#endif

#if defined(UNARY_MUL_FUSION)
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
const uint num_threads = 256;
#else
layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
#endif

float op_abs(float x) {
    return abs(x);
}

float op_sgn(float x) {
    return sign(x);
}

float op_neg(float x) {
    return -x;
}

float op_sqr(float x) {
    return x * x;
}

float op_sqrt(float x) {
    return sqrt(x);
}

float op_sin(float x) {
    return sin(x);
}

float op_cos(float x) {
    return cos(x);
}

float op_clamp(float x) {
    return clamp(x, p.param1, p.param2);
}

float op_leaky_relu(float x) {
    return max(x, 0.0f) + min(x, 0.0f) * p.param1;
}

float op_step(float x) {
    return x > 0.0f ? 1.0f : 0.0f;
}

float op_tanh(float x) {
    return 1.0f - 2.0f / (exp(2.0f*x) + 1.0f);
}

float op_elu(float x) {
    return x < 0.0f ? exp(x) - 1.0f : x;
}

float op_relu(float x) {
    return max(x, 0.0f);
}

float op_sigmoid(float x) {
    return 1.0f / (1.0f + exp(-x));
}

float op_gelu(float x) {
    const float GELU_COEF_A    = 0.044715f;
    const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
    const float val = SQRT_2_OVER_PI*x*(1.0f + GELU_COEF_A*x*x);
    return 0.5f*x*(2.0f - 2.0f / (exp(2.0f * val) + 1.0f));
}

float op_gelu_quick(float x) {
    const float GELU_QUICK_COEF = -1.702f;
    return x * (1.0f / (1.0f + exp(GELU_QUICK_COEF * x)));
}

float op_silu(float x) {
    return x / (1.0f + exp(-x));
}

float op_hardswish(float x) {
    return x * min(1.0f, max(0.0f, (x + 3.0f) / 6.0f));
}

float op_hardsigmoid(float x) {
    return min(1.0f, max(0.0f, (x + 3.0f) / 6.0f));
}

float op_exp(float x) {
    return exp(x);
}

float op_expm1(float x) {
    // exp(x) - 1 loses many ulps to cancellation near zero.  Use a degree-6
    // Taylor expansion for |x| <= 1/4: the omitted x^7/5040 term is < 1.3e-8,
    // about 0.5 ulp at expm1(0.25), and a host-side f32 model stays within
    // 2 ulps over the interval.  The first native exp(x)-1 values outside the
    // cutoff are about 1 ulp for +0.25 and 2 ulps for -0.25.
    if (abs(x) <= 0.25f) {
        return x * (1.0f + x * (0.5f + x * ((1.0f/6.0f) + x * ((1.0f/24.0f) + x * ((1.0f/120.0f) + x * (1.0f/720.0f))))));
    }
    return exp(x) - 1.0f;
}

float op_softplus(float x) {
    return (x > 20.0f) ? x : log(1.0f + exp(x));
}

float op_gelu_erf(float a) {
    // based on Abramowitz and Stegun formula 7.1.26 or similar Hastings' approximation
    const float p_erf  = 0.3275911f;
    const float a1_erf = 0.254829592f;
    const float a2_erf = -0.284496736f;
    const float a3_erf = 1.421413741f;
    const float a4_erf = -1.453152027f;
    const float a5_erf = 1.061405429f;

    const float SQRT_2_INV = 0.70710678118654752440084436210484f;
    const float a_div_sqr2 = a * SQRT_2_INV;
    const float sign_x = sign(a_div_sqr2);
    const float x = abs(a_div_sqr2);
    const float t = 1.0f / (1.0f + p_erf * x);
    const float y = 1.0f - (((((a5_erf * t + a4_erf) * t) + a3_erf) * t + a2_erf) * t + a1_erf) * t * exp(-x * x);
    return 0.5f * a * (1.0f + sign_x * y);
}

#if !defined(UNARY_MUL_FUSION)
float op_xielu(float x) {
    const float alpha_n = p.param1;
    const float alpha_p = p.param2;
    const float beta = p.param3;
    const float eps = p.param4;

    if (x > 0.0f) {
        return alpha_p * x * x + beta * x;
    }

    const float min_x_eps = min(x, eps);
    return (op_expm1(min_x_eps) - x) * alpha_n + beta * x;
}
#endif

float op_floor(float x) {
    return floor(x);
}

float op_ceil(float x) {
    return ceil(x);
}

float op_round(float x) {
    // Round halfway cases away from zero as roundf does.
    return x >= 0.0f ? floor(x + 0.5f) : ceil(x - 0.5f);
}

float op_trunc(float x) {
    return trunc(x);
}

void main() {
    uint idx = get_idx();

#if defined(UNARY_MUL_FUSION)
    // keep total threads at 512
    [[unroll]] for (uint iter = 0; iter < 2; ++iter) {
        if (idx >= p.ne) {
            continue;
        }
        uint i00, i01, i02, i03;
        get_indices(idx, i00, i01, i02, i03);

        if (op_on_b) {
            data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] =
                D_TYPE(FLOAT_TYPE(OP(float(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)]))) * FLOAT_TYPE(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]));
        } else {
            data_d[get_doffset() + dst_idx(i00, i01, i02, i03)] =
                D_TYPE(FLOAT_TYPE(OP(float(data_a[get_aoffset() + src0_idx(i00, i01, i02, i03)]))) * FLOAT_TYPE(data_b[get_boffset() + src1_idx(i00, i01, i02, i03)]));
        }

        idx += num_threads;
    }
#else
    if (idx >= p.ne) {
        return;
    }

    const uint a_idx = get_aoffset() + src0_idx(idx);
    const uint d_idx = get_doffset() + dst_idx(idx);

    data_d[d_idx] = D_TYPE(OP(float(data_a[a_idx])));
#endif
}