[mlir][Math] Fix NaN handling in Exp approximation

[akuegel]

Also add integration test.

[akuegel]

Rebase.

[bkramer]

looks good

mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp

ExpApproximation::matchAndRewrite(math::ExpOp op,
   Value cstCephesExpP1 = bcast(f32Cst(builder, 1.0));
   Value cstCephesExpP2 = bcast(f32Cst(builder, 0.49970514590562437052f));
   Value cstCephesExpP3 = bcast(f32Cst(builder, 0.16873890085469545053f));
   Value cstCephesExpP4 = bcast(f32Cst(builder, 0.03668965196652099192f));
   Value cstCephesExpP5 = bcast(f32Cst(builder, 0.01314350012789660196f));
 
   Value x = op.getOperand();
 
+  Value isNan = builder.create<arith::CmpFOp>(arith::CmpFPredicate::UNO, x, x);
+
   // Reduced y = x - floor(x / ln(2)) * ln(2) = x - k * ln(2)
   Value xL2Inv = mul(x, cstLog2E);
   Value kF32 = floor(xL2Inv);
   Value kLn2 = mul(kF32, cstLn2);
   Value y = sub(x, kLn2);
 
   // Use Estrin's evaluation scheme with 3 independent parts:
   // P(y)^y : (c0 + c1 y) + (c2 + c3 y) y^2 + (c4 + c5 y) y^4
Value isNegInfinityX = builder.create<arith::CmpFOp>(
       arith::CmpFPredicate::OEQ, x, constNegIfinity);
   Value isPosInfinityX = builder.create<arith::CmpFOp>(
       arith::CmpFPredicate::OEQ, x, constPosInfinity);
   Value isPostiveX =
       builder.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, x, zerof32Const);
   Value isComputable = builder.create<arith::AndIOp>(rightBound, leftBound);
 
   expY = builder.create<arith::SelectOp>(
+      isNan, x,
+      builder.create<arith::SelectOp>(
           isNegInfinityX, zerof32Const,
           builder.create<arith::SelectOp>(
               isPosInfinityX, constPosInfinity,
               builder.create<arith::SelectOp>(
                   isComputable, expY,
                   builder.create<arith::SelectOp>(isPostiveX, constPosInfinity,
-                                              underflow))));
+                                                  underflow)))));
 
   rewriter.replaceOp(op, expY);
 
   return success();
 }
 
 //----------------------------------------------------------------------------//
 // ExpM1 approximation.

mlir/test/Dialect/Math/polynomial-approximation.mlir

 // CHECK-DAG:           %[[VAL_7:.*]] = arith.constant 1.314350e-02 : f32
 // CHECK-DAG:           %[[VAL_8:.*]] = arith.constant 23 : i32
 // CHECK-DAG:           %[[VAL_9:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK-DAG:           %[[VAL_10:.*]] = arith.constant 0x7F800000 : f32
 // CHECK-DAG:           %[[VAL_11:.*]] = arith.constant 0xFF800000 : f32
 // CHECK-DAG:           %[[VAL_12:.*]] = arith.constant 1.17549435E-38 : f32
 // CHECK-DAG:           %[[VAL_13:.*]] = arith.constant 127 : i32
 // CHECK-DAG:           %[[VAL_14:.*]] = arith.constant -127 : i32
+// CHECK:           %[[IS_NAN:.*]] = arith.cmpf uno, %[[VAL_0]], %[[VAL_0]] : f32
 // CHECK:           %[[VAL_15:.*]] = arith.mulf %[[VAL_0]], %[[VAL_2]] : f32
 // CHECK:           %[[VAL_16:.*]] = math.floor %[[VAL_15]] : f32
 // CHECK:           %[[VAL_17:.*]] = arith.mulf %[[VAL_16]], %[[VAL_1]] : f32
 // CHECK:           %[[VAL_18:.*]] = arith.subf %[[VAL_0]], %[[VAL_17]] : f32
 // CHECK:           %[[VAL_19:.*]] = arith.mulf %[[VAL_18]], %[[VAL_18]] : f32
 // CHECK:           %[[VAL_20:.*]] = arith.mulf %[[VAL_19]], %[[VAL_19]] : f32
 // CHECK:           %[[VAL_21:.*]] = math.fma %[[VAL_3]], %[[VAL_18]], %[[VAL_3]] : f32
 // CHECK:           %[[VAL_22:.*]] = math.fma %[[VAL_5]], %[[VAL_18]], %[[VAL_4]] : f32

 // CHECK:           %[[VAL_33:.*]] = arith.cmpf oeq, %[[VAL_0]], %[[VAL_11]] : f32
 // CHECK:           %[[VAL_34:.*]] = arith.cmpf oeq, %[[VAL_0]], %[[VAL_10]] : f32
 // CHECK:           %[[VAL_35:.*]] = arith.cmpf ogt, %[[VAL_0]], %[[VAL_9]] : f32
 // CHECK:           %[[VAL_36:.*]] = arith.andi %[[VAL_31]], %[[VAL_32]] : i1
 // CHECK:           %[[VAL_37:.*]] = arith.select %[[VAL_35]], %[[VAL_10]], %[[VAL_12]] : f32
 // CHECK:           %[[VAL_38:.*]] = arith.select %[[VAL_36]], %[[VAL_30]], %[[VAL_37]] : f32
 // CHECK:           %[[VAL_39:.*]] = arith.select %[[VAL_34]], %[[VAL_10]], %[[VAL_38]] : f32
 // CHECK:           %[[VAL_40:.*]] = arith.select %[[VAL_33]], %[[VAL_9]], %[[VAL_39]] : f32
-// CHECK:           return %[[VAL_40]] : f32
+// CHECK:           %[[VAL_41:.*]] = arith.select %[[IS_NAN]], %[[VAL_0]], %[[VAL_40]] : f32
+// CHECK:           return %[[VAL_41]] : f32
 func @exp_scalar(%arg0: f32) -> f32 {
   %0 = math.exp %arg0 : f32
   return %0 : f32
 }
 
 // CHECK-LABEL:   func @exp_vector(
 // CHECK-SAME:                     %[[VAL_0:.*]]: vector<8xf32>) -> vector<8xf32> {
 // CHECK:           %[[VAL_1:.*]] = arith.constant dense<0.693147182> : vector<8xf32>
 // CHECK-NOT:       exp
-// CHECK-COUNT-3:   select
+// CHECK-COUNT-4:   select
 // CHECK:           %[[VAL_40:.*]] = arith.select
 // CHECK:           return %[[VAL_40]] : vector<8xf32>
 func @exp_vector(%arg0: vector<8xf32>) -> vector<8xf32> {
   %0 = math.exp %arg0 : vector<8xf32>
   return %0 : vector<8xf32>
 }
 
 // CHECK-LABEL:   func @expm1_scalar(
 // CHECK-SAME:                       %[[X:.*]]: f32) -> f32 {
 // CHECK-DAG:           %[[CST_MINUSONE:.*]] = arith.constant -1.000000e+00 : f32
 // CHECK-DAG:           %[[CST_LOG2E:.*]] = arith.constant 1.44269502 : f32
 // CHECK-DAG:           %[[CST_ONE:.*]] = arith.constant 1.000000e+00 : f32
 // CHECK:           %[[BEGIN_EXP_X:.*]] = arith.mulf %[[X]], %[[CST_LOG2E]] : f32
 // CHECK-NOT:       exp
-// CHECK-COUNT-3:   select
+// CHECK-COUNT-4:   select
 // CHECK:           %[[EXP_X:.*]] = arith.select
 // CHECK:           %[[IS_ONE_OR_NAN:.*]] = arith.cmpf ueq, %[[EXP_X]], %[[CST_ONE]] : f32
 // CHECK:           %[[VAL_59:.*]] = arith.subf %[[EXP_X]], %[[CST_ONE]] : f32
 // CHECK:           %[[VAL_60:.*]] = arith.cmpf oeq, %[[VAL_59]], %[[CST_MINUSONE]] : f32
 // CHECK-NOT:       log
 // CHECK-COUNT-5:   select
 // CHECK:           %[[LOG_U:.*]] = arith.select
 // CHECK:           %[[VAL_104:.*]] = arith.cmpf oeq, %[[LOG_U]], %[[EXP_X]] : f32
 // CHECK:           %[[VAL_105:.*]] = arith.divf %[[X]], %[[LOG_U]] : f32
 // CHECK:           %[[VAL_106:.*]] = arith.mulf %[[VAL_59]], %[[VAL_105]] : f32
 // CHECK:           %[[VAL_107:.*]] = arith.select %[[VAL_104]], %[[EXP_X]], %[[VAL_106]] : f32
 // CHECK:           %[[VAL_108:.*]] = arith.select %[[VAL_60]], %[[CST_MINUSONE]], %[[VAL_107]] : f32
 // CHECK:           %[[VAL_109:.*]] = arith.select %[[IS_ONE_OR_NAN]], %[[X]], %[[VAL_108]] : f32
 // CHECK:           return %[[VAL_109]] : f32
 // CHECK:         }
 func @expm1_scalar(%arg0: f32) -> f32 {
   %0 = math.expm1 %arg0 : f32
   return %0 : f32
 }
 
 // CHECK-LABEL:   func @expm1_vector(
 // CHECK-SAME:                       %[[VAL_0:.*]]: vector<8x8xf32>) -> vector<8x8xf32> {
 // CHECK:           %[[VAL_1:.*]] = arith.constant dense<-1.000000e+00> : vector<8x8xf32>
 // CHECK-NOT:       exp
-// CHECK-COUNT-4:   select
+// CHECK-COUNT-5:   select
 // CHECK-NOT:       log
 // CHECK-COUNT-5:   select
 // CHECK-NOT:       expm1
 // CHECK-COUNT-3:   select
 // CHECK:           %[[VAL_115:.*]] = arith.select
 // CHECK:           return %[[VAL_115]] : vector<8x8xf32>
 // CHECK:         }
 func @expm1_vector(%arg0: vector<8x8xf32>) -> vector<8x8xf32> {

mlir/test/mlir-cpu-runner/math-polynomial-approx.mlir

func @exp() {
   %exp_inf = math.exp %inf : f32
   vector.print %exp_inf : f32
 
   // CHECK: 0
   %negative_inf = arith.constant 0xff800000 : f32
   %exp_negative_inf = math.exp %negative_inf : f32
   vector.print %exp_negative_inf : f32
 
+  // CHECK: nan
+  %nan = arith.constant 0x7fc00000 : f32
+  %exp_nan = math.exp %nan : f32
+  vector.print %exp_nan : f32
+
   return
 }
 
 func @expm1() {
   // CHECK: 1e-10
   %0 = arith.constant 1.0e-10 : f32
   %1 = math.expm1 %0 : f32
   vector.print %1 : f32
func @expm1() {
   %expm1_inf = math.expm1 %inf : f32
   vector.print %expm1_inf : f32
 
   // CHECK: -1, inf, 1e-10
   %special_vec = arith.constant dense<[0xff800000, 0x7f800000, 1.0e-10]> : vector<3xf32>
   %log_special_vec = math.expm1 %special_vec : vector<3xf32>
   vector.print %log_special_vec : vector<3xf32>
 
+  // CHECK: nan
+  %nan = arith.constant 0x7fc00000 : f32
+  %exp_nan = math.expm1 %nan : f32
+  vector.print %exp_nan : f32
+
   return
 }
 // -------------------------------------------------------------------------- //
 // Sin.
 // -------------------------------------------------------------------------- //
 func @sin() {
   // CHECK: 0
   %0 = arith.constant 0.0 : f32