[mlir][math] Expand math.floorf to truncate, compares and increments

[bviyer]

Summary

Floorf are pushed directly to libm. This is problematic for
situations where libm is not available. This patch will break down
a floorf function to truncate followed by an increment for negative
values, if necessary.

[rsuderman]

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

192

Missing a newline at the end of the file.

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

140

For CHECK statements we typically do the following

CHECK-DAG for any constant values (e.g. arith.constant)
CHECK: for any computation.

In theory you can also do CHECK-DAG for everything but the return checker.

150

Missing a new line at the end of the file.

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

626

Could you include boundary checks for both positive and negative values? E.g. -2.0 and 2,0.

[bviyer]

Fixed all the changes requested by rsuderman.

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

 namespace mlir {
 
 class RewritePatternSet;
 
 void populateExpandCtlzPattern(RewritePatternSet &patterns);
 void populateExpandTanPattern(RewritePatternSet &patterns);
 void populateExpandTanhPattern(RewritePatternSet &patterns);
 void populateExpandFmaFPattern(RewritePatternSet &patterns);
-
+void populateExpandFloorFPattern(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 convertFmaFOp(math::FmaOp op, PatternRewriter &rewriter) {
   Value operandC = op.getOperand(2);
   Type type = op.getType();
   Value mult = b.create<arith::MulFOp>(type, operandA, operandB);
   Value add = b.create<arith::AddFOp>(type, mult, operandC);
   rewriter.replaceOp(op, add);
   return success();
 }
 
+// Converts a floorf() function to the following:
+// floorf(float x) ->
+//     y = (float)(int) x
+//     if (x < 0) then incr = -1 else incr = 0
+//     y = y + incr    <= replace this op with the floorf op.
+static LogicalResult convertFloorOp(math::FloorOp op,
+                                    PatternRewriter &rewriter) {
+  ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+  Value operand = op.getOperand();
+  Type opType = operand.getType();
+  Value fixedConvert = b.create<arith::FPToSIOp>(b.getI64Type(), operand);
+  Value fpFixedConvert = b.create<arith::SIToFPOp>(opType, fixedConvert);
+
+  // Creating constants for later use.
+  Value zero = createFloatConst(op->getLoc(), opType, 0.00, rewriter);
+  Value negOne = createFloatConst(op->getLoc(), opType, -1.00, rewriter);
+
+  Value negCheck =
+      b.create<arith::CmpFOp>(arith::CmpFPredicate::OLT, operand, zero);
+  Value incrValue =
+      b.create<arith::SelectOp>(op->getLoc(), negCheck, negOne, zero);
+  Value ret = b.create<arith::AddFOp>(opType, fpFixedConvert, incrValue);
+  rewriter.replaceOp(op, ret);
+  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::populateExpandTanhPattern(RewritePatternSet &patterns) {
   patterns.add(convertTanhOp);
 }
 
 void mlir::populateExpandFmaFPattern(RewritePatternSet &patterns) {
   patterns.add(convertFmaFOp);
 }
+void mlir::populateExpandFloorFPattern(RewritePatternSet &patterns) {
+  patterns.add(convertFloorOp);
+}

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

 // CHECK-SAME:     ([[ARG0:%.+]]: f64, [[ARG1:%.+]]: f64, [[ARG2:%.+]]: f64) -> f64
 func.func @fmaf_func(%a: f64, %b: f64, %c: f64) -> f64 {
   // CHECK-NEXT:     [[MULF:%.+]] = arith.mulf [[ARG0]], [[ARG1]]
   // CHECK-NEXT:     [[ADDF:%.+]] = arith.addf [[MULF]], [[ARG2]]
   // CHECK-NEXT:     return [[ADDF]]
   %ret = math.fma %a, %b, %c : f64
   return %ret : f64
 }
+
+// -----
+
+// CHECK-LABEL:     func @floorf_func
+// CHECK-SAME:      ([[ARG0:%.+]]: f64) -> f64
+func.func @floorf_func(%a: f64) -> f64 {
+  // CHECK-DAG:   [[CST:%.+]] = arith.constant 0.000
+  // CHECK-DAG:   [[CST_0:%.+]] = arith.constant -1.000
+  // CHECK-NEXT:   [[CVTI:%.+]] = arith.fptosi [[ARG0]]
+  // CHECK-NEXT:   [[CVTF:%.+]] = arith.sitofp [[CVTI]]
+  // CHECK-NEXT:   [[COMP:%.+]] = arith.cmpf olt, [[ARG0]], [[CST]]
+  // CHECK-NEXT:   [[INCR:%.+]] = arith.select [[COMP]], [[CST_0]], [[CST]]
+  // CHECK-NEXT:   [[ADDF:%.+]] = arith.addf [[CVTF]], [[INCR]]
+  // CHECK-NEXT:   return [[ADDF]]
+  %ret = math.floor %a : f64
+  return %ret : f64
+}

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

 } // namespace
 
 void TestExpandMathPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   populateExpandCtlzPattern(patterns);
   populateExpandTanPattern(patterns);
   populateExpandTanhPattern(patterns);
   populateExpandFmaFPattern(patterns);
+  populateExpandFloorFPattern(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 @cbrt() {
 
   // CHECK: 2.57128
   %g = arith.constant 17.0 : f32
   call @cbrt_f32(%g) : (f32) -> ()
 
   return
 }
 
+// -------------------------------------------------------------------------- //
+// floor.
+// -------------------------------------------------------------------------- //
+func.func @func_floorf32(%a : f32) {
+  %r = math.floor %a : f32
+  vector.print %r : f32
+  return
+}
+
+func.func @floorf() {
+  // CHECK: 3
+  %a = arith.constant 3.8 : f32
+  call @func_floorf32(%a) : (f32) -> ()
+
+  // CHECK: -4
+  %b = arith.constant -3.8 : f32
+  call @func_floorf32(%b) : (f32) -> ()
+
+  // CHECK: 0
+  %c = arith.constant 0.0 : f32
+  call @func_floorf32(%c) : (f32) -> ()
+
+  // CHECK: -5
+  %d = arith.constant -4.2 : f32
+  call @func_floorf32(%d) : (f32) -> ()
+
+  // CHECK: -2
+  %e = arith.constant -2.0 : f32
+  call @func_floorf32(%e) : (f32) -> ()
+
+  // CHECK: 2
+  %f = arith.constant 2.0 : f32
+  call @func_floorf32(%f) : (f32) -> ()
+
+  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() : () -> ()
   return
 }