device-libs: Move special case check in rsqrt f64 implementation

Move the edge case check to be on the original input argument,
instead of the output of the rsq. This simplifies optimizations
to strip out the check based on value tracking. This approximately
results in equivalently good code. On targets with v_fmac_f64,
the result looks worse in the most trivial example due to a
bad decision to not rewrite to v_fma_f64 on the final
fma (llvm#171891). The result
is equivalently good in other final use contexts.

Author: arsenm

amd/device-libs/ocml/src/rsqrtD.cl

@@ -11,8 +11,7 @@ CONSTATTR double
 MATH_MANGLE(rsqrt)(double x)
 {
     double y0 = BUILTIN_AMDGPU_RSQRT_F64(x);
-    double e = MATH_MAD(-x*y0, y0, 1.0);
-    double y1 = MATH_MAD(y0*e, MATH_MAD(e, 0.375, 0.5), y0);
-    return BUILTIN_CLASS_F64(y0, CLASS_PSUB|CLASS_PNOR) ? y1 : y0;
+    double e = MATH_MAD(-y0 * (x == PINF_F64 || x == 0.0 ? y0 : x), y0, 1.0);
+    return MATH_MAD(y0*e, MATH_MAD(e, 0.375, 0.5), y0);
 }
 

amd/device-libs/test/compile/rsqrt.cl

@@ -26,13 +26,14 @@ float test_rsqrt_f32(float x) {
 
 // CHECK-LABEL: {{^}}test_rsqrt_f64:
 // CHECK: v_rsq_f64
+// CHECK: v_cmp_class_f64
+// CHECK: v_cndmask_b32
+// CHECK: v_cndmask_b32
 // CHECK: v_mul_f64
 // CHECK: v_fma_f64
 // CHECK: v_mul_f64
 // CHECK: v_fma_f64
 // CHECK: v_fma_f64
-// CHECK: v_cndmask_b32
-// CHECK: v_cndmask_b32
 double test_rsqrt_f64(double x) {
     return rsqrt(x);
 }