1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
|
//===- UpliftToFMA.cpp - Arith to FMA uplifting ---------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements uplifting from arith ops to math.fma.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Math/Transforms/Passes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
namespace mlir::math {
#define GEN_PASS_DEF_MATHUPLIFTTOFMA
#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
} // namespace mlir::math
using namespace mlir;
template <typename Op>
static bool isValidForFMA(Op op) {
return static_cast<bool>(op.getFastmath() & arith::FastMathFlags::contract);
}
namespace {
struct UpliftFma final : OpRewritePattern<arith::AddFOp> {
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(arith::AddFOp op,
PatternRewriter &rewriter) const override {
if (!isValidForFMA(op))
return rewriter.notifyMatchFailure(op, "addf op is not suitable for fma");
Value c;
arith::MulFOp ab;
if ((ab = op.getLhs().getDefiningOp<arith::MulFOp>())) {
c = op.getRhs();
} else if ((ab = op.getRhs().getDefiningOp<arith::MulFOp>())) {
c = op.getLhs();
} else {
return rewriter.notifyMatchFailure(op, "no mulf op");
}
if (!isValidForFMA(ab))
return rewriter.notifyMatchFailure(ab, "mulf op is not suitable for fma");
Value a = ab.getLhs();
Value b = ab.getRhs();
arith::FastMathFlags fmf = op.getFastmath() & ab.getFastmath();
rewriter.replaceOpWithNewOp<math::FmaOp>(op, a, b, c, fmf);
return success();
}
};
struct MathUpliftToFMA final
: math::impl::MathUpliftToFMABase<MathUpliftToFMA> {
using MathUpliftToFMABase::MathUpliftToFMABase;
void runOnOperation() override {
RewritePatternSet patterns(&getContext());
populateUpliftToFMAPatterns(patterns);
if (failed(
applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
return signalPassFailure();
}
};
} // namespace
void mlir::populateUpliftToFMAPatterns(RewritePatternSet &patterns) {
patterns.insert<UpliftFma>(patterns.getContext());
}
|
