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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
|
//===- IndependenceTransforms.cpp - Make ops independent of values --------===//
//
// 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/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Transforms/Transforms.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
using namespace mlir;
using namespace mlir::memref;
/// Make the given OpFoldResult independent of all independencies.
static FailureOr<OpFoldResult> makeIndependent(OpBuilder &b, Location loc,
OpFoldResult ofr,
ValueRange independencies) {
if (ofr.is<Attribute>())
return ofr;
Value value = ofr.get<Value>();
AffineMap boundMap;
ValueDimList mapOperands;
if (failed(ValueBoundsConstraintSet::computeIndependentBound(
boundMap, mapOperands, presburger::BoundType::UB, value,
/*dim=*/std::nullopt, independencies, /*closedUB=*/true)))
return failure();
return affine::materializeComputedBound(b, loc, boundMap, mapOperands);
}
FailureOr<Value> memref::buildIndependentOp(OpBuilder &b,
memref::AllocaOp allocaOp,
ValueRange independencies) {
OpBuilder::InsertionGuard g(b);
b.setInsertionPoint(allocaOp);
Location loc = allocaOp.getLoc();
SmallVector<OpFoldResult> newSizes;
for (OpFoldResult ofr : allocaOp.getMixedSizes()) {
auto ub = makeIndependent(b, loc, ofr, independencies);
if (failed(ub))
return failure();
newSizes.push_back(*ub);
}
// Return existing memref::AllocaOp if nothing has changed.
if (llvm::equal(allocaOp.getMixedSizes(), newSizes))
return allocaOp.getResult();
// Create a new memref::AllocaOp.
Value newAllocaOp =
b.create<AllocaOp>(loc, newSizes, allocaOp.getType().getElementType());
// Create a memref::SubViewOp.
SmallVector<OpFoldResult> offsets(newSizes.size(), b.getIndexAttr(0));
SmallVector<OpFoldResult> strides(newSizes.size(), b.getIndexAttr(1));
return b
.create<SubViewOp>(loc, newAllocaOp, offsets, allocaOp.getMixedSizes(),
strides)
.getResult();
}
/// Push down an UnrealizedConversionCastOp past a SubViewOp.
static UnrealizedConversionCastOp
propagateSubViewOp(RewriterBase &rewriter,
UnrealizedConversionCastOp conversionOp, SubViewOp op) {
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(op);
auto newResultType = cast<MemRefType>(SubViewOp::inferRankReducedResultType(
op.getType().getShape(), op.getSourceType(), op.getMixedOffsets(),
op.getMixedSizes(), op.getMixedStrides()));
Value newSubview = rewriter.create<SubViewOp>(
op.getLoc(), newResultType, conversionOp.getOperand(0),
op.getMixedOffsets(), op.getMixedSizes(), op.getMixedStrides());
auto newConversionOp = rewriter.create<UnrealizedConversionCastOp>(
op.getLoc(), op.getType(), newSubview);
rewriter.replaceAllUsesWith(op.getResult(), newConversionOp->getResult(0));
return newConversionOp;
}
/// Given an original op and a new, modified op with the same number of results,
/// whose memref return types may differ, replace all uses of the original op
/// with the new op and propagate the new memref types through the IR.
///
/// Example:
/// %from = memref.alloca(%sz) : memref<?xf32>
/// %to = memref.subview ... : ... to memref<?xf32, strided<[1], offset: ?>>
/// memref.store %cst, %from[%c0] : memref<?xf32>
///
/// In the above example, all uses of %from are replaced with %to. This can be
/// done directly for ops such as memref.store. For ops that have memref results
/// (e.g., memref.subview), the result type may depend on the operand type, so
/// we cannot just replace all uses. There is special handling for common memref
/// ops. For all other ops, unrealized_conversion_cast is inserted.
static void replaceAndPropagateMemRefType(RewriterBase &rewriter,
Operation *from, Operation *to) {
assert(from->getNumResults() == to->getNumResults() &&
"expected same number of results");
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPointAfter(to);
// Wrap new results in unrealized_conversion_cast and replace all uses of the
// original op.
SmallVector<UnrealizedConversionCastOp> unrealizedConversions;
for (const auto &it :
llvm::enumerate(llvm::zip(from->getResults(), to->getResults()))) {
unrealizedConversions.push_back(rewriter.create<UnrealizedConversionCastOp>(
to->getLoc(), std::get<0>(it.value()).getType(),
std::get<1>(it.value())));
rewriter.replaceAllUsesWith(from->getResult(it.index()),
unrealizedConversions.back()->getResult(0));
}
// Push unrealized_conversion_cast ops further down in the IR. I.e., try to
// wrap results instead of operands in a cast.
for (int i = 0; i < static_cast<int>(unrealizedConversions.size()); ++i) {
UnrealizedConversionCastOp conversion = unrealizedConversions[i];
assert(conversion->getNumOperands() == 1 &&
conversion->getNumResults() == 1 &&
"expected single operand and single result");
SmallVector<Operation *> users = llvm::to_vector(conversion->getUsers());
for (Operation *user : users) {
// Handle common memref dialect ops that produce new memrefs and must
// be recreated with the new result type.
if (auto subviewOp = dyn_cast<SubViewOp>(user)) {
unrealizedConversions.push_back(
propagateSubViewOp(rewriter, conversion, subviewOp));
continue;
}
// TODO: Other memref ops such as memref.collapse_shape/expand_shape
// should also be handled here.
// Skip any ops that produce MemRef result or have MemRef region block
// arguments. These may need special handling (e.g., scf.for).
if (llvm::any_of(user->getResultTypes(),
[](Type t) { return isa<MemRefType>(t); }))
continue;
if (llvm::any_of(user->getRegions(), [](Region &r) {
return llvm::any_of(r.getArguments(), [](BlockArgument bbArg) {
return isa<MemRefType>(bbArg.getType());
});
}))
continue;
// For all other ops, we assume that we can directly replace the operand.
// This may have to be revised in the future; e.g., there may be ops that
// do not support non-identity layout maps.
for (OpOperand &operand : user->getOpOperands()) {
if (auto castOp =
operand.get().getDefiningOp<UnrealizedConversionCastOp>()) {
rewriter.updateRootInPlace(
user, [&]() { operand.set(conversion->getOperand(0)); });
}
}
}
}
// Erase all unrealized_conversion_cast ops without uses.
for (auto op : unrealizedConversions)
if (op->getUses().empty())
rewriter.eraseOp(op);
}
FailureOr<Value> memref::replaceWithIndependentOp(RewriterBase &rewriter,
memref::AllocaOp allocaOp,
ValueRange independencies) {
auto replacement =
memref::buildIndependentOp(rewriter, allocaOp, independencies);
if (failed(replacement))
return failure();
replaceAndPropagateMemRefType(rewriter, allocaOp,
replacement->getDefiningOp());
return replacement;
}
|
