[Mesa-dev] [PATCH v2] nir/spirv/glsl450: rewrite atan2 to deal with infinities

[Juan A. Suarez Romero]

Rewrite atan2(y,x) to cover (+/-)INF values.

This fixes several test cases in Vulkan CTS
(dEQP-VK.glsl.builtin.precision.atan2.*)

v2: do not flush denorms to 0 (jasuarez)
---
 src/compiler/spirv/vtn_glsl450.c | 48 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/compiler/spirv/vtn_glsl450.c b/src/compiler/spirv/vtn_glsl450.c
index 0d32fddbef..d52a22c0c3 100644
--- a/src/compiler/spirv/vtn_glsl450.c
+++ b/src/compiler/spirv/vtn_glsl450.c
@@ -299,18 +299,47 @@ build_atan(nir_builder *b, nir_ssa_def *y_over_x)
    return nir_fmul(b, tmp, nir_fsign(b, y_over_x));
 }
 
+/*
+ * Computes atan2(y,x)
+ */
 static nir_ssa_def *
 build_atan2(nir_builder *b, nir_ssa_def *y, nir_ssa_def *x)
 {
    nir_ssa_def *zero = nir_imm_float(b, 0.0f);
-
-   /* If |x| >= 1.0e-8 * |y|: */
-   nir_ssa_def *condition =
-      nir_fge(b, nir_fabs(b, x),
-              nir_fmul(b, nir_imm_float(b, 1.0e-8f), nir_fabs(b, y)));
+   nir_ssa_def *inf = nir_imm_float(b, INFINITY);
+   nir_ssa_def *minus_inf = nir_imm_float(b, -INFINITY);
+   nir_ssa_def *m_3_pi_4 = nir_fmul(b, nir_imm_float(b, 3.0f),
+                                       nir_imm_float(b, M_PI_4f));
+
+   /* if y == +-INF */
+   nir_ssa_def *y_is_inf = nir_feq(b, nir_fabs(b, y), inf);
+
+   /* if x == +-INF */
+   nir_ssa_def *x_is_inf = nir_feq(b, nir_fabs(b, x), inf);
+
+   /* Case: y is +-INF */
+   nir_ssa_def *y_is_inf_then =
+      nir_fmul(b, nir_fsign(b, y),
+                  nir_bcsel(b, nir_feq(b, x, inf),
+                               nir_imm_float(b, M_PI_4f),
+                               nir_bcsel(b, nir_feq(b, x, minus_inf),
+                                            m_3_pi_4,
+                                            nir_imm_float(b, M_PI_2f))));
+
+   /* Case: x is +-INF */
+   nir_ssa_def *x_is_inf_then =
+      nir_fmul(b, nir_fsign(b, y),
+                  nir_bcsel(b, nir_feq(b, x, inf),
+                               zero,
+                               nir_imm_float(b, M_PIf)));
+
+   /* If x > 0 */
+   nir_ssa_def *x_not_zero =
+      nir_fne(b, x, zero);
 
    /* Then...call atan(y/x) and fix it up: */
    nir_ssa_def *atan1 = build_atan(b, nir_fdiv(b, y, x));
+
    nir_ssa_def *r_then =
       nir_bcsel(b, nir_flt(b, x, zero),
                    nir_fadd(b, atan1,
@@ -323,7 +352,14 @@ build_atan2(nir_builder *b, nir_ssa_def *y, nir_ssa_def *x)
    nir_ssa_def *r_else =
       nir_fmul(b, nir_fsign(b, y), nir_imm_float(b, M_PI_2f));
 
-   return nir_bcsel(b, condition, r_then, r_else);
+   /* Everything together */
+   return nir_bcsel(b, y_is_inf,
+                       y_is_inf_then,
+                       nir_bcsel(b, x_is_inf,
+                                    x_is_inf_then,
+                                    nir_bcsel(b, x_not_zero,
+                                                 r_then,
+                                                 r_else)));
 }
 
 static nir_ssa_def *
-- 
2.11.0