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
|
//===- ReifyValueBounds.cpp --- Reify value bounds with arith ops -------*-===//
//
// 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/Arith/Transforms/Transforms.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
using namespace mlir;
using namespace mlir::arith;
/// Build Arith IR for the given affine map and its operands.
static Value buildArithValue(OpBuilder &b, Location loc, AffineMap map,
ValueRange operands) {
assert(map.getNumResults() == 1 && "multiple results not supported yet");
std::function<Value(AffineExpr)> buildExpr = [&](AffineExpr e) -> Value {
switch (e.getKind()) {
case AffineExprKind::Constant:
return b.create<ConstantIndexOp>(loc,
cast<AffineConstantExpr>(e).getValue());
case AffineExprKind::DimId:
return operands[cast<AffineDimExpr>(e).getPosition()];
case AffineExprKind::SymbolId:
return operands[cast<AffineSymbolExpr>(e).getPosition() +
map.getNumDims()];
case AffineExprKind::Add: {
auto binaryExpr = cast<AffineBinaryOpExpr>(e);
return b.create<AddIOp>(loc, buildExpr(binaryExpr.getLHS()),
buildExpr(binaryExpr.getRHS()));
}
case AffineExprKind::Mul: {
auto binaryExpr = cast<AffineBinaryOpExpr>(e);
return b.create<MulIOp>(loc, buildExpr(binaryExpr.getLHS()),
buildExpr(binaryExpr.getRHS()));
}
case AffineExprKind::FloorDiv: {
auto binaryExpr = cast<AffineBinaryOpExpr>(e);
return b.create<DivSIOp>(loc, buildExpr(binaryExpr.getLHS()),
buildExpr(binaryExpr.getRHS()));
}
case AffineExprKind::CeilDiv: {
auto binaryExpr = cast<AffineBinaryOpExpr>(e);
return b.create<CeilDivSIOp>(loc, buildExpr(binaryExpr.getLHS()),
buildExpr(binaryExpr.getRHS()));
}
case AffineExprKind::Mod: {
auto binaryExpr = cast<AffineBinaryOpExpr>(e);
return b.create<RemSIOp>(loc, buildExpr(binaryExpr.getLHS()),
buildExpr(binaryExpr.getRHS()));
}
}
llvm_unreachable("unsupported AffineExpr kind");
};
return buildExpr(map.getResult(0));
}
static FailureOr<OpFoldResult>
reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
Value value, std::optional<int64_t> dim,
ValueBoundsConstraintSet::StopConditionFn stopCondition,
bool closedUB) {
// Compute bound.
AffineMap boundMap;
ValueDimList mapOperands;
if (failed(ValueBoundsConstraintSet::computeBound(
boundMap, mapOperands, type, value, dim, stopCondition, closedUB)))
return failure();
// Materialize tensor.dim/memref.dim ops.
SmallVector<Value> operands;
for (auto valueDim : mapOperands) {
Value value = valueDim.first;
std::optional<int64_t> dim = valueDim.second;
if (!dim.has_value()) {
// This is an index-typed value.
assert(value.getType().isIndex() && "expected index type");
operands.push_back(value);
continue;
}
assert(cast<ShapedType>(value.getType()).isDynamicDim(*dim) &&
"expected dynamic dim");
if (isa<RankedTensorType>(value.getType())) {
// A tensor dimension is used: generate a tensor.dim.
operands.push_back(b.create<tensor::DimOp>(loc, value, *dim));
} else if (isa<MemRefType>(value.getType())) {
// A memref dimension is used: generate a memref.dim.
operands.push_back(b.create<memref::DimOp>(loc, value, *dim));
} else {
llvm_unreachable("cannot generate DimOp for unsupported shaped type");
}
}
// Check for special cases where no arith ops are needed.
if (boundMap.isSingleConstant()) {
// Bound is a constant: return an IntegerAttr.
return static_cast<OpFoldResult>(
b.getIndexAttr(boundMap.getSingleConstantResult()));
}
// No arith ops are needed if the bound is a single SSA value.
if (auto expr = dyn_cast<AffineDimExpr>(boundMap.getResult(0)))
return static_cast<OpFoldResult>(operands[expr.getPosition()]);
if (auto expr = dyn_cast<AffineSymbolExpr>(boundMap.getResult(0)))
return static_cast<OpFoldResult>(
operands[expr.getPosition() + boundMap.getNumDims()]);
// General case: build Arith ops.
return static_cast<OpFoldResult>(buildArithValue(b, loc, boundMap, operands));
}
FailureOr<OpFoldResult> mlir::arith::reifyShapedValueDimBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
int64_t dim, ValueBoundsConstraintSet::StopConditionFn stopCondition,
bool closedUB) {
auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
ValueBoundsConstraintSet &cstr) {
// We are trying to reify a bound for `value` in terms of the owning op's
// operands. Construct a stop condition that evaluates to "true" for any SSA
// value expect for `value`. I.e., the bound will be computed in terms of
// any SSA values expect for `value`. The first such values are operands of
// the owner of `value`.
return v != value;
};
return reifyValueBound(b, loc, type, value, dim,
stopCondition ? stopCondition : reifyToOperands,
closedUB);
}
FailureOr<OpFoldResult> mlir::arith::reifyIndexValueBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) {
auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
ValueBoundsConstraintSet &cstr) {
return v != value;
};
return reifyValueBound(b, loc, type, value, /*dim=*/std::nullopt,
stopCondition ? stopCondition : reifyToOperands,
closedUB);
}
|
