MLIR  16.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 
22 #include "mlir/IR/Builders.h"
23 #include "mlir/IR/IntegerSet.h"
24 #include "mlir/IR/Operation.h"
25 #include "mlir/IR/TypeUtilities.h"
26 #include "mlir/Support/LLVM.h"
28 #include <numeric>
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 (source.getType().isa<UnrankedMemRefType, MemRefType>())
40  return b.createOrFold<memref::DimOp>(loc, source, dim);
41  if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
42  return b.createOrFold<tensor::DimOp>(loc, source, dim);
43  llvm_unreachable("Expected MemRefType or TensorType");
44 }
45 
46 /// Return the number of elements of basis, `0` if empty.
48  if (basis.empty())
49  return 0;
50  return std::accumulate(basis.begin(), basis.end(), 1,
51  std::multiplies<int64_t>());
52 }
53 
55  ArrayRef<int64_t> sizes) {
56  int64_t rank = shape.size();
57  // Compute the count for each dimension.
58  SmallVector<int64_t, 4> sliceDimCounts(rank);
59  for (int64_t r = 0; r < rank; ++r)
60  sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
61  // Use that to compute the slice stride for each dimension.
62  SmallVector<int64_t, 4> sliceStrides(rank);
63  sliceStrides[rank - 1] = 1;
64  for (int64_t r = rank - 2; r >= 0; --r)
65  sliceStrides[r] = sliceStrides[r + 1] * sliceDimCounts[r + 1];
66  return sliceStrides;
67 }
68 
70  ArrayRef<int64_t> sizes, ArrayRef<int64_t> vectorOffsets) {
72  for (auto it : llvm::zip(vectorOffsets, sizes))
73  result.push_back(std::get<0>(it) * std::get<1>(it));
74  return result;
75 }
76 
78  ArrayRef<int64_t> subShape) {
79  if (superShape.size() < subShape.size()) {
81  }
82 
83  // Starting from the end, compute the integer divisors.
84  std::vector<int64_t> result;
85  result.reserve(superShape.size());
86  for (auto [superSize, subSize] :
87  llvm::zip(llvm::reverse(superShape), llvm::reverse(subShape))) {
88  assert(superSize > 0 && "superSize must be > 0");
89  assert(subSize > 0 && "subSize must be > 0");
90 
91  // If integral division does not occur, return and let the caller decide.
92  if (superSize % subSize != 0)
93  return None;
94  result.push_back(superSize / subSize);
95  }
96 
97  // At this point we computed the ratio (in reverse) for the common
98  // size. Fill with the remaining entries from the super-vector shape (still in
99  // reverse).
100  int commonSize = subShape.size();
101  std::copy(superShape.rbegin() + commonSize, superShape.rend(),
102  std::back_inserter(result));
103 
104  assert(result.size() == superShape.size() &&
105  "super to sub shape ratio is not of the same size as the super rank");
106 
107  // Reverse again to get it back in the proper order and return.
108  return SmallVector<int64_t, 4>{result.rbegin(), result.rend()};
109 }
110 
112  VectorType subVectorType) {
113  assert(superVectorType.getElementType() == subVectorType.getElementType() &&
114  "vector types must be of the same elemental type");
115  return shapeRatio(superVectorType.getShape(), subVectorType.getShape());
116 }
117 
118 /// Constructs a permutation map from memref indices to vector dimension.
119 ///
120 /// The implementation uses the knowledge of the mapping of enclosing loop to
121 /// vector dimension. `enclosingLoopToVectorDim` carries this information as a
122 /// map with:
123 /// - keys representing "vectorized enclosing loops";
124 /// - values representing the corresponding vector dimension.
125 /// The algorithm traverses "vectorized enclosing loops" and extracts the
126 /// at-most-one MemRef index that is invariant along said loop. This index is
127 /// guaranteed to be at most one by construction: otherwise the MemRef is not
128 /// vectorizable.
129 /// If this invariant index is found, it is added to the permutation_map at the
130 /// proper vector dimension.
131 /// If no index is found to be invariant, 0 is added to the permutation_map and
132 /// corresponds to a vector broadcast along that dimension.
133 ///
134 /// Returns an empty AffineMap if `enclosingLoopToVectorDim` is empty,
135 /// signalling that no permutation map can be constructed given
136 /// `enclosingLoopToVectorDim`.
137 ///
138 /// Examples can be found in the documentation of `makePermutationMap`, in the
139 /// header file.
141  ArrayRef<Value> indices,
142  const DenseMap<Operation *, unsigned> &enclosingLoopToVectorDim) {
143  if (enclosingLoopToVectorDim.empty())
144  return AffineMap();
145  MLIRContext *context =
146  enclosingLoopToVectorDim.begin()->getFirst()->getContext();
147  SmallVector<AffineExpr, 4> perm(enclosingLoopToVectorDim.size(),
148  getAffineConstantExpr(0, context));
149 
150  for (auto kvp : enclosingLoopToVectorDim) {
151  assert(kvp.second < perm.size());
152  auto invariants = getInvariantAccesses(
153  cast<AffineForOp>(kvp.first).getInductionVar(), indices);
154  unsigned numIndices = indices.size();
155  unsigned countInvariantIndices = 0;
156  for (unsigned dim = 0; dim < numIndices; ++dim) {
157  if (!invariants.count(indices[dim])) {
158  assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
159  "permutationMap already has an entry along dim");
160  perm[kvp.second] = getAffineDimExpr(dim, context);
161  } else {
162  ++countInvariantIndices;
163  }
164  }
165  assert((countInvariantIndices == numIndices ||
166  countInvariantIndices == numIndices - 1) &&
167  "Vectorization prerequisite violated: at most 1 index may be "
168  "invariant wrt a vectorized loop");
169  }
170  return AffineMap::get(indices.size(), 0, perm, context);
171 }
172 
173 /// Implementation detail that walks up the parents and records the ones with
174 /// the specified type.
175 /// TODO: could also be implemented as a collect parents followed by a
176 /// filter and made available outside this file.
177 template <typename T>
180  auto *current = block->getParentOp();
181  while (current) {
182  if (auto typedParent = dyn_cast<T>(current)) {
183  assert(res.count(current) == 0 && "Already inserted");
184  res.insert(current);
185  }
186  current = current->getParentOp();
187  }
188  return res;
189 }
190 
191 /// Returns the enclosing AffineForOp, from closest to farthest.
193  return getParentsOfType<AffineForOp>(block);
194 }
195 
197  Block *insertPoint, ArrayRef<Value> indices,
198  const DenseMap<Operation *, unsigned> &loopToVectorDim) {
199  DenseMap<Operation *, unsigned> enclosingLoopToVectorDim;
200  auto enclosingLoops = getEnclosingforOps(insertPoint);
201  for (auto *forInst : enclosingLoops) {
202  auto it = loopToVectorDim.find(forInst);
203  if (it != loopToVectorDim.end()) {
204  enclosingLoopToVectorDim.insert(*it);
205  }
206  }
207  return ::makePermutationMap(indices, enclosingLoopToVectorDim);
208 }
209 
211  Operation *op, ArrayRef<Value> indices,
212  const DenseMap<Operation *, unsigned> &loopToVectorDim) {
213  return makePermutationMap(op->getBlock(), indices, loopToVectorDim);
214 }
215 
216 bool matcher::operatesOnSuperVectorsOf(Operation &op,
217  VectorType subVectorType) {
218  // First, extract the vector type and distinguish between:
219  // a. ops that *must* lower a super-vector (i.e. vector.transfer_read,
220  // vector.transfer_write); and
221  // b. ops that *may* lower a super-vector (all other ops).
222  // The ops that *may* lower a super-vector only do so if the super-vector to
223  // sub-vector ratio exists. The ops that *must* lower a super-vector are
224  // explicitly checked for this property.
225  /// TODO: there should be a single function for all ops to do this so we
226  /// do not have to special case. Maybe a trait, or just a method, unclear atm.
227  bool mustDivide = false;
228  (void)mustDivide;
229  VectorType superVectorType;
230  if (auto transfer = dyn_cast<VectorTransferOpInterface>(op)) {
231  superVectorType = transfer.getVectorType();
232  mustDivide = true;
233  } else if (op.getNumResults() == 0) {
234  if (!isa<func::ReturnOp>(op)) {
235  op.emitError("NYI: assuming only return operations can have 0 "
236  " results at this point");
237  }
238  return false;
239  } else if (op.getNumResults() == 1) {
240  if (auto v = op.getResult(0).getType().dyn_cast<VectorType>()) {
241  superVectorType = v;
242  } else {
243  // Not a vector type.
244  return false;
245  }
246  } else {
247  // Not a vector.transfer and has more than 1 result, fail hard for now to
248  // wake us up when something changes.
249  op.emitError("NYI: operation has more than 1 result");
250  return false;
251  }
252 
253  // Get the ratio.
254  auto ratio = shapeRatio(superVectorType, subVectorType);
255 
256  // Sanity check.
257  assert((ratio || !mustDivide) &&
258  "vector.transfer operation in which super-vector size is not an"
259  " integer multiple of sub-vector size");
260 
261  // This catches cases that are not strictly necessary to have multiplicity but
262  // still aren't divisible by the sub-vector shape.
263  // This could be useful information if we wanted to reshape at the level of
264  // the vector type (but we would have to look at the compute and distinguish
265  // between parallel, reduction and possibly other cases.
266  return ratio.has_value();
267 }
Include the generated interface declarations.
static SetVector< Operation * > getEnclosingforOps(Block *block)
Returns the enclosing AffineForOp, from closest to farthest.
SmallVector< int64_t, 4 > computeStrides(ArrayRef< int64_t > shape, ArrayRef< int64_t > sizes)
Given the shape and sizes of a vector, returns the corresponding strides for each dimension...
Definition: VectorUtils.cpp:54
Optional< SmallVector< int64_t, 4 > > shapeRatio(ArrayRef< int64_t > superShape, ArrayRef< int64_t > subShape)
Computes and returns the multi-dimensional ratio of superShape to subShape.
Definition: VectorUtils.cpp:77
SmallVector< int64_t, 4 > computeElementOffsetsFromVectorSliceOffsets(ArrayRef< int64_t > sizes, ArrayRef< int64_t > vectorOffsets)
Given the target sizes of a vector, together with vector-space offsets, returns the element-space off...
Definition: VectorUtils.cpp:69
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:466
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:28
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
Definition: Block.cpp:30
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:514
Block represents an ordered list of Operations.
Definition: Block.h:29
static void copy(Location loc, Value dst, Value src, Value size, OpBuilder &builder)
Copies the given number of bytes from src to dst pointers.
static unsigned perm(const SparseTensorEncodingAttr &enc, unsigned d)
Helper method to apply dimension ordering permutation.
Block * getBlock()
Returns the operation block that contains this operation.
Definition: Operation.h:144
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:48
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
U dyn_cast() const
Definition: Types.h:270
int64_t ceilDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR&#39;s ceildiv operation on constants.
Definition: MathExtras.h:23
OpResult getResult(unsigned idx)
Get the &#39;idx&#39;th result of this operation.
Definition: Operation.h:324
static AffineMap makePermutationMap(ArrayRef< Value > indices, const DenseMap< Operation *, unsigned > &enclosingLoopToVectorDim)
Constructs a permutation map from memref indices to vector dimension.
int64_t computeMaxLinearIndex(ArrayRef< int64_t > basis)
Return the number of elements of basis, 0 if empty.
Definition: VectorUtils.cpp:47
A multi-dimensional affine map Affine map&#39;s are immutable like Type&#39;s, and they are uniqued...
Definition: AffineMap.h:42
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Definition: AffineExpr.cpp:489
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 class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:85
Type getType() const
Return the type of this value.
Definition: Value.h:118
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:55
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:321
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 ...
static SetVector< Operation * > getParentsOfType(Block *block)
Implementation detail that walks up the parents and records the ones with the specified type...
bool isa() const
Definition: Types.h:254
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:221
This class helps build Operations.
Definition: Builders.h:192