81 lines · cpp
1//===- SincosFusion.cpp - Fuse sin/cos into sincos -----------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/Math/IR/Math.h"10#include "mlir/Dialect/Math/Transforms/Passes.h"11#include "mlir/IR/PatternMatch.h"12#include "mlir/Transforms/GreedyPatternRewriteDriver.h"13 14using namespace mlir;15using namespace mlir::math;16 17namespace {18 19/// Fuse a math.sin and math.cos in the same block that use the same operand and20/// have identical fastmath flags into a single math.sincos.21struct SincosFusionPattern : OpRewritePattern<math::SinOp> {22 using Base::Base;23 24 LogicalResult matchAndRewrite(math::SinOp sinOp,25 PatternRewriter &rewriter) const override {26 Value operand = sinOp.getOperand();27 mlir::arith::FastMathFlags sinFastMathFlags = sinOp.getFastmath();28 29 math::CosOp cosOp = nullptr;30 sinOp->getBlock()->walk([&](math::CosOp op) {31 if (op.getOperand() == operand && op.getFastmath() == sinFastMathFlags) {32 cosOp = op;33 return WalkResult::interrupt();34 }35 return WalkResult::advance();36 });37 38 if (!cosOp)39 return failure();40 41 Operation *firstOp = sinOp->isBeforeInBlock(cosOp) ? sinOp.getOperation()42 : cosOp.getOperation();43 rewriter.setInsertionPoint(firstOp);44 45 Type elemType = sinOp.getType();46 auto sincos = math::SincosOp::create(rewriter, firstOp->getLoc(),47 TypeRange{elemType, elemType}, operand,48 sinOp.getFastmathAttr());49 50 rewriter.replaceOp(sinOp, sincos.getSin());51 rewriter.replaceOp(cosOp, sincos.getCos());52 return success();53 }54};55 56} // namespace57 58namespace mlir::math {59#define GEN_PASS_DEF_MATHSINCOSFUSIONPASS60#include "mlir/Dialect/Math/Transforms/Passes.h.inc"61} // namespace mlir::math62 63namespace {64 65struct MathSincosFusionPass final66 : math::impl::MathSincosFusionPassBase<MathSincosFusionPass> {67 using MathSincosFusionPassBase::MathSincosFusionPassBase;68 69 void runOnOperation() override {70 RewritePatternSet patterns(&getContext());71 patterns.add<SincosFusionPattern>(&getContext());72 73 GreedyRewriteConfig config;74 if (failed(75 applyPatternsGreedily(getOperation(), std::move(patterns), config)))76 return signalPassFailure();77 }78};79 80} // namespace81