[mlir][math] Expand math.exp2 to use math.exp.

[bviyer]

Summary

Exp2 functions are pushed directly to libm. This is problematic for
situations where libm is not available. This patch will expand the exp2
function to use exp2 with the input multiplied by ln2 (natural log).

[rsuderman]

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

25

Rather than self defining we can just using stdlib's <numbers> library:

https://en.cppreference.com/w/cpp/header/numbers

mlir/test/Dialect/Math/expand-math.mlir

191

Missing newline at the end of the file.

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

700

Could we check values that fall outside the boundaries? We should verify that the edgecase behavior is approximately correct (e.g. very negative or very positive numbers should have a reasonable behavior).

Also we could we include nan, inf, and -inf cases.

[bviyer]

Added all the changes requested by rsuderman

[bviyer]

Moved the approximation tests to a new file.

mlir/include/mlir/Dialect/Math/Transforms/Passes.h

 class RewritePatternSet;
 
 void populateExpandCtlzPattern(RewritePatternSet &patterns);
 void populateExpandTanPattern(RewritePatternSet &patterns);
 void populateExpandTanhPattern(RewritePatternSet &patterns);
 void populateExpandFmaFPattern(RewritePatternSet &patterns);
 void populateExpandFloorFPattern(RewritePatternSet &patterns);
 void populateExpandCeilFPattern(RewritePatternSet &patterns);
+void populateExpandExp2FPattern(RewritePatternSet &patterns);
 void populateMathAlgebraicSimplificationPatterns(RewritePatternSet &patterns);
 
 struct MathPolynomialApproximationOptions {
   // Enables the use of AVX2 intrinsics in some of the approximations.
   bool enableAvx2 = false;
 };
 
 void populateMathPolynomialApproximationPatterns(
     RewritePatternSet &patterns,
     const MathPolynomialApproximationOptions &options = {});
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_

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

 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
+#define LN_2 (0.693147180559945309417232121458)
+
 /// Create a float constant.
 static Value createFloatConst(Location loc, Type type, double value,
                               OpBuilder &b) {
   auto attr = b.getFloatAttr(getElementTypeOrSelf(type), value);
   if (auto shapedTy = dyn_cast<ShapedType>(type)) {
     return b.create<arith::ConstantOp>(loc,
                                        DenseElementsAttr::get(shapedTy, attr));
   }
Value gtCheck = b.create<arith::CmpFOp>(arith::CmpFPredicate::OGT, operand,
                                           fpFixedConvert);
   Value incrValue = b.create<arith::SelectOp>(op->getLoc(), gtCheck, one, zero);
 
   Value ret = b.create<arith::AddFOp>(opType, fpFixedConvert, incrValue);
   rewriter.replaceOp(op, ret);
   return success();
 }
 
+// exp2f(float x) -> exp(x * ln(2))
+//   Proof: Let's say 2^x = y
+//   ln(2^x) = ln(y)
+//   x * ln(2) = ln(y) => e ^(x*ln(2)) = y
+static LogicalResult convertExp2fOp(math::Exp2Op op,
+                                    PatternRewriter &rewriter) {
+  ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+  Value operand = op.getOperand();
+  Type opType = operand.getType();
+  Value ln2 = createFloatConst(op->getLoc(), opType, LN_2, b);
+  Value mult = b.create<arith::MulFOp>(opType, operand, ln2);
+  Value exp = b.create<math::ExpOp>(op->getLoc(), mult);
+  rewriter.replaceOp(op, exp);
+  return success();
+}
+
 // Converts math.ctlz to scf and arith operations. This is done
 // by performing a binary search on the bits.
 static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
                                    PatternRewriter &rewriter) {
   auto operand = op.getOperand();
   auto operandTy = operand.getType();
   auto eTy = getElementTypeOrSelf(operandTy);
   Location loc = op.getLoc();

 void mlir::populateExpandFmaFPattern(RewritePatternSet &patterns) {
   patterns.add(convertFmaFOp);
 }
 
 void mlir::populateExpandCeilFPattern(RewritePatternSet &patterns) {
   patterns.add(convertCeilOp);
 }
 
+void mlir::populateExpandExp2FPattern(RewritePatternSet &patterns) {
+  patterns.add(convertExp2fOp);
+}
+
 void mlir::populateExpandFloorFPattern(RewritePatternSet &patterns) {
   patterns.add(convertFloorOp);
 }

mlir/test/Dialect/Math/expand-math.mlir

func.func @ceilf_func(%a: f64) -> f64 {
   // CHECK-NEXT:   [[CVTF:%.+]] = arith.sitofp [[CVTI]]
   // CHECK-NEXT:   [[COMP:%.+]] = arith.cmpf ogt, [[ARG0]], [[CVTF]]
   // CHECK-NEXT:   [[INCR:%.+]] = arith.select [[COMP]], [[CST_0]], [[CST]]
   // CHECK-NEXT:   [[ADDF:%.+]] = arith.addf [[CVTF]], [[INCR]]
   // CHECK-NEXT:   return [[ADDF]]
   %ret = math.ceil %a : f64
   return %ret : f64
 }
+
+// -----
+
+// CHECK-LABEL:     func @exp2f_func
+// CHECK-SAME:      ([[ARG0:%.+]]: f64) -> f64
+func.func @exp2f_func(%a: f64) -> f64 {
+  // CHECK-DAG:     [[CST:%.+]]  = arith.constant 0.69314718055994529
+  // CHECK-NEXT:    [[MULF:%.+]] = arith.mulf [[ARG0]], [[CST]]
+  // CHECK-NEXT:    [[EXP:%.+]]  = math.exp [[MULF]]
+  // CHECK-NEXT:    return [[EXP]]
+  %ret = math.exp2 %a : f64
+  return %ret : f64
+}
+
+// CHECK-LABEL:     func @exp2f_func_tensor
+// CHECK-SAME:      ([[ARG0:%.+]]: tensor<1xf32>) -> tensor<1xf32>
+func.func @exp2f_func_tensor(%a: tensor<1xf32>) -> tensor<1xf32> {
+  // CHECK-DAG:     [[CST:%.+]]  = arith.constant dense<0.693147182>
+  // CHECK-NEXT:    [[MULF:%.+]] = arith.mulf [[ARG0]], [[CST]]
+  // CHECK-NEXT:    [[EXP:%.+]]  = math.exp [[MULF]]
+  // CHECK-NEXT:    return [[EXP]]
+  %ret = math.exp2 %a : tensor<1xf32>
+  return %ret : tensor<1xf32>
+}
 No newline at end of file

mlir/test/lib/Dialect/Math/TestExpandMath.cpp

struct TestExpandMathPass
   }
   StringRef getDescription() const final { return "Test expanding math"; }
 };
 } // namespace
 
 void TestExpandMathPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   populateExpandCtlzPattern(patterns);
+  populateExpandExp2FPattern(patterns);
   populateExpandTanPattern(patterns);
   populateExpandTanhPattern(patterns);
   populateExpandFmaFPattern(patterns);
   populateExpandFloorFPattern(patterns);
   populateExpandCeilFPattern(patterns);
   (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }
 
 namespace mlir {
 namespace test {
 void registerTestExpandMathPass() { PassRegistration<TestExpandMathPass>(); }
 } // namespace test
 } // namespace mlir

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

func.func @ceilf() {
 
   // CHECK: 495
   %f = arith.constant 495.0 : f32
   call @func_ceilf32(%f) : (f32) -> ()
 
   return
 }
 
+// -------------------------------------------------------------------------- //
+// exp2f.
+// -------------------------------------------------------------------------- //
+func.func @func_exp2f(%a : f64) {
+  %r = math.exp2 %a : f64
+  vector.print %r : f64
+  return
+}
+
+func.func @exp2f() {
+  // CHECK: 2
+  %a = arith.constant 1.0 : f64
+  call @func_exp2f(%a) : (f64) -> ()  
+
+  // CHECK: 4
+  %b = arith.constant 2.0 : f64
+  call @func_exp2f(%b) : (f64) -> ()
+
+  // CHECK: 5.65685
+  %c = arith.constant 2.5 : f64
+  call @func_exp2f(%c) : (f64) -> ()
+
+  // CHECK: 0.29730
+  %d = arith.constant -1.75 : f64
+  call @func_exp2f(%d) : (f64) -> ()
+
+  // CHECK: 1.09581
+  %e = arith.constant 0.132 : f64
+  call @func_exp2f(%e) : (f64) -> ()
+  return
+}
+
 func.func @main() {
   call @tanh(): () -> ()
   call @log(): () -> ()
   call @log2(): () -> ()
   call @log1p(): () -> ()
   call @erf(): () -> ()
   call @exp(): () -> ()
   call @expm1(): () -> ()
   call @sin(): () -> ()
   call @cos(): () -> ()
   call @atan() : () -> ()
   call @atan2() : () -> ()
   call @cbrt() : () -> ()
   call @floorf() : () -> ()
   call @ceilf() : () -> ()
+  call @exp2f() : () -> ()
   return
 }