[mlir][math] Lower cos,sin to libm

[Lewuathe]

Summary

Lower math.cos and math.sin to libm

[ftynse]

mlir/test/Conversion/MathToLibm/convert-to-libm.mlir

160

Please add a new line

mlir/lib/Conversion/MathToLibm/MathToLibm.cpp

rewriter.replaceOpWithNewOp<func::CallOp>(op, name, op.getType(),
                                             op->getOperands());
 
   return success();
 }
 
 void mlir::populateMathToLibmConversionPatterns(RewritePatternSet &patterns,
                                                 PatternBenefit benefit) {
   patterns.add<VecOpToScalarOp<math::Atan2Op>, VecOpToScalarOp<math::ExpM1Op>,
-               VecOpToScalarOp<math::TanhOp>>(patterns.getContext(), benefit);
+               VecOpToScalarOp<math::TanhOp>, VecOpToScalarOp<math::CosOp>,
+               VecOpToScalarOp<math::SinOp>>(patterns.getContext(), benefit);
   patterns.add<PromoteOpToF32<math::Atan2Op>, PromoteOpToF32<math::ExpM1Op>,
-               PromoteOpToF32<math::TanhOp>>(patterns.getContext(), benefit);
+               PromoteOpToF32<math::TanhOp>, PromoteOpToF32<math::CosOp>,
+               PromoteOpToF32<math::SinOp>>(patterns.getContext(), benefit);
   patterns.add<ScalarOpToLibmCall<math::Atan2Op>>(patterns.getContext(),
                                                   "atan2f", "atan2", benefit);
   patterns.add<ScalarOpToLibmCall<math::ErfOp>>(patterns.getContext(), "erff",
                                                 "erf", benefit);
   patterns.add<ScalarOpToLibmCall<math::ExpM1Op>>(patterns.getContext(),
                                                   "expm1f", "expm1", benefit);
   patterns.add<ScalarOpToLibmCall<math::TanhOp>>(patterns.getContext(), "tanhf",
                                                  "tanh", benefit);
   patterns.add<ScalarOpToLibmCall<math::RoundOp>>(patterns.getContext(),
                                                   "roundf", "round", benefit);
+  patterns.add<ScalarOpToLibmCall<math::CosOp>>(patterns.getContext(), "cosf",
+                                                "cos", benefit);
+  patterns.add<ScalarOpToLibmCall<math::SinOp>>(patterns.getContext(), "sinf",
+                                                "sin", benefit);
 }
 
 namespace {
 struct ConvertMathToLibmPass
     : public ConvertMathToLibmBase<ConvertMathToLibmPass> {
   void runOnOperation() override;
 };
 } // namespace

mlir/test/Conversion/MathToLibm/convert-to-libm.mlir

 // RUN: mlir-opt %s -convert-math-to-libm -canonicalize | FileCheck %s
 
 // CHECK-DAG: @erf(f64) -> f64
 // CHECK-DAG: @erff(f32) -> f32
 // CHECK-DAG: @expm1(f64) -> f64
 // CHECK-DAG: @expm1f(f32) -> f32
 // CHECK-DAG: @atan2(f64, f64) -> f64
 // CHECK-DAG: @atan2f(f32, f32) -> f32
 // CHECK-DAG: @tanh(f64) -> f64
 // CHECK-DAG: @tanhf(f32) -> f32
 // CHECK-DAG: @round(f64) -> f64
 // CHECK-DAG: @roundf(f32) -> f32
+// CHECK-DAG: @cos(f64) -> f64
+// CHECK-DAG: @cosf(f32) -> f32
+// CHECK-DAG: @sin(f64) -> f64
+// CHECK-DAG: @sinf(f32) -> f32
 
 // CHECK-LABEL: func @tanh_caller
 // CHECK-SAME: %[[FLOAT:.*]]: f32
 // CHECK-SAME: %[[DOUBLE:.*]]: f64
 func.func @tanh_caller(%float: f32, %double: f64) -> (f32, f64)  {
   // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @tanhf(%[[FLOAT]]) : (f32) -> f32
   %float_result = math.tanh %float : f32
   // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @tanh(%[[DOUBLE]]) : (f64) -> f64

 func.func @round_caller(%float: f32, %double: f64) -> (f32, f64) {
   // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @roundf(%[[FLOAT]]) : (f32) -> f32
   %float_result = math.round %float : f32
   // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @round(%[[DOUBLE]]) : (f64) -> f64
   %double_result = math.round %double : f64
   // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
   return %float_result, %double_result : f32, f64
 }
+
+// CHECK-LABEL: func @cos_caller
+// CHECK-SAME: %[[FLOAT:.*]]: f32
+// CHECK-SAME: %[[DOUBLE:.*]]: f64
+func.func @cos_caller(%float: f32, %double: f64) -> (f32, f64)  {
+  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @cosf(%[[FLOAT]]) : (f32) -> f32
+  %float_result = math.cos %float : f32
+  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @cos(%[[DOUBLE]]) : (f64) -> f64
+  %double_result = math.cos %double : f64
+  // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
+  return %float_result, %double_result : f32, f64
+}
+
+// CHECK-LABEL: func @sin_caller
+// CHECK-SAME: %[[FLOAT:.*]]: f32
+// CHECK-SAME: %[[DOUBLE:.*]]: f64
+func.func @sin_caller(%float: f32, %double: f64) -> (f32, f64)  {
+  // CHECK-DAG: %[[FLOAT_RESULT:.*]] = call @sinf(%[[FLOAT]]) : (f32) -> f32
+  %float_result = math.sin %float : f32
+  // CHECK-DAG: %[[DOUBLE_RESULT:.*]] = call @sin(%[[DOUBLE]]) : (f64) -> f64
+  %double_result = math.sin %double : f64
+  // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
+  return %float_result, %double_result : f32, f64
+}