Function atan2_approx

src/jit/impl/mlir/ir/numerical.cpp:27–66

polynomial approximation of arctangent atan(t) = t + c3 * t^3 + c5 * t^5 + ... + c17 * t^17 original paper: https://arxiv.org/pdf/1508.03211.pdf

// original paper:
// https://arxiv.org/pdf/1508.03211.pdf
mlir::Value atan2_approx(ValueBuilderHelper& helper, mlir::Value y, mlir::Value x) {
    auto atan_poly = [&](mlir::Value t) {
        std::vector<mlir::Value> coeff = {
                helper.const_f32(2.90188402868807315826416015625E-3),
                helper.const_f32(-1.62907354533672332763671875E-2),
                helper.const_f32(4.3082617223262786865234375E-2),
                helper.const_f32(-7.5408883392810821533203125E-2),
                helper.const_f32(0.1066047251224517822265625),
                helper.const_f32(-0.14209578931331634521484375),
                helper.const_f32(0.19993579387664794921875),
                helper.const_f32(-0.3333314359188079833984375)};
        auto t2 = helper.mul(t, t);
        auto p = polynomial(helper, t2, coeff);
        return helper.add(helper.mul(helper.mul(p, t2), t), t);
    };

    // constants
    auto zero = helper.const_f32(0);
    auto pi = helper.const_f32(3.141592653589793);
    auto pi_over_2 = helper.const_f32(1.570796326794897);

    // transform the angle into interval [0, pi/4]
    auto ax = helper.abs(x);
    auto ay = helper.abs(y);
    auto q = helper.div(helper.min(ax, ay), helper.max(ax, ay));

    // get approximation for interval [0, pi/4]
    auto r = atan_poly(q);

    // [0, pi/4] => [0, pi/2]
    r = helper.select(helper.le(ax, ay), helper.sub(pi_over_2, r), r);

    // [0, pi/2] => [0, pi]
    r = helper.select(helper.le(x, zero), helper.sub(pi, r), r);

    // [0, pi] => [-pi, pi]
    r = helper.select(helper.le(y, zero), helper.sub(zero, r), r);

    return r;
}

// numerical approximation of gauss error function
// https://en.wikipedia.org/wiki/Error_function#Polynomial