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
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
|
//===- TestReifyValueBounds.cpp - Test value bounds reification -----------===//
//
// 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 "TestDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Affine/Transforms/Transforms.h"
#include "mlir/Dialect/Arith/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "mlir/Pass/Pass.h"
#define PASS_NAME "test-affine-reify-value-bounds"
using namespace mlir;
using namespace mlir::affine;
using mlir::presburger::BoundType;
namespace {
/// This pass applies the permutation on the first maximal perfect nest.
struct TestReifyValueBounds
: public PassWrapper<TestReifyValueBounds, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestReifyValueBounds)
StringRef getArgument() const final { return PASS_NAME; }
StringRef getDescription() const final {
return "Tests ValueBoundsOpInterface with affine dialect reification";
}
TestReifyValueBounds() = default;
TestReifyValueBounds(const TestReifyValueBounds &pass) : PassWrapper(pass){};
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<affine::AffineDialect, tensor::TensorDialect,
memref::MemRefDialect>();
}
void runOnOperation() override;
private:
Option<bool> reifyToFuncArgs{
*this, "reify-to-func-args",
llvm::cl::desc("Reify in terms of function args"), llvm::cl::init(false)};
Option<bool> useArithOps{*this, "use-arith-ops",
llvm::cl::desc("Reify with arith dialect ops"),
llvm::cl::init(false)};
};
} // namespace
static ValueBoundsConstraintSet::ComparisonOperator
invertComparisonOperator(ValueBoundsConstraintSet::ComparisonOperator cmp) {
if (cmp == ValueBoundsConstraintSet::ComparisonOperator::LT)
return ValueBoundsConstraintSet::ComparisonOperator::GE;
if (cmp == ValueBoundsConstraintSet::ComparisonOperator::LE)
return ValueBoundsConstraintSet::ComparisonOperator::GT;
if (cmp == ValueBoundsConstraintSet::ComparisonOperator::GT)
return ValueBoundsConstraintSet::ComparisonOperator::LE;
if (cmp == ValueBoundsConstraintSet::ComparisonOperator::GE)
return ValueBoundsConstraintSet::ComparisonOperator::LT;
llvm_unreachable("unsupported comparison operator");
}
/// Look for "test.reify_bound" ops in the input and replace their results with
/// the reified values.
static LogicalResult testReifyValueBounds(func::FuncOp funcOp,
bool reifyToFuncArgs,
bool useArithOps) {
IRRewriter rewriter(funcOp.getContext());
WalkResult result = funcOp.walk([&](test::ReifyBoundOp op) {
auto boundType = op.getBoundType();
Value value = op.getVar();
std::optional<int64_t> dim = op.getDim();
bool constant = op.getConstant();
bool scalable = op.getScalable();
// Prepare stop condition. By default, reify in terms of the op's
// operands. No stop condition is used when a constant was requested.
std::function<bool(Value, std::optional<int64_t>,
ValueBoundsConstraintSet & cstr)>
stopCondition = [&](Value v, std::optional<int64_t> d,
ValueBoundsConstraintSet &cstr) {
// Reify in terms of SSA values that are different from `value`.
return v != value;
};
if (reifyToFuncArgs) {
// Reify in terms of function block arguments.
stopCondition = [](Value v, std::optional<int64_t> d,
ValueBoundsConstraintSet &cstr) {
auto bbArg = dyn_cast<BlockArgument>(v);
if (!bbArg)
return false;
return isa<FunctionOpInterface>(bbArg.getParentBlock()->getParentOp());
};
}
// Reify value bound
rewriter.setInsertionPointAfter(op);
FailureOr<OpFoldResult> reified = failure();
if (constant) {
auto reifiedConst = ValueBoundsConstraintSet::computeConstantBound(
boundType, {value, dim}, /*stopCondition=*/nullptr);
if (succeeded(reifiedConst))
reified = FailureOr<OpFoldResult>(rewriter.getIndexAttr(*reifiedConst));
} else if (scalable) {
auto loc = op->getLoc();
auto reifiedScalable =
vector::ScalableValueBoundsConstraintSet::computeScalableBound(
value, dim, *op.getVscaleMin(), *op.getVscaleMax(), boundType);
if (succeeded(reifiedScalable)) {
SmallVector<std::pair<Value, std::optional<int64_t>>, 1> vscaleOperand;
if (reifiedScalable->map.getNumInputs() == 1) {
// The only possible input to the bound is vscale.
vscaleOperand.push_back(std::make_pair(
rewriter.create<vector::VectorScaleOp>(loc), std::nullopt));
}
reified = affine::materializeComputedBound(
rewriter, loc, reifiedScalable->map, vscaleOperand);
}
} else {
if (useArithOps) {
reified = arith::reifyValueBound(rewriter, op->getLoc(), boundType,
op.getVariable(), stopCondition);
} else {
reified = reifyValueBound(rewriter, op->getLoc(), boundType,
op.getVariable(), stopCondition);
}
}
if (failed(reified)) {
op->emitOpError("could not reify bound");
return WalkResult::interrupt();
}
// Replace the op with the reified bound.
if (auto val = llvm::dyn_cast_if_present<Value>(*reified)) {
rewriter.replaceOp(op, val);
return WalkResult::skip();
}
Value constOp = rewriter.create<arith::ConstantIndexOp>(
op->getLoc(), cast<IntegerAttr>(reified->get<Attribute>()).getInt());
rewriter.replaceOp(op, constOp);
return WalkResult::skip();
});
return failure(result.wasInterrupted());
}
/// Look for "test.compare" ops and emit errors/remarks.
static LogicalResult testEquality(func::FuncOp funcOp) {
IRRewriter rewriter(funcOp.getContext());
WalkResult result = funcOp.walk([&](test::CompareOp op) {
auto cmpType = op.getComparisonOperator();
if (op.getCompose()) {
if (cmpType != ValueBoundsConstraintSet::EQ) {
op->emitOpError(
"comparison operator must be EQ when 'composed' is specified");
return WalkResult::interrupt();
}
FailureOr<int64_t> delta = affine::fullyComposeAndComputeConstantDelta(
op->getOperand(0), op->getOperand(1));
if (failed(delta)) {
op->emitError("could not determine equality");
} else if (*delta == 0) {
op->emitRemark("equal");
} else {
op->emitRemark("different");
}
return WalkResult::advance();
}
auto compare = [&](ValueBoundsConstraintSet::ComparisonOperator cmp) {
return ValueBoundsConstraintSet::compare(op.getLhs(), cmp, op.getRhs());
};
if (compare(cmpType)) {
op->emitRemark("true");
} else if (cmpType != ValueBoundsConstraintSet::EQ &&
compare(invertComparisonOperator(cmpType))) {
op->emitRemark("false");
} else if (cmpType == ValueBoundsConstraintSet::EQ &&
(compare(ValueBoundsConstraintSet::ComparisonOperator::LT) ||
compare(ValueBoundsConstraintSet::ComparisonOperator::GT))) {
op->emitRemark("false");
} else {
op->emitError("unknown");
}
return WalkResult::advance();
});
return failure(result.wasInterrupted());
}
void TestReifyValueBounds::runOnOperation() {
if (failed(
testReifyValueBounds(getOperation(), reifyToFuncArgs, useArithOps)))
signalPassFailure();
if (failed(testEquality(getOperation())))
signalPassFailure();
}
namespace mlir {
void registerTestAffineReifyValueBoundsPass() {
PassRegistration<TestReifyValueBounds>();
}
} // namespace mlir
|
