[mlir][math] Math expansion for math.tan

[rsuderman]

Summary

We can implement a polynomial approximation of math.tan by
decomposing to math.sin and math.cos. While it is not
technically a polynomial approximation it should be the most
straight forward approximation.

[jpienaar]

Could we just have it in its own group/approximations pass? (Well and we can populate with patterns from polynomial if needed). Just feels wrong to have this in polynomial

[rsuderman]

Moved decomposition to ExpandPattterns

[rsuderman]

In D144980#4159308, @jpienaar wrote:

Could we just have it in its own group/approximations pass? (Well and we can populate with patterns from polynomial if needed). Just feels wrong to have this in polynomial

I moved it to the ExpandMath patterns. Should be reasonable enough and it can be included optionally this way.

[jpienaar]

Thanks!

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

 //===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_
 
 namespace mlir {
 
 class RewritePatternSet;
 
 void populateExpandCtlzPattern(RewritePatternSet &patterns);
+void populateExpandTanPattern(RewritePatternSet &patterns);
 void populateExpandTanhPattern(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

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/Math/Transforms/Passes.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 /// Expands tanh op into
 ///   1) 1-exp^{-2x} / 1+exp^{-2x}, if x => 0
 ///   2) exp^{2x}-1 / exp^{2x}+1  , if x < 0
 static LogicalResult convertTanhOp(math::TanhOp op, PatternRewriter &rewriter) {
static LogicalResult convertTanhOp(math::TanhOp op, PatternRewriter &rewriter) {
   Value zero = rewriter.create<arith::ConstantOp>(loc, floatZero);
   Value cmpRes = rewriter.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OGE,
                                                 op.getOperand(), zero);
   rewriter.replaceOpWithNewOp<arith::SelectOp>(op, cmpRes, positiveRes,
                                                negativeRes);
   return success();
 }
 
+static LogicalResult convertTanOp(math::TanOp op, PatternRewriter &rewriter) {
+  ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+  Value operand = op.getOperand();
+  Type type = operand.getType();
+  Value sin = b.create<math::SinOp>(type, operand);
+  Value cos = b.create<math::CosOp>(type, operand);
+  Value div = b.create<arith::DivFOp>(type, sin, cos);
+  rewriter.replaceOp(op, div);
+  return success();
+}
+
 static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
                                    PatternRewriter &rewriter) {
   auto operand = op.getOperand();
   auto elementTy = operand.getType();
   auto resultTy = op.getType();
   Location loc = op.getLoc();
 
   int bitWidth = elementTy.getIntOrFloatBitWidth();
static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
   rewriter.replaceOp(op, whileOp->getResult(1));
   return success();
 }
 
 void mlir::populateExpandCtlzPattern(RewritePatternSet &patterns) {
   patterns.add(convertCtlzOp);
 }
 
+void mlir::populateExpandTanPattern(RewritePatternSet &patterns) {
+  patterns.add(convertTanOp);
+}
+
 void mlir::populateExpandTanhPattern(RewritePatternSet &patterns) {
   patterns.add(convertTanhOp);
 }

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

 // CHECK: %[[DIVISOR2:.+]] = arith.addf %[[EXP2]], %[[ONE]] : f32
 // CHECK: %[[RES2:.+]] = arith.divf %[[DIVIDEND2]], %[[DIVISOR2]] : f32
 // CHECK: %[[COND:.+]] = arith.cmpf oge, %arg0, %[[ZERO]] : f32
 // CHECK: %[[RESULT:.+]] = arith.select %[[COND]], %[[RES1]], %[[RES2]] : f32
 // CHECK: return %[[RESULT]]
 
 // -----
 
+// CHECK-LABEL: func @tan
+func.func @tan(%arg: f32) -> f32 {
+  %res = math.tan %arg : f32
+  return %res : f32
+}
+
+// CHECK-SAME: %[[ARG0:.+]]: f32
+// CHECK: %[[SIN:.+]] = math.sin %[[ARG0]]
+// CHECK: %[[COS:.+]] = math.cos %[[ARG0]]
+// CEHCK: %[[DIV:.+]] = arith.div %[[SIN]] %[[COS]]
+
+// -----
+
 // CHECK-LABEL: func @ctlz
 func.func @ctlz(%arg: i32) -> i32 {
   // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : i32
   // CHECK-DAG: %[[C32:.+]] = arith.constant 32 : i32
   // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : i32
   // CHECK: %[[WHILE:.+]]:3 = scf.while (%[[A1:.+]] = %arg0, %[[A2:.+]] = %[[C32]], %[[A3:.+]] = %[[C0]])
   // CHECK:   %[[CMP:.+]] = arith.cmpi ne, %[[A1]], %[[A3]]
   // CHECK:   scf.condition(%[[CMP]]) %[[A1]], %[[A2]], %[[A3]]
   // CHECK:   %[[SHR:.+]] = arith.shrui %[[A1]], %[[C1]]
   // CHECK:   %[[SUB:.+]] = arith.subi %[[A2]], %[[C1]]
   // CHECK:   scf.yield %[[SHR]], %[[SUB]], %[[A3]]
   %res = math.ctlz %arg : i32
 
   // CHECK: return %[[WHILE]]#1
   return %res : i32
 }

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

struct TestExpandMathPass
   }
   StringRef getDescription() const final { return "Test expanding math"; }
 };
 } // namespace
 
 void TestExpandMathPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   populateExpandCtlzPattern(patterns);
+  populateExpandTanPattern(patterns);
   populateExpandTanhPattern(patterns);
   (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }
 
 namespace mlir {
 namespace test {
 void registerTestExpandMathPass() { PassRegistration<TestExpandMathPass>(); }
 } // namespace test
 } // namespace mlir