Avoid infinity arithmetics when computing exp approximations

[asaadaldien]

Summary

Otherwise this can result a poison value on some platforms see https://bugs.llvm.org/show_bug.cgi?id=51204

[jonpa]

Thanks for fixing this!

I can confirm that the test now passes on SystemZ with this patch. So please also remove the XFAIL:

mlir/test/mlir-cpu-runner/math_polynomial_approx.mlir

-// XFAIL: s390x
-// (see https://bugs.llvm.org/show_bug.cgi?id=51204)
-

I cannot approve of the patch since I am not that familiar with mlir, I hope someone else can review it fully...

[asaadaldien]

remove XFAIL

[asaadaldien]

check vecor type

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

ExpApproximation::matchAndRewrite(math::ExpOp op,
   Value kMaxNegConst = bcast(i32Cst(builder, -127));
   Value rightBound =
       builder.create<arith::CmpIOp>(arith::CmpIPredicate::sle, k, kMaxConst);
   Value leftBound =
       builder.create<arith::CmpIOp>(arith::CmpIPredicate::sge, k, kMaxNegConst);
 
   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<SelectOp>(
-      isComputable, expY,
+      isNegInfinityX, zerof32Const,
       builder.create<SelectOp>(
-          isPostiveX, constPosInfinity,
-          builder.create<SelectOp>(isNegInfinityX, zerof32Const, underflow)));
+          isPosInfinityX, constPosInfinity,
+          builder.create<SelectOp>(isComputable, expY,
+                                   builder.create<SelectOp>(isPostiveX,
+                                                            constPosInfinity,
+                                                            underflow))));
 
   rewriter.replaceOp(op, expY);
 
   return success();
 }
 
 //----------------------------------------------------------------------------//
 // ExpM1 approximation.

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

 // CHECK:           %[[VAL_26:.*]] = arith.fptosi %[[VAL_16]] : f32 to i32
 // CHECK:           %[[VAL_27:.*]] = arith.addi %[[VAL_26]], %[[VAL_13]] : i32
 // CHECK:           %[[VAL_28:.*]] = arith.shli %[[VAL_27]], %[[VAL_8]] : i32
 // CHECK:           %[[VAL_29:.*]] = arith.bitcast %[[VAL_28]] : i32 to f32
 // CHECK:           %[[VAL_30:.*]] = arith.mulf %[[VAL_25]], %[[VAL_29]] : f32
 // CHECK:           %[[VAL_31:.*]] = arith.cmpi sle, %[[VAL_26]], %[[VAL_13]] : i32
 // CHECK:           %[[VAL_32:.*]] = arith.cmpi sge, %[[VAL_26]], %[[VAL_14]] : i32
 // CHECK:           %[[VAL_33:.*]] = arith.cmpf oeq, %[[VAL_0]], %[[VAL_11]] : f32
-// CHECK:           %[[VAL_34:.*]] = arith.cmpf ogt, %[[VAL_0]], %[[VAL_9]] : f32
-// CHECK:           %[[VAL_35:.*]] = arith.andi %[[VAL_31]], %[[VAL_32]] : i1
-// CHECK:           %[[VAL_36:.*]] = select %[[VAL_33]], %[[VAL_9]], %[[VAL_12]] : f32
-// CHECK:           %[[VAL_37:.*]] = select %[[VAL_34]], %[[VAL_10]], %[[VAL_36]] : f32
-// CHECK:           %[[VAL_38:.*]] = select %[[VAL_35]], %[[VAL_30]], %[[VAL_37]] : f32
-// CHECK:           return %[[VAL_38]] : f32
-// CHECK:         }
+// 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:.*]] = select %[[VAL_35]], %[[VAL_10]], %[[VAL_12]] : f32
+// CHECK:           %[[VAL_38:.*]] = select %[[VAL_36]], %[[VAL_30]], %[[VAL_37]] : f32
+// CHECK:           %[[VAL_39:.*]] = select %[[VAL_34]], %[[VAL_10]], %[[VAL_38]] : f32
+// CHECK:           %[[VAL_40:.*]] = select %[[VAL_33]], %[[VAL_9]], %[[VAL_39]] : f32
+// CHECK:           return %[[VAL_40]] : 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-2:   select
-// CHECK:           %[[VAL_38:.*]] = select
-// CHECK:           return %[[VAL_38]] : vector<8xf32>
-// CHECK:         }
+// CHECK-COUNT-3:   select
+// CHECK:           %[[VAL_40:.*]] = select
+// CHECK:           return %[[VAL_40]]
 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-2:   select
+// CHECK-COUNT-3:   select
 // CHECK:           %[[EXP_X:.*]] = select
 // CHECK:           %[[VAL_58:.*]] = 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:.*]] = 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:.*]] = select %[[VAL_104]], %[[EXP_X]], %[[VAL_106]] : f32
 // CHECK:           %[[VAL_108:.*]] = select %[[VAL_60]], %[[CST_MINUSONE]], %[[VAL_107]] : f32
 // CHECK:           %[[VAL_109:.*]] = select %[[VAL_58]], %[[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<8xf32>) -> vector<8xf32> {
 // CHECK:           %[[VAL_1:.*]] = arith.constant dense<-1.000000e+00> : vector<8xf32>
 // CHECK-NOT:       exp
-// CHECK-COUNT-3:   select
+// CHECK-COUNT-4:   select
 // CHECK-NOT:       log
 // CHECK-COUNT-5:   select
 // CHECK-NOT:       expm1
 // CHECK-COUNT-3:   select
 // CHECK:           %[[VAL_115:.*]] = select
 // CHECK:           return %[[VAL_115]] : vector<8xf32>
 // CHECK:         }
 func @expm1_vector(%arg0: vector<8xf32>) -> vector<8xf32> {

mlir/test/mlir-cpu-runner/math_polynomial_approx.mlir

 // RUN:   mlir-opt %s -test-math-polynomial-approximation                      \
 // RUN:               -convert-arith-to-llvm                                   \
 // RUN:               -convert-vector-to-llvm                                  \
 // RUN:               -convert-math-to-llvm                                    \
 // RUN:               -convert-std-to-llvm                                     \
 // RUN:               -reconcile-unrealized-casts                              \
 // RUN: | mlir-cpu-runner                                                      \
 // RUN:     -e main -entry-point-result=void -O0                               \
 // RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_c_runner_utils%shlibext  \
 // RUN:     -shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext    \
 // RUN: | FileCheck %s
 
-// XFAIL: s390x
-// (see https://bugs.llvm.org/show_bug.cgi?id=51204)
-
 // -------------------------------------------------------------------------- //
 // Tanh.
 // -------------------------------------------------------------------------- //
 func @tanh() {
   // CHECK: 0.848284
   %0 = arith.constant 1.25 : f32
   %1 = math.tanh %0 : f32
   vector.print %1 : f32