[mlir][math] Expand math.powf to exp, log and multiply

[bviyer]

Summary

Powf functions are pushed directly to libm. This is problematic for
situations where libm is not available. This patch will decompose the
powf function into log of exponent multiplied by log of base and raise
it to the exp.

[rsuderman]

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

167

Do we have a polynomial approximation for log? Curious whether it currently exists or not. This is not a blocking issue, just wondering if we should add it on our polynomial approximation list.

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

177

This should just be CHECK

[bviyer]

Replaced CHECK-NEXT with CHECK.

[bviyer]

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

167

yes we do. We have for Log base 2 and log base e.

[bviyer]

Moved the test to a different file.

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

 
 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 populateExpandPowFPattern(RewritePatternSet &patterns);
 void populateExpandRoundFPattern(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

static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) {
   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();
 }
+// Converts  Powf(float a, float b) (meaning a^b) to exp^(b * ln(a))
+static LogicalResult convertPowfOp(math::PowFOp op, PatternRewriter &rewriter) {
+  ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+  Value operandA = op.getOperand(0);
+  Value operandB = op.getOperand(1);
+  Type opType = operandA.getType();
+
+  Value logA = b.create<math::LogOp>(opType, operandA);
+  Value mult = b.create<arith::MulFOp>(opType, logA, operandB);
+  Value expResult = b.create<math::ExpOp>(opType, mult);
+  rewriter.replaceOp(op, expResult);
+  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);

 void mlir::populateExpandCeilFPattern(RewritePatternSet &patterns) {
   patterns.add(convertCeilOp);
 }
 
 void mlir::populateExpandExp2FPattern(RewritePatternSet &patterns) {
   patterns.add(convertExp2fOp);
 }
 
+void mlir::populateExpandPowFPattern(RewritePatternSet &patterns) {
+  patterns.add(convertPowfOp);
+}
+
 void mlir::populateExpandRoundFPattern(RewritePatternSet &patterns) {
   patterns.add(convertRoundOp);
 }
 
 void mlir::populateExpandFloorFPattern(RewritePatternSet &patterns) {
   patterns.add(convertFloorOp);
 }

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

 // -----
 
 // 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:         [[MULF:%.+]] = arith.mulf [[ARG0]], [[CST]]
   // CHECK:         [[EXP:%.+]]  = math.exp [[MULF]]
   // CHECK:         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>
func.func @roundf_func(%a: f64) -> f64 {
   // CHECK-DAG:  [[SEL:%.+]] = arith.select [[COMP]], [[CST_0]], [[CST_1]]
   // CHECK-DAG:  [[ADDF:%.+]] = arith.addf [[ARG0]], [[SEL]]
   // CHECK-DAG:   [[CVTI:%.+]] = arith.fptosi [[ADDF]]
   // CHECK-DAG:   [[CVTF:%.+]] = arith.sitofp [[CVTI]]
   // CHECK:   return [[CVTF]]
   %ret = math.round %a : f64
   return %ret : f64
 }
+
+// -----
+
+// CHECK-LABEL:   func @powf_func
+// CHECK-SAME:    ([[ARG0:%.+]]: f64, [[ARG1:%.+]]: f64)
+func.func @powf_func(%a: f64, %b: f64) ->f64 {
+  // CHECK-DAG: [[LOG:%.+]] = math.log [[ARG0]]
+  // CHECK-DAG: [[MULT:%.+]] = arith.mulf [[LOG]], [[ARG1]]
+  // CHECK-DAG: [[EXPR:%.+]] = math.exp [[MULT]]
+  // CHECK: return [[EXPR]]
+  %ret = math.powf %a, %b : f64
+  return %ret : f64
+}

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

void TestExpandMathPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   populateExpandCtlzPattern(patterns);
   populateExpandExp2FPattern(patterns);
   populateExpandTanPattern(patterns);
   populateExpandTanhPattern(patterns);
   populateExpandFmaFPattern(patterns);
   populateExpandFloorFPattern(patterns);
   populateExpandCeilFPattern(patterns);
+  populateExpandPowFPattern(patterns);
   populateExpandRoundFPattern(patterns);
   (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }
 
 namespace mlir {
 namespace test {
 void registerTestExpandMathPass() { PassRegistration<TestExpandMathPass>(); }
 } // namespace test
 } // namespace mlir

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

func.func @roundf() {
 
   // CHECK: -9
   %h = arith.constant -8.5 : f32
   call @func_roundf(%h) : (f32) -> ()
 
   return
 }
 
+// -------------------------------------------------------------------------- //
+// pow.
+// -------------------------------------------------------------------------- //
+func.func @func_powff64(%a : f64, %b : f64) {
+  %r = math.powf %a, %b : f64
+  vector.print %r : f64
+  return
+}
+
+func.func @powf() {
+  // CHECK: 16
+  %a   = arith.constant 4.0 : f64
+  %a_p = arith.constant 2.0 : f64
+  call @func_powff64(%a, %a_p) : (f64, f64) -> ()
+
+  // CHECK: -nan
+  %b   = arith.constant -3.0 : f64
+  %b_p = arith.constant 3.0 : f64
+  call @func_powff64(%b, %b_p) : (f64, f64) -> ()
+
+  // CHECK: 2.343
+  %c   = arith.constant 2.343 : f64
+  %c_p = arith.constant 1.000 : f64
+  call @func_powff64(%c, %c_p) : (f64, f64) -> ()
+
+  // CHECK: 0.176171
+  %d   = arith.constant 4.25 : f64
+  %d_p = arith.constant -1.2  : f64
+  call @func_powff64(%d, %d_p) : (f64, f64) -> ()
+
+  // CHECK: 1
+  %e   = arith.constant 4.385 : f64
+  %e_p = arith.constant 0.00 : f64
+  call @func_powff64(%e, %e_p) : (f64, f64) -> ()
+
+  // CHECK: 6.62637
+  %f    = arith.constant 4.835 : f64
+  %f_p  = arith.constant 1.2 : f64
+  call @func_powff64(%f, %f_p) : (f64, f64) -> ()
+
+  // CHECK: -nan
+  %g    = arith.constant 0xff80000000000000 : f64
+  call @func_powff64(%g, %g) : (f64, f64) -> ()
+
+  // CHECK: nan
+  %h = arith.constant 0x7fffffffffffffff : f64
+  call @func_powff64(%h, %h) : (f64, f64) -> ()
+
+  // CHECK: nan
+  %i = arith.constant 1.0 : f64
+  call @func_powff64(%i, %h) : (f64, f64) -> ()
+
+  // CHECK: inf
+  %j   = arith.constant 29385.0 : f64
+  %j_p = arith.constant 23598.0 : f64
+  call @func_powff64(%j, %j_p) : (f64, f64) -> () 
+  return
+}
 
 func.func @main() {
   call @exp2f() : () -> ()
   call @roundf() : () -> ()
+  call @powf() : () -> ()
   return
 }