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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
|
//===- FoldAddIntoDest.cpp ---------------------------------------*- 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
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/IR/Dominance.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
using namespace mlir;
// Determine whether the value is defined to be zero.
static bool isDefinedAsZero(Value val) {
if (!val)
return false;
// Check whether val is a constant scalar / vector splat / tensor splat float
// or integer zero.
if (matchPattern(val, m_AnyZeroFloat()) || matchPattern(val, m_Zero()))
return true;
return TypeSwitch<Operation *, bool>(val.getDefiningOp())
.Case<linalg::FillOp, linalg::CopyOp>([&](auto op) {
return op && op.getInputs().size() == 1 &&
isDefinedAsZero(op.getInputs()[0]);
})
.Default([&](auto) { return false; });
}
/// Replace a linalg.add with one operand the single user of a contraction,
/// which has a zero-filled, "identity-mapped" destination and is dominated by
/// the `other` operand, by the contraction with `other` as its dest.
///
/// As an example, the following pseudo-code will be rewritten
/// %cst = arith.constant 0.000000e+00
/// %empty = tensor.empty()
/// %zeroed = linalg.fill ins(%cst : f32) outs(%empty : !type) -> !type
/// %C = linalg.matmul ins(%A, %B) outs(%zeroed)
/// %empty2 = tensor.empty()
/// %zeroed2 = linalg.fill ins(%cst : f32) outs(%empty2 : !type) -> !type
/// %F = linalg.matmul ins(%D, %E) outs(%zeroed2)
/// %out = linalg.add ins(%C, %F) outs(%empty)
/// to:
/// %cst = arith.constant 0.000000e+00
/// %empty = tensor.empty()
/// %zeroed = linalg.fill ins(%cst : f32) outs(%empty : !type) -> !type
/// %C = linalg.matmul ins(%A, %B) outs(%zeroed)
/// %out = linalg.matmul ins(%D, %E) outs(%C)
///
struct FoldAddIntoDest final : public OpRewritePattern<linalg::AddOp> {
using OpRewritePattern<linalg::AddOp>::OpRewritePattern;
LogicalResult matchAndRewrite(linalg::AddOp addOp,
PatternRewriter &rewriter) const override {
// For now, pattern only applies on tensor types (memref support is TODO).
if (!addOp.hasPureTensorSemantics())
return failure();
Value dominatingOperand = nullptr;
linalg::LinalgOp dominatedOp = nullptr;
{ // We will forget about which operand was left or right after this block.
Value lhs = addOp.getInputs()[0];
Value rhs = addOp.getInputs()[1];
// Can only put one of addOp's operands in the dest/out arg of the other's
// defining op based on suitable dominance.
// TODO: Can be generalized to move ops around as long as that still
// respects use-def chains and doesn't affect side-effects.
if (auto rhsOp = rhs.getDefiningOp<linalg::LinalgOp>()) {
DominanceInfo domInfo(rhsOp);
if (domInfo.properlyDominates(lhs, rhsOp)) {
dominatingOperand = lhs;
dominatedOp = rhsOp;
}
}
if (auto lhsOp = lhs.getDefiningOp<linalg::LinalgOp>()) {
DominanceInfo domInfo(lhsOp);
if (domInfo.properlyDominates(rhs, lhsOp)) {
dominatingOperand = rhs;
dominatedOp = lhsOp;
}
}
if (!dominatingOperand || !dominatedOp)
return failure();
// NB: As linalg.add's generalisation ignores the out argument in its
// region there is no need to perform checks on addOp's out argument.
}
// When dominated op is a contraction we know it accumulates on its out arg.
// E.g., AddOp is not a contraction and hence ignores its out arg's value.
// TODO: Generalize check to also pass in case of other LinalgOps that
// accumulate on their out arg but are not (binary) contraction ops.
auto dominatedDestOp =
dyn_cast<DestinationStyleOpInterface>((Operation *)dominatedOp);
if (dominatedOp->getNumResults() != 1 ||
!linalg::isaContractionOpInterface(dominatedOp) ||
(!dominatedDestOp || dominatedDestOp.getNumDpsInits() != 1))
return rewriter.notifyMatchFailure(
dominatedOp, "expected dominated op to be single-result "
"destination-passing contraction");
// To change the contraction's result, `addOp` must be its only user.
if (!dominatedOp->getResult(0).hasOneUse())
return rewriter.notifyMatchFailure(
dominatedOp,
"expected linalg.add to be single user of contraction's result");
// As `dominatedOp` was already accumulating on its out argument, it is only
// safe to no longer use its current out arg when it is the additive ident.
auto *destOperand = dominatedDestOp.getDpsInitOperand(0);
if (!isDefinedAsZero(destOperand->get()))
return rewriter.notifyMatchFailure(
dominatedOp, "expected dominated op's dest to be additive zero");
// TODO: If the other op is a contraction and has additive ident as dest, we
// can swap the dests and achieve the proper sum, given suitable dominance.
// As an operand to `addOp`, `dominatingOperand` has an identity affine_map.
// Hence, we can only substitute `dominatingOperand` for the dest of the
// contraction when dest's indexing_map corresponds to an identity map
// w.r.t. just the dimensions of dest, i.e. is an ordered projection.
SmallVector<AffineMap> indexMaps = dominatedOp.getIndexingMapsArray();
int prevDimPos = -1;
for (auto expr : indexMaps[destOperand->getOperandNumber()].getResults()) {
auto dim = dyn_cast<AffineDimExpr>(expr);
if (!dim || prevDimPos > static_cast<int>(dim.getPosition()))
return rewriter.notifyMatchFailure(
dominatedOp, "expected index_map for contraction's dest to be an "
"ordered projection");
prevDimPos = dim.getPosition();
}
// Replace the additive-ident, i.e. zero, out arg of the dominated op by the
// dominating summand. This makes the dominated op's result the sum of both
// of addOp's arguments - therefore we replace addOp and it uses by it.
rewriter.modifyOpInPlace(
dominatedOp, [&]() { dominatedOp->setOperand(2, dominatingOperand); });
rewriter.replaceAllOpUsesWith(addOp, dominatedOp->getResult(0));
return success();
}
};
void linalg::populateFoldAddIntoDestPatterns(RewritePatternSet &patterns) {
// Replace linalg.add when destination passing suffices for achieving the sum.
patterns.add<FoldAddIntoDest>(patterns.getContext());
}
|
