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
|
//===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
//
// 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/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
using namespace mlir;
namespace mlir {
namespace memref {
namespace {
template <typename OpTy>
struct AllocOpInterface
: public ValueBoundsOpInterface::ExternalModel<AllocOpInterface<OpTy>,
OpTy> {
void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
ValueBoundsConstraintSet &cstr) const {
auto allocOp = cast<OpTy>(op);
assert(value == allocOp.getResult() && "invalid value");
cstr.bound(value)[dim] == allocOp.getMixedSizes()[dim];
}
};
struct CastOpInterface
: public ValueBoundsOpInterface::ExternalModel<CastOpInterface, CastOp> {
void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
ValueBoundsConstraintSet &cstr) const {
auto castOp = cast<CastOp>(op);
assert(value == castOp.getResult() && "invalid value");
if (llvm::isa<MemRefType>(castOp.getResult().getType()) &&
llvm::isa<MemRefType>(castOp.getSource().getType())) {
cstr.bound(value)[dim] == cstr.getExpr(castOp.getSource(), dim);
}
}
};
struct DimOpInterface
: public ValueBoundsOpInterface::ExternalModel<DimOpInterface, DimOp> {
void populateBoundsForIndexValue(Operation *op, Value value,
ValueBoundsConstraintSet &cstr) const {
auto dimOp = cast<DimOp>(op);
assert(value == dimOp.getResult() && "invalid value");
auto constIndex = dimOp.getConstantIndex();
if (!constIndex.has_value())
return;
cstr.bound(value) == cstr.getExpr(dimOp.getSource(), *constIndex);
}
};
struct GetGlobalOpInterface
: public ValueBoundsOpInterface::ExternalModel<GetGlobalOpInterface,
GetGlobalOp> {
void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
ValueBoundsConstraintSet &cstr) const {
auto getGlobalOp = cast<GetGlobalOp>(op);
assert(value == getGlobalOp.getResult() && "invalid value");
auto type = getGlobalOp.getType();
assert(!type.isDynamicDim(dim) && "expected static dim");
cstr.bound(value)[dim] == type.getDimSize(dim);
}
};
struct RankOpInterface
: public ValueBoundsOpInterface::ExternalModel<RankOpInterface, RankOp> {
void populateBoundsForIndexValue(Operation *op, Value value,
ValueBoundsConstraintSet &cstr) const {
auto rankOp = cast<RankOp>(op);
assert(value == rankOp.getResult() && "invalid value");
auto memrefType = llvm::dyn_cast<MemRefType>(rankOp.getMemref().getType());
if (!memrefType)
return;
cstr.bound(value) == memrefType.getRank();
}
};
struct SubViewOpInterface
: public ValueBoundsOpInterface::ExternalModel<SubViewOpInterface,
SubViewOp> {
void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
ValueBoundsConstraintSet &cstr) const {
auto subViewOp = cast<SubViewOp>(op);
assert(value == subViewOp.getResult() && "invalid value");
llvm::SmallBitVector dropped = subViewOp.getDroppedDims();
int64_t ctr = -1;
for (int64_t i = 0, e = subViewOp.getMixedSizes().size(); i < e; ++i) {
// Skip over rank-reduced dimensions.
if (!dropped.test(i))
++ctr;
if (ctr == dim) {
cstr.bound(value)[dim] == subViewOp.getMixedSizes()[i];
return;
}
}
llvm_unreachable("could not find non-rank-reduced dim");
}
};
} // namespace
} // namespace memref
} // namespace mlir
void mlir::memref::registerValueBoundsOpInterfaceExternalModels(
DialectRegistry ®istry) {
registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
memref::AllocOp::attachInterface<memref::AllocOpInterface<memref::AllocOp>>(
*ctx);
memref::AllocaOp::attachInterface<
memref::AllocOpInterface<memref::AllocaOp>>(*ctx);
memref::CastOp::attachInterface<memref::CastOpInterface>(*ctx);
memref::DimOp::attachInterface<memref::DimOpInterface>(*ctx);
memref::GetGlobalOp::attachInterface<memref::GetGlobalOpInterface>(*ctx);
memref::RankOp::attachInterface<memref::RankOpInterface>(*ctx);
memref::SubViewOp::attachInterface<memref::SubViewOpInterface>(*ctx);
});
}
|
