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
|
//===- Utils.cpp - Utilities to support the Tensor dialect ----------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements utilities for the Tensor dialect.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Tensor/Utils/Utils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
using namespace mlir;
using namespace mlir::tensor;
PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
Value pad, bool nofold, Location loc,
OpBuilder &b) {
SmallVector<OpFoldResult> low(type.getRank(), b.getIndexAttr(0));
SmallVector<OpFoldResult> high(type.getRank(), b.getIndexAttr(0));
for (const auto &en : enumerate(type.getShape())) {
// Pad only the static dimensions of the result tensor type.
if (ShapedType::isDynamic(en.value()))
continue;
// Compute the padding width.
AffineExpr d0;
bindDims(b.getContext(), d0);
OpFoldResult sz = tensor::getMixedSize(b, loc, source, en.index());
high[en.index()] =
affine::makeComposedFoldedAffineApply(b, loc, en.value() - d0, {sz});
}
return b.create<PadOp>(loc, type, source, low, high, pad, nofold);
}
SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
Location loc,
Value rankedTensor) {
auto tensorTy = cast<RankedTensorType>(rankedTensor.getType());
SmallVector<Value> dynamicDims;
for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
if (en.value() == ShapedType::kDynamic)
dynamicDims.push_back(
b.create<tensor::DimOp>(loc, rankedTensor, en.index()));
}
return dynamicDims;
}
FailureOr<RankedTensorType>
mlir::tensor::computeTransposedType(RankedTensorType rankedTensorType,
ArrayRef<int64_t> transposeVector) {
if (transposeVector.empty())
return rankedTensorType;
if (!isPermutationVector(transposeVector) ||
transposeVector.size() != static_cast<size_t>(rankedTensorType.getRank()))
return failure();
SmallVector<int64_t> transposedShape(rankedTensorType.getShape().begin(),
rankedTensorType.getShape().end());
applyPermutationToVector(transposedShape, transposeVector);
using RTTBuilder = RankedTensorType::Builder;
RankedTensorType transposedTensorType =
RTTBuilder(rankedTensorType).setShape(transposedShape);
return transposedTensorType;
}
bool mlir::tensor::isCastLikeInsertSliceOp(InsertSliceOp op) {
llvm::SmallBitVector droppedDims = op.getDroppedDims();
int64_t srcDim = 0;
// Source dims and destination dims (apart from dropped dims) must have the
// same size.
for (int64_t resultDim = 0; resultDim < op.getDestType().getRank();
++resultDim) {
if (droppedDims.test(resultDim)) {
continue;
}
FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(
op.getSource(), op.getResult(), srcDim, resultDim);
if (failed(equalDimSize) || !*equalDimSize)
return false;
++srcDim;
}
return true;
}
bool mlir::tensor::isCastLikeExtractSliceOp(ExtractSliceOp op) {
llvm::SmallBitVector droppedDims = op.getDroppedDims();
int64_t resultDim = 0;
// Source dims and result dims (apart from dropped dims) must have the same
// size.
for (int64_t dim = 0; dim < op.getSourceType().getRank(); ++dim) {
if (droppedDims.test(dim)) {
continue;
}
FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(
op.getSource(), op.getResult(), dim, resultDim);
if (failed(equalDimSize) || !*equalDimSize)
return false;
++resultDim;
}
return true;
}
|
