Add polynomial approximation for trigonometric sine and cosine functions

[asaadaldien]

Summary

The approximation relays on range reduced version y \in [0, pi/2]. An input x will have
the property that sin(x) = sin(y), -sin(y), cos(y), -cos(y) depends on which quadrable x
is in, where sin(y) and cos(y) are approximated with 5th degree polynomial (of x^2).
As a result a single pattern can be used to compute approximation for both sine and cosine.

[ezhulenev]

There are bunch of variable name warning, I'd say it's ok to keep CC/KR/etc upper case but cast_f32_to_i32 and such are better to keep consistent

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

639

private:
  static constexpr bool isSine = std::is_same<OpTy, math::SinOp>::value;

+ static_assert that OpTy is Sin or Cos

[asaadaldien]

consistant constant / utility func naming

[asaadaldien]

assert cosOp&sinOp types and fix clang-tidy warning

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

ExpM1Approximation::matchAndRewrite(math::ExpM1Op op,
   expm1 = builder.create<SelectOp>(isInf, u, expm1);
   Value approximation = builder.create<SelectOp>(
       uEqOne, x, builder.create<SelectOp>(uMinusOneEqNegOne, cstNegOne, expm1));
   rewriter.replaceOp(op, approximation);
   return success();
 }
 
 //----------------------------------------------------------------------------//
+// Sin and Cos approximation.
+//----------------------------------------------------------------------------//
+
+namespace {
+
+template <bool isSine, typename OpTy>
+struct SinAndCosApproximation : public OpRewritePattern<OpTy> {
+public:
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op, PatternRewriter &rewriter) const final;
+};
+} // namespace
+
+#define TWO_OVER_PI                                                            \
+  0.6366197723675813430755350534900574481378385829618257949906693762L
+#define PI_OVER_2                                                              \
+  1.5707963267948966192313216916397514420985846996875529104874722961L
+
+// Approximates sin(x) or cos(x) by finding the best approximation polynomial in
+// the reduced range [0, pi/2] for both sin(x) and cos(x). Then given y in the
+// reduced range sin(x) will be computed as sin(y), -sin(y), cos(y) or -cos(y).
+template <bool isSine, typename OpTy>
+LogicalResult SinAndCosApproximation<isSine, OpTy>::matchAndRewrite(
+    OpTy op, PatternRewriter &rewriter) const {
+  auto width = vectorWidth(op.operand().getType(), isF32);
+  if (!width.hasValue())
+    return rewriter.notifyMatchFailure(op, "unsupported operand type");
+
+  ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
+  auto bcast = [&](Value value) -> Value {
+    return broadcast(builder, value, *width);
+  };
+  auto mul = [&](Value a, Value b) -> Value {
+    return builder.create<MulFOp>(a, b);
+  };
+  auto sub = [&](Value a, Value b) -> Value {
+    return builder.create<SubFOp>(a, b);
+  };
+  auto floor = [&](Value a) { return builder.create<FloorFOp>(a); };
+
+  auto i32Vec = broadcast(builder.getI32Type(), *width);
+  auto cast_f32_to_i32 = [&](Value a) -> Value {
+    return builder.create<FPToSIOp>(a, i32Vec);
+  };
+
+  auto modulo_4 = [&](Value a) -> Value {
+    return builder.create<AndOp>(a, bcast(i32Cst(builder, 3)));
+  };
+
+  auto isEqualTo = [&](Value a, Value b) -> Value {
+    return builder.create<CmpIOp>(CmpIPredicate::eq, a, b);
+  };
+
+  auto isGreaterThan = [&](Value a, Value b) -> Value {
+    return builder.create<CmpIOp>(CmpIPredicate::sgt, a, b);
+  };
+
+  auto select = [&](Value cond, Value t, Value f) -> Value {
+    return builder.create<SelectOp>(cond, t, f);
+  };
+
+  auto fmla = [&](Value a, Value b, Value c) {
+    return builder.create<FmaFOp>(a, b, c);
+  };
+
+  auto Or = [&](Value a, Value b) { return builder.create<OrOp>(a, b); };
+
+  Value twoOverPi = bcast(f32Cst(builder, TWO_OVER_PI));
+  Value piOverTwo = bcast(f32Cst(builder, PI_OVER_2));
+
+  Value x = op.operand();
+
+  Value k = floor(mul(x, twoOverPi));
+
+  Value y = sub(x, mul(k, piOverTwo));
+
+  Value One = bcast(f32Cst(builder, 1.0));
+  Value nOne = bcast(f32Cst(builder, -1.0));
+
+  Value SC2 = bcast(f32Cst(builder, -0.16666667163372039794921875f));
+  Value SC4 = bcast(f32Cst(builder, 8.333347737789154052734375e-3f));
+  Value SC6 = bcast(f32Cst(builder, -1.9842604524455964565277099609375e-4f));
+  Value SC8 = bcast(f32Cst(builder, 2.760012648650445044040679931640625e-6f));
+  Value SC10 =
+      bcast(f32Cst(builder, -2.50293279435709337121807038784027099609375e-8f));
+
+  Value CC2 = bcast(f32Cst(builder, -0.5f));
+  Value CC4 = bcast(f32Cst(builder, 4.166664183139801025390625e-2f));
+  Value CC6 = bcast(f32Cst(builder, -1.388833043165504932403564453125e-3f));
+  Value CC8 = bcast(f32Cst(builder, 2.47562347794882953166961669921875e-5f));
+  Value CC10 =
+      bcast(f32Cst(builder, -2.59630184018533327616751194000244140625e-7f));
+
+  Value kmod4 = modulo_4(cast_f32_to_i32(k));
+
+  Value KR0 = isEqualTo(kmod4, bcast(i32Cst(builder, 0)));
+  Value KR1 = isEqualTo(kmod4, bcast(i32Cst(builder, 1)));
+  Value KR2 = isEqualTo(kmod4, bcast(i32Cst(builder, 2)));
+  Value KR3 = isEqualTo(kmod4, bcast(i32Cst(builder, 3)));
+
+  Value sinuseCos = isSine ? Or(KR1, KR3) : Or(KR0, KR2);
+  Value negativeRange =
+      isSine ? isGreaterThan(kmod4, bcast(i32Cst(builder, 1))) : Or(KR1, KR2);
+
+  Value y2 = mul(y, y);
+
+  Value base = select(sinuseCos, One, y);
+  Value C2 = select(sinuseCos, CC2, SC2);
+  Value C4 = select(sinuseCos, CC4, SC4);
+  Value C6 = select(sinuseCos, CC6, SC6);
+  Value C8 = select(sinuseCos, CC8, SC8);
+  Value C10 = select(sinuseCos, CC10, SC10);
+
+  Value v1 = fmla(y2, C10, C8);
+  Value v2 = fmla(y2, v1, C6);
+  Value v3 = fmla(y2, v2, C4);
+  Value v4 = fmla(y2, v3, C2);
+  Value v5 = fmla(y2, v4, One);
+  Value v6 = mul(base, v5);
+
+  Value approximation = select(negativeRange, mul(nOne, v6), v6);
+
+  rewriter.replaceOp(op, approximation);
+
+  return success();
+}
+
+//----------------------------------------------------------------------------//
 
 void mlir::populateMathPolynomialApproximationPatterns(
     RewritePatternSet &patterns) {
   patterns.add<TanhApproximation, LogApproximation, Log2Approximation,
-               Log1pApproximation, ExpApproximation, ExpM1Approximation>(
+               Log1pApproximation, ExpApproximation, ExpM1Approximation,
+               SinAndCosApproximation<true, math::SinOp>,
+               SinAndCosApproximation<false, math::CosOp>>(
       patterns.getContext());
 }

mlir/test/mlir-cpu-runner/math_polynomial_approx.mlir

func @expm1() {
 
   // CHECK: -1, inf, 1e-10
   %special_vec = constant dense<[0xff800000, 0x7f800000, 1.0e-10]> : vector<3xf32>
   %log_special_vec = math.expm1 %special_vec : vector<3xf32>
   vector.print %log_special_vec : vector<3xf32>
 
   return
 }
+// -------------------------------------------------------------------------- //
+// Sin.
+// -------------------------------------------------------------------------- //
+func @sin() {
+  // CHECK: 0
+  %0 = constant 0.0 : f32
+  %sin_0 = math.sin %0 : f32
+  vector.print %sin_0 : f32
+
+  // CHECK: 0.707107
+  %pi_over_4 = constant 0.78539816339 : f32
+  %sin_pi_over_4 = math.sin %pi_over_4 : f32
+  vector.print %sin_pi_over_4 : f32
+
+  // CHECK: 1
+  %pi_over_2 = constant 1.57079632679 : f32
+  %sin_pi_over_2 = math.sin %pi_over_2 : f32
+  vector.print %sin_pi_over_2 : f32
+
+
+  // CHECK: 0
+  %pi = constant 3.14159265359 : f32
+  %sin_pi = math.sin %pi : f32
+  vector.print %sin_pi : f32
+
+  // CHECK: -1
+  %pi_3_over_2 = constant 4.71238898038 : f32
+  %sin_pi_3_over_2 = math.sin %pi_3_over_2 : f32
+  vector.print %sin_pi_3_over_2 : f32
+
+  // CHECK: 0, 0.866025, -1
+  %vec_x = constant dense<[9.42477796077, 2.09439510239, -1.57079632679]> : vector<3xf32>
+  %sin_vec_x = math.sin %vec_x : vector<3xf32>
+  vector.print %sin_vec_x : vector<3xf32>
+
+  return
+}
+
+// -------------------------------------------------------------------------- //
+// cos.
+// -------------------------------------------------------------------------- //
+
+func @cos() {
+  // CHECK: 1
+  %0 = constant 0.0 : f32
+  %cos_0 = math.cos %0 : f32
+  vector.print %cos_0 : f32
+
+  // CHECK: 0.707107
+  %pi_over_4 = constant 0.78539816339 : f32
+  %cos_pi_over_4 = math.cos %pi_over_4 : f32
+  vector.print %cos_pi_over_4 : f32
+
+  //// CHECK: 0
+  %pi_over_2 = constant 1.57079632679 : f32
+  %cos_pi_over_2 = math.cos %pi_over_2 : f32
+  vector.print %cos_pi_over_2 : f32
+
+  /// CHECK: -1
+  %pi = constant 3.14159265359 : f32
+  %cos_pi = math.cos %pi : f32
+  vector.print %cos_pi : f32
+
+  // CHECK: 0
+  %pi_3_over_2 = constant 4.71238898038 : f32
+  %cos_pi_3_over_2 = math.cos %pi_3_over_2 : f32
+  vector.print %cos_pi_3_over_2 : f32
+
+  // CHECK: -1, -0.5, 0
+  %vec_x = constant dense<[9.42477796077, 2.09439510239, -1.57079632679]> : vector<3xf32>
+  %cos_vec_x = math.cos %vec_x : vector<3xf32>
+  vector.print %cos_vec_x : vector<3xf32>
+
+
+  return
+}
+
 
 func @main() {
   call @tanh(): () -> ()
   call @log(): () -> ()
   call @log2(): () -> ()
   call @log1p(): () -> ()
   call @exp(): () -> ()
   call @expm1(): () -> ()
+  call @sin(): () -> ()
+  call @cos(): () -> ()
   return
 }