MLIR  19.0.0git
Utils.cpp
Go to the documentation of this file.
1 //===- Utils.cpp - Utilities to support the Tensor dialect ----------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements utilities for the Tensor dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
14 
20 
21 using namespace mlir;
22 using namespace mlir::tensor;
23 
24 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
25  Value pad, bool nofold, Location loc,
26  OpBuilder &b) {
27  SmallVector<OpFoldResult> low(type.getRank(), b.getIndexAttr(0));
28  SmallVector<OpFoldResult> high(type.getRank(), b.getIndexAttr(0));
29  for (const auto &en : enumerate(type.getShape())) {
30  // Pad only the static dimensions of the result tensor type.
31  if (ShapedType::isDynamic(en.value()))
32  continue;
33  // Compute the padding width.
34  AffineExpr d0;
35  bindDims(b.getContext(), d0);
36  OpFoldResult sz = tensor::getMixedSize(b, loc, source, en.index());
37  high[en.index()] =
38  affine::makeComposedFoldedAffineApply(b, loc, en.value() - d0, {sz});
39  }
40  return b.create<PadOp>(loc, type, source, low, high, pad, nofold);
41 }
42 
44  Location loc,
45  Value rankedTensor) {
46  auto tensorTy = cast<RankedTensorType>(rankedTensor.getType());
47  SmallVector<Value> dynamicDims;
48  for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
49  if (en.value() == ShapedType::kDynamic)
50  dynamicDims.push_back(
51  b.create<tensor::DimOp>(loc, rankedTensor, en.index()));
52  }
53  return dynamicDims;
54 }
55 
56 FailureOr<RankedTensorType>
57 mlir::tensor::computeTransposedType(RankedTensorType rankedTensorType,
58  ArrayRef<int64_t> transposeVector) {
59  if (transposeVector.empty())
60  return rankedTensorType;
61 
62  if (!isPermutationVector(transposeVector) ||
63  transposeVector.size() != static_cast<size_t>(rankedTensorType.getRank()))
64  return failure();
65 
66  SmallVector<int64_t> transposedShape(rankedTensorType.getShape().begin(),
67  rankedTensorType.getShape().end());
68  applyPermutationToVector(transposedShape, transposeVector);
69 
70  using RTTBuilder = RankedTensorType::Builder;
71  RankedTensorType transposedTensorType =
72  RTTBuilder(rankedTensorType).setShape(transposedShape);
73  return transposedTensorType;
74 }
75 
76 /// The permutation can be obtained from two permutations:
77 /// a) Compute the permutation vector to move the last `numPackedDims` into
78 /// the `innerPosDims` of a shape of rank `rank`.
79 /// b) Compute the permutation vector to move outer dims if the
80 /// `outerPerm` parameter is not empty.
81 /// Apply (b) permutation on (a) permutation to get the final permutation.
83 computePackUnPackPerm(int64_t rank, ArrayRef<int64_t> &innerDimsPos,
84  ArrayRef<int64_t> &outerPerm,
85  PackingMetadata &packingMetadata) {
86  int64_t numPackedDims = innerDimsPos.size();
87  auto lastDims =
88  llvm::to_vector(llvm::seq<int64_t>(rank - numPackedDims, rank));
89  packingMetadata = computePackingMetadata(rank, innerDimsPos);
90  SmallVector<int64_t> innerPositionsPerm =
91  computePermutationVector(rank, lastDims, packingMetadata.insertPositions);
92 
93  SmallVector<int64_t> outerPos = packingMetadata.outerPositions;
94  if (!outerPerm.empty())
95  applyPermutationToVector(outerPos, outerPerm);
96  SmallVector<int64_t> outerPositionPerm =
97  computePermutationVector(rank, packingMetadata.outerPositions, outerPos);
98 
99  SmallVector<int64_t> packInverseDestPermutation = innerPositionsPerm;
100  applyPermutationToVector(packInverseDestPermutation, outerPositionPerm);
101  return packInverseDestPermutation;
102 }
103 
105 
106  PackingMetadata pMetadata;
107  int64_t packedRank = packOp.getDestType().getRank();
108  ArrayRef<int64_t> innerDimPos = packOp.getInnerDimsPos();
109  ArrayRef<int64_t> outerPerm = packOp.getOuterDimsPerm();
110  SmallVector<int64_t> packInvDestPerm =
111  computePackUnPackPerm(packedRank, innerDimPos, outerPerm, pMetadata);
112  return packInvDestPerm;
113 }
114 
116  PackingMetadata metadata;
117  return mlir::tensor::getUnPackInverseSrcPerm(unpackOp, metadata);
118 }
119 
121 mlir::tensor::getUnPackInverseSrcPerm(UnPackOp unpackOp,
122  PackingMetadata &metadata) {
123  int64_t unpackRank = unpackOp.getSourceType().getRank();
124  ArrayRef<int64_t> innerDimPos = unpackOp.getInnerDimsPos();
125  ArrayRef<int64_t> outerPerm = unpackOp.getOuterDimsPerm();
126  SmallVector<int64_t> unpackInvSrcPerm =
127  computePackUnPackPerm(unpackRank, innerDimPos, outerPerm, metadata);
128  return unpackInvSrcPerm;
129 }
130 
132  llvm::SmallBitVector droppedDims = op.getDroppedDims();
133  int64_t srcDim = 0;
134  RankedTensorType resultType = op.getDestType();
135  // Source dims and destination dims (apart from dropped dims) must have the
136  // same size.
137  for (int64_t resultDim = 0; resultDim < resultType.getRank(); ++resultDim) {
138  if (droppedDims.test(resultDim)) {
139  // InsertSlice may expand unit dimensions that result from inserting a
140  // size-1 slice into a non-size-1 result dimension.
141  if (resultType.getDimSize(resultDim) != 1)
142  return false;
143  continue;
144  }
145  FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(
146  {op.getSource(), srcDim}, {op.getResult(), resultDim});
147  if (failed(equalDimSize) || !*equalDimSize)
148  return false;
149  ++srcDim;
150  }
151 
152  return true;
153 }
154 
155 bool mlir::tensor::isCastLikeExtractSliceOp(ExtractSliceOp op) {
156  llvm::SmallBitVector droppedDims = op.getDroppedDims();
157  int64_t resultDim = 0;
158  // Source dims and result dims (apart from dropped dims) must have the same
159  // size.
160  RankedTensorType sourceType = op.getSourceType();
161  for (int64_t dim = 0, e = sourceType.getRank(); dim < e; ++dim) {
162  if (droppedDims.test(dim)) {
163  // ExtractSlice may drop unit dimensions that result from taking a size-1
164  // slice from a non-size-1 source dimension.
165  if (sourceType.getDimSize(dim) != 1)
166  return false;
167  continue;
168  }
169  FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(
170  {op.getSource(), dim}, {op.getResult(), resultDim});
171  if (failed(equalDimSize) || !*equalDimSize)
172  return false;
173  ++resultDim;
174  }
175 
176  return true;
177 }
static SmallVector< int64_t > computePackUnPackPerm(int64_t rank, ArrayRef< int64_t > &innerDimsPos, ArrayRef< int64_t > &outerPerm, PackingMetadata &packingMetadata)
The permutation can be obtained from two permutations: a) Compute the permutation vector to move the ...
Definition: Utils.cpp:83
Base type for affine expression.
Definition: AffineExpr.h:68
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:124
MLIRContext * getContext() const
Definition: Builders.h:55
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
This class helps build Operations.
Definition: Builders.h:209
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:464
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:402
This is a builder type that keeps local references to arguments.
Definition: BuiltinTypes.h:254
Builder & setShape(ArrayRef< int64_t > newShape)
Definition: BuiltinTypes.h:265
static FailureOr< bool > areEqual(const Variable &var1, const Variable &var2)
Compute whether the given variables are equal.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:129
OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineApplyOp that applies map to operands after composing the map with the maps of any...
Definition: AffineOps.cpp:1192
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
SmallVector< int64_t > getUnPackInverseSrcPerm(tensor::UnPackOp unpackOp)
Shell function to compute the Source Permutation of unPackOp.
SmallVector< Value > createDynamicDimValues(OpBuilder &b, Location loc, Value rankedTensor)
Definition: Utils.cpp:43
bool isCastLikeInsertSliceOp(InsertSliceOp op)
A tensor.insert_slice is a cast-like operation if it merely rank-extends the source tensor or inserts...
Definition: Utils.cpp:131
bool isCastLikeExtractSliceOp(ExtractSliceOp op)
A tensor.extract_slice is a cast-like operation if it merely rank-reduces unit dimensions of the sour...
Definition: Utils.cpp:155
OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value, int64_t dim)
Return the dimension of the given tensor value.
Definition: TensorOps.cpp:55
SmallVector< int64_t > getPackInverseDestPerm(tensor::PackOp packOp)
Shell function to compute the Destination Permutation of PackOp This function uses the helper functio...
PadOp createPadHighOp(RankedTensorType type, Value source, Value pad, bool nofold, Location loc, OpBuilder &builder)
Definition: Utils.cpp:24
FailureOr< RankedTensorType > computeTransposedType(RankedTensorType rankedTensorType, ArrayRef< int64_t > transposeVector)
Returns the transposed rankedTensorType if transposeVector is non-empty.
Definition: Utils.cpp:57
Include the generated interface declarations.
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
Definition: AffineExpr.h:348
SmallVector< int64_t > computePermutationVector(int64_t permSize, ArrayRef< int64_t > positions, ArrayRef< int64_t > desiredPositions)
Return a permutation vector of size permSize that would result in moving positions into desiredPositi...
void applyPermutationToVector(SmallVector< T, N > &inVec, ArrayRef< int64_t > permutation)
Apply the permutation defined by permutation to inVec.
bool isPermutationVector(ArrayRef< int64_t > interchange)
Method to check if an interchange vector is a permutation.