[Mlir-commits] [mlir] [MLIR] Add sincos fusion pass (PR #161413)

[Slava Zakharin]

================
@@ -0,0 +1,77 @@
+//===- SincosFusion.cpp - Fuse sin/cos into sincos -----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::math;
+
+namespace {
+
+/// Fuse a math.sin and math.cos in the same block that use the same operand and
+/// have identical fastmath flags into a single math.sincos.
+struct SincosFusionPattern : OpRewritePattern<math::SinOp> {
+  using OpRewritePattern<math::SinOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(math::SinOp sinOp,
+                                PatternRewriter &rewriter) const override {
+    Value operand = sinOp.getOperand();
+    mlir::arith::FastMathFlags sinFastMathFlags = sinOp.getFastmath();
+
+    math::CosOp cosOp = nullptr;
+    sinOp->getBlock()->walk([&](math::CosOp op) {
+      if (op.getOperand() == operand && op.getFastmath() == sinFastMathFlags) {
+        cosOp = op;
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    });
+
+    if (!cosOp)
+      return failure();
+
+    Type elemType = sinOp.getType();
+    auto sincos = rewriter.create<math::SincosOp>(
+        sinOp.getLoc(), TypeRange{elemType, elemType}, operand,
+        sinOp.getFastmathAttr());
+
+    rewriter.replaceOp(sinOp, sincos.getSin());
----------------
vzakhari wrote:

Please try an example where `cos` precedes `sin`, and there is a use of `cos` before `sin`.

https://github.com/llvm/llvm-project/pull/161413