MLIR  17.0.0git
VectorUtils.cpp
Go to the documentation of this file.
1 //===- VectorUtils.cpp - MLIR Utilities for VectorOps ------------------===//
2 //
3 // Part of the MLIR 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 utility methods for working with the Vector dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
14 
23 #include "mlir/IR/Builders.h"
24 #include "mlir/IR/IntegerSet.h"
25 #include "mlir/IR/Operation.h"
26 #include "mlir/IR/TypeUtilities.h"
27 #include "mlir/Support/LLVM.h"
29 
30 #include "llvm/ADT/DenseSet.h"
31 #include "llvm/ADT/SetVector.h"
32 
33 using namespace mlir;
34 
35 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
36 /// the type of `source`.
38  int64_t dim) {
39  if (isa<UnrankedMemRefType, MemRefType>(source.getType()))
40  return b.createOrFold<memref::DimOp>(loc, source, dim);
41  if (isa<UnrankedTensorType, RankedTensorType>(source.getType()))
42  return b.createOrFold<tensor::DimOp>(loc, source, dim);
43  llvm_unreachable("Expected MemRefType or TensorType");
44 }
45 
46 /// Given the n-D transpose pattern 'transp', return true if 'dim0' and 'dim1'
47 /// should be transposed with each other within the context of their 2D
48 /// transposition slice.
49 ///
50 /// Example 1: dim0 = 0, dim1 = 2, transp = [2, 1, 0]
51 /// Return true: dim0 and dim1 are transposed within the context of their 2D
52 /// transposition slice ([1, 0]).
53 ///
54 /// Example 2: dim0 = 0, dim1 = 1, transp = [2, 1, 0]
55 /// Return true: dim0 and dim1 are transposed within the context of their 2D
56 /// transposition slice ([1, 0]). Paradoxically, note how dim1 (1) is *not*
57 /// transposed within the full context of the transposition.
58 ///
59 /// Example 3: dim0 = 0, dim1 = 1, transp = [2, 0, 1]
60 /// Return false: dim0 and dim1 are *not* transposed within the context of
61 /// their 2D transposition slice ([0, 1]). Paradoxically, note how dim0 (0)
62 /// and dim1 (1) are transposed within the full context of the of the
63 /// transposition.
64 static bool areDimsTransposedIn2DSlice(int64_t dim0, int64_t dim1,
65  ArrayRef<int64_t> transp) {
66  // Perform a linear scan along the dimensions of the transposed pattern. If
67  // dim0 is found first, dim0 and dim1 are not transposed within the context of
68  // their 2D slice. Otherwise, 'dim1' is found first and they are transposed.
69  for (int64_t permDim : transp) {
70  if (permDim == dim0)
71  return false;
72  if (permDim == dim1)
73  return true;
74  }
75 
76  llvm_unreachable("Ill-formed transpose pattern");
77 }
78 
80 mlir::vector::isTranspose2DSlice(vector::TransposeOp op) {
81  VectorType srcType = op.getSourceVectorType();
82  SmallVector<int64_t> srcGtOneDims;
83  for (auto [index, size] : llvm::enumerate(srcType.getShape()))
84  if (size > 1)
85  srcGtOneDims.push_back(index);
86 
87  if (srcGtOneDims.size() != 2)
88  return failure();
89 
90  SmallVector<int64_t> transp;
91  for (auto attr : op.getTransp())
92  transp.push_back(cast<IntegerAttr>(attr).getInt());
93 
94  // Check whether the two source vector dimensions that are greater than one
95  // must be transposed with each other so that we can apply one of the 2-D
96  // transpose pattens. Otherwise, these patterns are not applicable.
97  if (!areDimsTransposedIn2DSlice(srcGtOneDims[0], srcGtOneDims[1], transp))
98  return failure();
99 
100  return std::pair<int, int>(srcGtOneDims[0], srcGtOneDims[1]);
101 }
102 
103 /// Constructs a permutation map from memref indices to vector dimension.
104 ///
105 /// The implementation uses the knowledge of the mapping of enclosing loop to
106 /// vector dimension. `enclosingLoopToVectorDim` carries this information as a
107 /// map with:
108 /// - keys representing "vectorized enclosing loops";
109 /// - values representing the corresponding vector dimension.
110 /// The algorithm traverses "vectorized enclosing loops" and extracts the
111 /// at-most-one MemRef index that is invariant along said loop. This index is
112 /// guaranteed to be at most one by construction: otherwise the MemRef is not
113 /// vectorizable.
114 /// If this invariant index is found, it is added to the permutation_map at the
115 /// proper vector dimension.
116 /// If no index is found to be invariant, 0 is added to the permutation_map and
117 /// corresponds to a vector broadcast along that dimension.
118 ///
119 /// Returns an empty AffineMap if `enclosingLoopToVectorDim` is empty,
120 /// signalling that no permutation map can be constructed given
121 /// `enclosingLoopToVectorDim`.
122 ///
123 /// Examples can be found in the documentation of `makePermutationMap`, in the
124 /// header file.
126  ArrayRef<Value> indices,
127  const DenseMap<Operation *, unsigned> &enclosingLoopToVectorDim) {
128  if (enclosingLoopToVectorDim.empty())
129  return AffineMap();
130  MLIRContext *context =
131  enclosingLoopToVectorDim.begin()->getFirst()->getContext();
132  SmallVector<AffineExpr> perm(enclosingLoopToVectorDim.size(),
133  getAffineConstantExpr(0, context));
134 
135  for (auto kvp : enclosingLoopToVectorDim) {
136  assert(kvp.second < perm.size());
137  auto invariants = affine::getInvariantAccesses(
138  cast<affine::AffineForOp>(kvp.first).getInductionVar(), indices);
139  unsigned numIndices = indices.size();
140  unsigned countInvariantIndices = 0;
141  for (unsigned dim = 0; dim < numIndices; ++dim) {
142  if (!invariants.count(indices[dim])) {
143  assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
144  "permutationMap already has an entry along dim");
145  perm[kvp.second] = getAffineDimExpr(dim, context);
146  } else {
147  ++countInvariantIndices;
148  }
149  }
150  assert((countInvariantIndices == numIndices ||
151  countInvariantIndices == numIndices - 1) &&
152  "Vectorization prerequisite violated: at most 1 index may be "
153  "invariant wrt a vectorized loop");
154  (void)countInvariantIndices;
155  }
156  return AffineMap::get(indices.size(), 0, perm, context);
157 }
158 
159 /// Implementation detail that walks up the parents and records the ones with
160 /// the specified type.
161 /// TODO: could also be implemented as a collect parents followed by a
162 /// filter and made available outside this file.
163 template <typename T>
166  auto *current = block->getParentOp();
167  while (current) {
168  if (auto typedParent = dyn_cast<T>(current)) {
169  assert(res.count(current) == 0 && "Already inserted");
170  res.insert(current);
171  }
172  current = current->getParentOp();
173  }
174  return res;
175 }
176 
177 /// Returns the enclosing AffineForOp, from closest to farthest.
179  return getParentsOfType<affine::AffineForOp>(block);
180 }
181 
183  Block *insertPoint, ArrayRef<Value> indices,
184  const DenseMap<Operation *, unsigned> &loopToVectorDim) {
185  DenseMap<Operation *, unsigned> enclosingLoopToVectorDim;
186  auto enclosingLoops = getEnclosingforOps(insertPoint);
187  for (auto *forInst : enclosingLoops) {
188  auto it = loopToVectorDim.find(forInst);
189  if (it != loopToVectorDim.end()) {
190  enclosingLoopToVectorDim.insert(*it);
191  }
192  }
193  return ::makePermutationMap(indices, enclosingLoopToVectorDim);
194 }
195 
197  Operation *op, ArrayRef<Value> indices,
198  const DenseMap<Operation *, unsigned> &loopToVectorDim) {
199  return makePermutationMap(op->getBlock(), indices, loopToVectorDim);
200 }
201 
202 bool matcher::operatesOnSuperVectorsOf(Operation &op,
203  VectorType subVectorType) {
204  // First, extract the vector type and distinguish between:
205  // a. ops that *must* lower a super-vector (i.e. vector.transfer_read,
206  // vector.transfer_write); and
207  // b. ops that *may* lower a super-vector (all other ops).
208  // The ops that *may* lower a super-vector only do so if the super-vector to
209  // sub-vector ratio exists. The ops that *must* lower a super-vector are
210  // explicitly checked for this property.
211  /// TODO: there should be a single function for all ops to do this so we
212  /// do not have to special case. Maybe a trait, or just a method, unclear atm.
213  bool mustDivide = false;
214  (void)mustDivide;
215  VectorType superVectorType;
216  if (auto transfer = dyn_cast<VectorTransferOpInterface>(op)) {
217  superVectorType = transfer.getVectorType();
218  mustDivide = true;
219  } else if (op.getNumResults() == 0) {
220  if (!isa<func::ReturnOp>(op)) {
221  op.emitError("NYI: assuming only return operations can have 0 "
222  " results at this point");
223  }
224  return false;
225  } else if (op.getNumResults() == 1) {
226  if (auto v = dyn_cast<VectorType>(op.getResult(0).getType())) {
227  superVectorType = v;
228  } else {
229  // Not a vector type.
230  return false;
231  }
232  } else {
233  // Not a vector.transfer and has more than 1 result, fail hard for now to
234  // wake us up when something changes.
235  op.emitError("NYI: operation has more than 1 result");
236  return false;
237  }
238 
239  // Get the ratio.
240  auto ratio =
241  computeShapeRatio(superVectorType.getShape(), subVectorType.getShape());
242 
243  // Sanity check.
244  assert((ratio || !mustDivide) &&
245  "vector.transfer operation in which super-vector size is not an"
246  " integer multiple of sub-vector size");
247 
248  // This catches cases that are not strictly necessary to have multiplicity but
249  // still aren't divisible by the sub-vector shape.
250  // This could be useful information if we wanted to reshape at the level of
251  // the vector type (but we would have to look at the compute and distinguish
252  // between parallel, reduction and possibly other cases.
253  return ratio.has_value();
254 }
static bool areDimsTransposedIn2DSlice(int64_t dim0, int64_t dim1, ArrayRef< int64_t > transp)
Given the n-D transpose pattern 'transp', return true if 'dim0' and 'dim1' should be transposed with ...
Definition: VectorUtils.cpp:64
static SetVector< Operation * > getEnclosingforOps(Block *block)
Returns the enclosing AffineForOp, from closest to farthest.
static AffineMap makePermutationMap(ArrayRef< Value > indices, const DenseMap< Operation *, unsigned > &enclosingLoopToVectorDim)
Constructs a permutation map from memref indices to vector dimension.
static SetVector< Operation * > getParentsOfType(Block *block)
Implementation detail that walks up the parents and records the ones with the specified type.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:44
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
Block represents an ordered list of Operations.
Definition: Block.h:30
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
Definition: Block.cpp:30
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
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class helps build Operations.
Definition: Builders.h:202
void createOrFold(SmallVectorImpl< Value > &results, Location location, Args &&...args)
Create an operation of specific op type at the current insertion point, and immediately try to fold i...
Definition: Builders.h:501
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:402
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:266
Block * getBlock()
Returns the operation block that contains this operation.
Definition: Operation.h:213
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:399
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:93
Type getType() const
Return the type of this value.
Definition: Value.h:122
DenseSet< Value, DenseMapInfo< Value > > getInvariantAccesses(Value iv, ArrayRef< Value > indices)
Given an induction variable iv of type AffineForOp and indices of type IndexType, returns the set of ...
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:262
FailureOr< std::pair< int, int > > isTranspose2DSlice(vector::TransposeOp op)
Returns two dims that are greater than one if the transposition is applied on a 2D slice.
Definition: VectorUtils.cpp:80
Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim)
Helper function that creates a memref::DimOp or tensor::DimOp depending on the type of source.
Definition: VectorUtils.cpp:37
This header declares functions that assit transformations in the MemRef dialect.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:527
std::optional< SmallVector< int64_t > > computeShapeRatio(ArrayRef< int64_t > shape, ArrayRef< int64_t > subShape)
Return the multi-dimensional integral ratio of subShape to the trailing dimensions of shape.
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Definition: AffineExpr.cpp:502