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 
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 /// The permutation can be obtained from two permutations:
76 /// a) Compute the permutation vector to move the last `numPackedDims` into
77 /// the `innerPosDims` of a shape of rank `rank`.
78 /// b) Compute the permutation vector to move outer dims if the
79 /// `outerPerm` parameter is not empty.
80 /// Apply (b) permutation on (a) permutation to get the final permutation.
82 computePackUnPackPerm(int64_t rank, ArrayRef<int64_t> &innerDimsPos,
83  ArrayRef<int64_t> &outerPerm,
84  PackingMetadata &packingMetadata) {
85  int64_t numPackedDims = innerDimsPos.size();
86  auto lastDims =
87  llvm::to_vector(llvm::seq<int64_t>(rank - numPackedDims, rank));
88  packingMetadata = computePackingMetadata(rank, innerDimsPos);
89  SmallVector<int64_t> innerPositionsPerm =
90  computePermutationVector(rank, lastDims, packingMetadata.insertPositions);
91 
92  SmallVector<int64_t> outerPos = packingMetadata.outerPositions;
93  if (!outerPerm.empty())
94  applyPermutationToVector(outerPos, outerPerm);
95  SmallVector<int64_t> outerPositionPerm =
96  computePermutationVector(rank, packingMetadata.outerPositions, outerPos);
97 
98  SmallVector<int64_t> packInverseDestPermutation = innerPositionsPerm;
99  applyPermutationToVector(packInverseDestPermutation, outerPositionPerm);
100  return packInverseDestPermutation;
101 }
102 
103 /// Shell function to compute the Destination Permutation of PackOp
104 /// This function uses the helper function `computePackUnPackPerm` to get
105 /// the permutation vector. Only major difference between UnPack and Pack is
106 /// that packOp uses destination rank whereas unpack Uses source rank.
108 
109  PackingMetadata pMetadata;
110  int64_t packedRank = packOp.getDestType().getRank();
111  ArrayRef<int64_t> innerDimPos = packOp.getInnerDimsPos();
112  ArrayRef<int64_t> outerPerm = packOp.getOuterDimsPerm();
113  SmallVector<int64_t> packInvDestPerm =
114  computePackUnPackPerm(packedRank, innerDimPos, outerPerm, pMetadata);
115  return packInvDestPerm;
116 }
117 
118 /// Shell function to compute the Source Permutation of unPackOp.
119 /// This function, like the getPackInverseDestPerm uses the helper function
120 /// computePackUnPackPerm` to get the permutation vector.
121 /// Only major difference between UnPack and Pack is that packOp uses
122 /// destination rank whereas unpack Uses source rank.
124  PackingMetadata metadata;
125  return mlir::tensor::getUnPackInverseSrcPerm(unpackOp, metadata);
126 }
127 
128 /// Shell function to compute the Source rank permutation for unpackOp
129 /// Unpack requires some packing metadata data information, so created
130 /// another function where this value is passed by reference.
132 mlir::tensor::getUnPackInverseSrcPerm(UnPackOp unpackOp,
133  PackingMetadata &metadata) {
134  int64_t unpackRank = unpackOp.getSourceType().getRank();
135  ArrayRef<int64_t> innerDimPos = unpackOp.getInnerDimsPos();
136  ArrayRef<int64_t> outerPerm = unpackOp.getOuterDimsPerm();
137  SmallVector<int64_t> unpackInvSrcPerm =
138  computePackUnPackPerm(unpackRank, innerDimPos, outerPerm, metadata);
139  return unpackInvSrcPerm;
140 }
141 
143  llvm::SmallBitVector droppedDims = op.getDroppedDims();
144  int64_t srcDim = 0;
145  // Source dims and destination dims (apart from dropped dims) must have the
146  // same size.
147  for (int64_t resultDim = 0; resultDim < op.getDestType().getRank();
148  ++resultDim) {
149  if (droppedDims.test(resultDim)) {
150  continue;
151  }
153  op.getSource(), op.getResult(), srcDim, resultDim);
154  if (failed(equalDimSize) || !*equalDimSize)
155  return false;
156  ++srcDim;
157  }
158 
159  return true;
160 }
161 
162 bool mlir::tensor::isCastLikeExtractSliceOp(ExtractSliceOp op) {
163  llvm::SmallBitVector droppedDims = op.getDroppedDims();
164  int64_t resultDim = 0;
165  // Source dims and result dims (apart from dropped dims) must have the same
166  // size.
167  RankedTensorType sourceType = op.getSourceType();
168  for (int64_t dim = 0, e = sourceType.getRank(); dim < e; ++dim) {
169  if (droppedDims.test(dim)) {
170  // ExtractSlice may drop unit dimensions that result from taking a size-1
171  // slice from a non-size-1 source dimension.
172  if (sourceType.getDimSize(dim) != 1)
173  return false;
174  continue;
175  }
177  op.getSource(), op.getResult(), dim, resultDim);
178  if (failed(equalDimSize) || !*equalDimSize)
179  return false;
180  ++resultDim;
181  }
182 
183  return true;
184 }
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:82
Base type for affine expression.
Definition: AffineExpr.h:69
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:124
MLIRContext * getContext() const
Definition: Builders.h:55
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
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:266
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:249
Builder & setShape(ArrayRef< int64_t > newShape)
Definition: BuiltinTypes.h:260
static FailureOr< bool > areEqual(Value value1, Value value2, std::optional< int64_t > dim1=std::nullopt, std::optional< int64_t > dim2=std::nullopt)
Compute whether the given values/dimensions 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:125
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:1185
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
SmallVector< int64_t > getUnPackInverseSrcPerm(tensor::UnPackOp 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:142
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:162
OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value, int64_t dim)
Return the dimension of the given tensor value.
Definition: TensorOps.cpp:51
SmallVector< int64_t > getPackInverseDestPerm(tensor::PackOp packOp)
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.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
Definition: AffineExpr.h:349
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.
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72