[mlir][math] Expand math.round to truncate, compare and increment.

[bviyer]

Summary

Round functions are pushed directly to libm. This is problematic for
situations where libm is not available. This patch will decompose the
roundf function by adding 0.5 to positive number to input
(subtracting for negative) following by a truncate.

[rsuderman]

Just the one fix and we should be good.

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

177

Rather than CHECK-NEXT just use CHECK-DAG on everything except the return. That way we don't enforce strict ordering on operations when not necessary.

[bviyer]

Replaced all-but-last CHECK-NEXT with CHECK-DAG as requested by rsuderman

[rsuderman]

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

182

This should be CHECK

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

690 ↗

(On Diff #512940)

So I realized a problem. These tests are in polynomial approximation but should be in a separate CPU runner test. Specifically one that validates expand-math separately works. Before landing these two CLs, could have split out the arith lowering tests to a separate file and substitute test-math-polynomial-approximation for test-expand-math and add the equivalent arith expansion.

[rsuderman]

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

690 ↗

(On Diff #512940)

So just to note, we will want both one for arith's expand patterns and one for math's expand patterns.

[bviyer]

Added the test to the new file and also replaced CHECK-NEXT with CHECK.

[rsuderman]

Just one small change.

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

687–731 ↗

(On Diff #513304)

We should remove this as it is in the expand math tests now.

[bviyer]

Removed func_roundf and roundf

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 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 convertExp2fOp(math::Exp2Op op,
   Type opType = operand.getType();
   Value ln2 = createFloatConst(op->getLoc(), opType, llvm::numbers::ln2, 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();
 }
 
+static LogicalResult convertRoundOp(math::RoundOp op,
+                                    PatternRewriter &rewriter) {
+  ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+  Value operand = op.getOperand();
+  Type opType = operand.getType();
+
+  // Creating constants for later use.
+  Value zero = createFloatConst(op->getLoc(), opType, 0.00, rewriter);
+  Value half = createFloatConst(op->getLoc(), opType, 0.5, rewriter);
+  Value negHalf = createFloatConst(op->getLoc(), opType, -0.5, rewriter);
+
+  Value posCheck =
+      b.create<arith::CmpFOp>(arith::CmpFPredicate::OGE, operand, zero);
+  Value incrValue =
+      b.create<arith::SelectOp>(op->getLoc(), posCheck, half, negHalf);
+  Value add = b.create<arith::AddFOp>(opType, operand, incrValue);
+
+  Value fpFixedConvert = createTruncatedFPValue(add, b);
+  rewriter.replaceOp(op, fpFixedConvert);
+  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::populateExpandCeilFPattern(RewritePatternSet &patterns) {
   patterns.add(convertCeilOp);
 }
 
 void mlir::populateExpandExp2FPattern(RewritePatternSet &patterns) {
   patterns.add(convertExp2fOp);
 }
 
+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>
   // CHECK:         [[MULF:%.+]] = arith.mulf [[ARG0]], [[CST]]
   // CHECK:         [[EXP:%.+]]  = math.exp [[MULF]]
   // CHECK:         return [[EXP]]
   %ret = math.exp2 %a : tensor<1xf32>
   return %ret : tensor<1xf32>
 }
+
+// -----
+
+// CHECK-LABEL:      func @roundf_func
+// CHECK-SAME:      ([[ARG0:%.+]]: f64) -> f64
+func.func @roundf_func(%a: f64) -> f64 {
+  // CHECK-DAG:   [[CST:%.+]] = arith.constant 0.000
+  // CHECK-DAG:   [[CST_0:%.+]] = arith.constant 5.000000e-01
+  // CHECK-DAG:   [[CST_1:%.+]] = arith.constant -5.000000e-01
+  // CHECK-DAG:  [[COMP:%.+]] = arith.cmpf oge, [[ARG0]], [[CST]]
+  // 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
+}

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);
+  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 @exp2f() {
   call @func_exp2f(%neg_inf) : (f64) -> ()
 
   // CHECK: inf
   %i = arith.constant 0x7fc0000000000000 : f64
   call @func_exp2f(%i) : (f64) -> ()
   return
 }
 
+// -------------------------------------------------------------------------- //
+// round.
+// -------------------------------------------------------------------------- //
+func.func @func_roundf(%a : f32) {
+  %r = math.round %a : f32
+  vector.print %r : f32
+  return
+}
+
+func.func @roundf() {
+  // CHECK: 4
+  %a = arith.constant 3.8 : f32
+  call @func_roundf(%a) : (f32) -> ()
+
+  // CHECK: -4
+  %b = arith.constant -3.8 : f32
+  call @func_roundf(%b) : (f32) -> ()
+
+  // CHECK: 0
+  %c = arith.constant 0.0 : f32
+  call @func_roundf(%c) : (f32) -> ()
+
+  // CHECK: -4
+  %d = arith.constant -4.2 : f32
+  call @func_roundf(%d) : (f32) -> ()
+
+  // CHECK: -495
+  %e = arith.constant -495.0 : f32
+  call @func_roundf(%e) : (f32) -> ()
+
+  // CHECK: 495
+  %f = arith.constant 495.0 : f32
+  call @func_roundf(%f) : (f32) -> ()
+
+  // CHECK: 9
+  %g = arith.constant 8.5 : f32
+  call @func_roundf(%g) : (f32) -> ()
+
+  // CHECK: -9
+  %h = arith.constant -8.5 : f32
+  call @func_roundf(%h) : (f32) -> ()
+
+  return
+}
+
+
 func.func @main() {
   call @exp2f() : () -> ()
+  call @roundf() : () -> ()
   return
 }