[mlir][Math] Fix NaN handling in ExpM1 approximation.

[akuegel]

Summary

For NaN input, output should be NaN as well.

[bkramer]

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

1038

Would changing this to UEQ do the same thing?

[akuegel]

Replace UNO with UEQ.

[akuegel]

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

1038

Yes, that should work as well to detect NaN. Changed.

[bkramer]

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

1038

Sorry for being unclear. My idea was to do the one check and the NaN check with a single UEQ instead of two fcmps.

[akuegel]

Actually do the suggested change.

[bkramer]

thanks

[akuegel]

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

1038

Ah, thanks, that is a nice observation :)
And my previous change was wrong, anyway, I can use UNE or UNO when comparing x to itself to detect NaN, but not UEQ.

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

ExpM1Approximation::matchAndRewrite(math::ExpM1Op op,
   };
 
   // expm1(x) = exp(x) - 1 = u - 1.
   // We have to handle it carefully when x is near 0, i.e. u ~= 1,
   // and when the input is ~= -inf, i.e. u - 1 ~= -1.
   Value cstOne = bcast(f32Cst(builder, 1.0f));
   Value cstNegOne = bcast(f32Cst(builder, -1.0f));
   Value x = op.getOperand();
+  Value xIsNaN = builder.create<arith::CmpFOp>(arith::CmpFPredicate::UNO, x, x);
   Value u = builder.create<math::ExpOp>(x);
   Value uEqOne =
       builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, u, cstOne);
   Value uMinusOne = builder.create<arith::SubFOp>(u, cstOne);
   Value uMinusOneEqNegOne = builder.create<arith::CmpFOp>(
       arith::CmpFPredicate::OEQ, uMinusOne, cstNegOne);
   // logU = log(u) ~= x
   Value logU = builder.create<math::LogOp>(u);
 
   // Detect exp(x) = +inf; written this way to avoid having to form +inf.
   Value isInf =
       builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, logU, u);
 
   // (u - 1) * (x / ~x)
   Value expm1 = builder.create<arith::MulFOp>(
       uMinusOne, builder.create<arith::DivFOp>(x, logU));
   expm1 = builder.create<arith::SelectOp>(isInf, u, expm1);
+  Value uEqOneOrNaN = builder.create<arith::OrIOp>(uEqOne, xIsNaN);
   Value approximation = builder.create<arith::SelectOp>(
-      uEqOne, x,
+      uEqOneOrNaN, x,
       builder.create<arith::SelectOp>(uMinusOneEqNegOne, cstNegOne, expm1));
   rewriter.replaceOp(op, approximation);
   return success();
 }
 
 //----------------------------------------------------------------------------//
 // Sin and Cos approximation.
 //----------------------------------------------------------------------------//

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

func @exp_vector(%arg0: vector<8xf32>) -> 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-DAG:           %[[IS_NAN:.*]] = arith.cmpf uno, %[[X]], %[[X]] : f32
 // CHECK:           %[[BEGIN_EXP_X:.*]] = arith.mulf %[[X]], %[[CST_LOG2E]] : f32
 // CHECK-NOT:       exp
 // CHECK-COUNT-3:   select
 // CHECK:           %[[EXP_X:.*]] = arith.select
-// CHECK:           %[[VAL_58:.*]] = arith.cmpf oeq, %[[EXP_X]], %[[CST_ONE]] : f32
+// CHECK:           %[[IS_ONE:.*]] = arith.cmpf oeq, %[[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:           %[[IS_ONE_OR_NAN:.*]] = arith.ori %[[IS_ONE]], %[[IS_NAN]] : i1
 // CHECK:           %[[VAL_108:.*]] = arith.select %[[VAL_60]], %[[CST_MINUSONE]], %[[VAL_107]] : f32
-// CHECK:           %[[VAL_109:.*]] = arith.select %[[VAL_58]], %[[X]], %[[VAL_108]] : 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(