MLIR  14.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 
21 #include "mlir/IR/Builders.h"
22 #include "mlir/IR/IntegerSet.h"
23 #include "mlir/IR/Operation.h"
24 #include "mlir/Support/LLVM.h"
26 #include <numeric>
27 
28 #include "llvm/ADT/DenseSet.h"
29 #include "llvm/ADT/SetVector.h"
30 
31 using namespace mlir;
32 
33 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
34 /// the type of `source`.
36  int64_t dim) {
37  if (source.getType().isa<UnrankedMemRefType, MemRefType>())
38  return b.createOrFold<memref::DimOp>(loc, source, dim);
39  if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
40  return b.createOrFold<tensor::DimOp>(loc, source, dim);
41  llvm_unreachable("Expected MemRefType or TensorType");
42 }
43 
44 /// Return the number of elements of basis, `0` if empty.
46  if (basis.empty())
47  return 0;
48  return std::accumulate(basis.begin(), basis.end(), 1,
49  std::multiplies<int64_t>());
50 }
51 
53  ArrayRef<int64_t> sizes) {
54  int64_t rank = shape.size();
55  // Compute the count for each dimension.
56  SmallVector<int64_t, 4> sliceDimCounts(rank);
57  for (int64_t r = 0; r < rank; ++r)
58  sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
59  // Use that to compute the slice stride for each dimension.
60  SmallVector<int64_t, 4> sliceStrides(rank);
61  sliceStrides[rank - 1] = 1;
62  for (int64_t r = rank - 2; r >= 0; --r)
63  sliceStrides[r] = sliceStrides[r + 1] * sliceDimCounts[r + 1];
64  return sliceStrides;
65 }
66 
68  assert(offsets.size() == basis.size());
69  int64_t linearIndex = 0;
70  for (unsigned idx = 0, e = basis.size(); idx < e; ++idx)
71  linearIndex += offsets[idx] * basis[idx];
72  return linearIndex;
73 }
74 
76  int64_t index) {
77  int64_t rank = sliceStrides.size();
78  SmallVector<int64_t, 4> vectorOffsets(rank);
79  for (int64_t r = 0; r < rank; ++r) {
80  assert(sliceStrides[r] > 0);
81  vectorOffsets[r] = index / sliceStrides[r];
82  index %= sliceStrides[r];
83  }
84  return vectorOffsets;
85 }
86 
88  ArrayRef<int64_t> sizes, ArrayRef<int64_t> vectorOffsets) {
90  for (auto it : llvm::zip(vectorOffsets, sizes))
91  result.push_back(std::get<0>(it) * std::get<1>(it));
92  return result;
93 }
94 
96  ArrayRef<int64_t> subShape) {
97  if (superShape.size() < subShape.size()) {
99  }
100 
101  // Starting from the end, compute the integer divisors.
102  std::vector<int64_t> result;
103  result.reserve(superShape.size());
104  int64_t superSize = 0, subSize = 0;
105  for (auto it :
106  llvm::zip(llvm::reverse(superShape), llvm::reverse(subShape))) {
107  std::tie(superSize, subSize) = it;
108  assert(superSize > 0 && "superSize must be > 0");
109  assert(subSize > 0 && "subSize must be > 0");
110 
111  // If integral division does not occur, return and let the caller decide.
112  if (superSize % subSize != 0)
113  return None;
114  result.push_back(superSize / subSize);
115  }
116 
117  // At this point we computed the ratio (in reverse) for the common
118  // size. Fill with the remaining entries from the super-vector shape (still in
119  // reverse).
120  int commonSize = subShape.size();
121  std::copy(superShape.rbegin() + commonSize, superShape.rend(),
122  std::back_inserter(result));
123 
124  assert(result.size() == superShape.size() &&
125  "super to sub shape ratio is not of the same size as the super rank");
126 
127  // Reverse again to get it back in the proper order and return.
128  return SmallVector<int64_t, 4>{result.rbegin(), result.rend()};
129 }
130 
132  VectorType subVectorType) {
133  assert(superVectorType.getElementType() == subVectorType.getElementType() &&
134  "vector types must be of the same elemental type");
135  return shapeRatio(superVectorType.getShape(), subVectorType.getShape());
136 }
137 
138 /// Constructs a permutation map from memref indices to vector dimension.
139 ///
140 /// The implementation uses the knowledge of the mapping of enclosing loop to
141 /// vector dimension. `enclosingLoopToVectorDim` carries this information as a
142 /// map with:
143 /// - keys representing "vectorized enclosing loops";
144 /// - values representing the corresponding vector dimension.
145 /// The algorithm traverses "vectorized enclosing loops" and extracts the
146 /// at-most-one MemRef index that is invariant along said loop. This index is
147 /// guaranteed to be at most one by construction: otherwise the MemRef is not
148 /// vectorizable.
149 /// If this invariant index is found, it is added to the permutation_map at the
150 /// proper vector dimension.
151 /// If no index is found to be invariant, 0 is added to the permutation_map and
152 /// corresponds to a vector broadcast along that dimension.
153 ///
154 /// Returns an empty AffineMap if `enclosingLoopToVectorDim` is empty,
155 /// signalling that no permutation map can be constructed given
156 /// `enclosingLoopToVectorDim`.
157 ///
158 /// Examples can be found in the documentation of `makePermutationMap`, in the
159 /// header file.
161  ArrayRef<Value> indices,
162  const DenseMap<Operation *, unsigned> &enclosingLoopToVectorDim) {
163  if (enclosingLoopToVectorDim.empty())
164  return AffineMap();
165  MLIRContext *context =
166  enclosingLoopToVectorDim.begin()->getFirst()->getContext();
167  SmallVector<AffineExpr, 4> perm(enclosingLoopToVectorDim.size(),
168  getAffineConstantExpr(0, context));
169 
170  for (auto kvp : enclosingLoopToVectorDim) {
171  assert(kvp.second < perm.size());
172  auto invariants = getInvariantAccesses(
173  cast<AffineForOp>(kvp.first).getInductionVar(), indices);
174  unsigned numIndices = indices.size();
175  unsigned countInvariantIndices = 0;
176  for (unsigned dim = 0; dim < numIndices; ++dim) {
177  if (!invariants.count(indices[dim])) {
178  assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
179  "permutationMap already has an entry along dim");
180  perm[kvp.second] = getAffineDimExpr(dim, context);
181  } else {
182  ++countInvariantIndices;
183  }
184  }
185  assert((countInvariantIndices == numIndices ||
186  countInvariantIndices == numIndices - 1) &&
187  "Vectorization prerequisite violated: at most 1 index may be "
188  "invariant wrt a vectorized loop");
189  }
190  return AffineMap::get(indices.size(), 0, perm, context);
191 }
192 
193 /// Implementation detail that walks up the parents and records the ones with
194 /// the specified type.
195 /// TODO: could also be implemented as a collect parents followed by a
196 /// filter and made available outside this file.
197 template <typename T>
200  auto *current = block->getParentOp();
201  while (current) {
202  if (auto typedParent = dyn_cast<T>(current)) {
203  assert(res.count(current) == 0 && "Already inserted");
204  res.insert(current);
205  }
206  current = current->getParentOp();
207  }
208  return res;
209 }
210 
211 /// Returns the enclosing AffineForOp, from closest to farthest.
213  return getParentsOfType<AffineForOp>(block);
214 }
215 
217  Block *insertPoint, ArrayRef<Value> indices,
218  const DenseMap<Operation *, unsigned> &loopToVectorDim) {
219  DenseMap<Operation *, unsigned> enclosingLoopToVectorDim;
220  auto enclosingLoops = getEnclosingforOps(insertPoint);
221  for (auto *forInst : enclosingLoops) {
222  auto it = loopToVectorDim.find(forInst);
223  if (it != loopToVectorDim.end()) {
224  enclosingLoopToVectorDim.insert(*it);
225  }
226  }
227  return ::makePermutationMap(indices, enclosingLoopToVectorDim);
228 }
229 
231  Operation *op, ArrayRef<Value> indices,
232  const DenseMap<Operation *, unsigned> &loopToVectorDim) {
233  return makePermutationMap(op->getBlock(), indices, loopToVectorDim);
234 }
235 
236 AffineMap mlir::getTransferMinorIdentityMap(ShapedType shapedType,
237  VectorType vectorType) {
238  int64_t elementVectorRank = 0;
239  VectorType elementVectorType =
240  shapedType.getElementType().dyn_cast<VectorType>();
241  if (elementVectorType)
242  elementVectorRank += elementVectorType.getRank();
243  // 0-d transfers are to/from tensor<t>/memref<t> and vector<1xt>.
244  // TODO: replace once we have 0-d vectors.
245  if (shapedType.getRank() == 0 &&
246  vectorType.getShape() == ArrayRef<int64_t>{1})
247  return AffineMap::get(
248  /*numDims=*/0, /*numSymbols=*/0,
249  getAffineConstantExpr(0, shapedType.getContext()));
251  shapedType.getRank(), vectorType.getRank() - elementVectorRank,
252  shapedType.getContext());
253 }
254 
255 bool matcher::operatesOnSuperVectorsOf(Operation &op,
256  VectorType subVectorType) {
257  // First, extract the vector type and distinguish between:
258  // a. ops that *must* lower a super-vector (i.e. vector.transfer_read,
259  // vector.transfer_write); and
260  // b. ops that *may* lower a super-vector (all other ops).
261  // The ops that *may* lower a super-vector only do so if the super-vector to
262  // sub-vector ratio exists. The ops that *must* lower a super-vector are
263  // explicitly checked for this property.
264  /// TODO: there should be a single function for all ops to do this so we
265  /// do not have to special case. Maybe a trait, or just a method, unclear atm.
266  bool mustDivide = false;
267  (void)mustDivide;
268  VectorType superVectorType;
269  if (auto transfer = dyn_cast<VectorTransferOpInterface>(op)) {
270  superVectorType = transfer.getVectorType();
271  mustDivide = true;
272  } else if (op.getNumResults() == 0) {
273  if (!isa<ReturnOp>(op)) {
274  op.emitError("NYI: assuming only return operations can have 0 "
275  " results at this point");
276  }
277  return false;
278  } else if (op.getNumResults() == 1) {
279  if (auto v = op.getResult(0).getType().dyn_cast<VectorType>()) {
280  superVectorType = v;
281  } else {
282  // Not a vector type.
283  return false;
284  }
285  } else {
286  // Not a vector.transfer and has more than 1 result, fail hard for now to
287  // wake us up when something changes.
288  op.emitError("NYI: operation has more than 1 result");
289  return false;
290  }
291 
292  // Get the ratio.
293  auto ratio = shapeRatio(superVectorType, subVectorType);
294 
295  // Sanity check.
296  assert((ratio.hasValue() || !mustDivide) &&
297  "vector.transfer operation in which super-vector size is not an"
298  " integer multiple of sub-vector size");
299 
300  // This catches cases that are not strictly necessary to have multiplicity but
301  // still aren't divisible by the sub-vector shape.
302  // This could be useful information if we wanted to reshape at the level of
303  // the vector type (but we would have to look at the compute and distinguish
304  // between parallel, reduction and possibly other cases.
305  return ratio.hasValue();
306 }
307 
308 bool mlir::isDisjointTransferIndices(VectorTransferOpInterface transferA,
309  VectorTransferOpInterface transferB) {
310  // For simplicity only look at transfer of same type.
311  if (transferA.getVectorType() != transferB.getVectorType())
312  return false;
313  unsigned rankOffset = transferA.getLeadingShapedRank();
314  for (unsigned i = 0, e = transferA.indices().size(); i < e; i++) {
315  auto indexA = transferA.indices()[i].getDefiningOp<arith::ConstantOp>();
316  auto indexB = transferB.indices()[i].getDefiningOp<arith::ConstantOp>();
317  // If any of the indices are dynamic we cannot prove anything.
318  if (!indexA || !indexB)
319  continue;
320 
321  if (i < rankOffset) {
322  // For leading dimensions, if we can prove that index are different we
323  // know we are accessing disjoint slices.
324  if (indexA.getValue().cast<IntegerAttr>().getInt() !=
325  indexB.getValue().cast<IntegerAttr>().getInt())
326  return true;
327  } else {
328  // For this dimension, we slice a part of the memref we need to make sure
329  // the intervals accessed don't overlap.
330  int64_t distance =
331  std::abs(indexA.getValue().cast<IntegerAttr>().getInt() -
332  indexB.getValue().cast<IntegerAttr>().getInt());
333  if (distance >= transferA.getVectorType().getDimSize(i - rankOffset))
334  return true;
335  }
336  }
337  return false;
338 }
339 
340 bool mlir::isDisjointTransferSet(VectorTransferOpInterface transferA,
341  VectorTransferOpInterface transferB) {
342  if (transferA.source() != transferB.source())
343  return false;
344  return isDisjointTransferIndices(transferA, transferB);
345 }
346 
347 bool mlir::checkSameValueRAW(vector::TransferWriteOp defWrite,
348  vector::TransferReadOp read) {
349  return !defWrite.hasOutOfBoundsDim() && !defWrite.mask() && !read.mask() &&
350  defWrite.indices() == read.indices() &&
351  defWrite.getVectorType() == read.getVectorType() &&
352  defWrite.permutation_map() == read.permutation_map();
353 }
354 
355 bool mlir::checkSameValueWAW(vector::TransferWriteOp write,
356  vector::TransferWriteOp priorWrite) {
357  return priorWrite.indices() == write.indices() &&
358  priorWrite.mask() == write.mask() &&
359  priorWrite.getVectorType() == write.getVectorType() &&
360  priorWrite.permutation_map() == write.permutation_map();
361 }
362 
363 SmallVector<int64_t, 4> mlir::getI64SubArray(ArrayAttr arrayAttr,
364  unsigned dropFront,
365  unsigned dropBack) {
366  assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds");
367  auto range = arrayAttr.getAsRange<IntegerAttr>();
369  res.reserve(arrayAttr.size() - dropFront - dropBack);
370  for (auto it = range.begin() + dropFront, eit = range.end() - dropBack;
371  it != eit; ++it)
372  res.push_back((*it).getValue().getSExtValue());
373  return res;
374 }
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:52
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:95
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:87
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:444
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:516
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:96
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:244
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:276
AffineMap getTransferMinorIdentityMap(ShapedType shapedType, VectorType vectorType)
Build the default minor identity map suitable for a vector transfer.
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:45
A multi-dimensional affine map Affine map&#39;s are immutable like Type&#39;s, and they are uniqued...
Definition: AffineMap.h:38
static AffineMap getMinorIdentityMap(unsigned dims, unsigned results, MLIRContext *context)
Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
Definition: AffineMap.cpp:102
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Definition: AffineExpr.cpp:491
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:35
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:84
SmallVector< int64_t, 4 > delinearize(ArrayRef< int64_t > strides, int64_t linearIndex)
Given the strides together with a linear index in the dimension space, returns the vector-space offse...
Definition: VectorUtils.cpp:75
void dropFront(int64_t arr[N], int64_t *res)
Definition: CRunnerUtils.h:117
Type getType() const
Return the type of this value.
Definition: Value.h:117
static VectorType vectorType(CodeGen &codegen, Type etp)
Constructs vector type.
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:273
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:234
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:231
This class helps build Operations.
Definition: Builders.h:177
int64_t linearize(ArrayRef< int64_t > offsets, ArrayRef< int64_t > basis)
Computes and returns the linearized index of &#39;offsets&#39; w.r.t. &#39;basis&#39;.
Definition: VectorUtils.cpp:67