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 
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 (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 /// Constructs a permutation map from memref indices to vector dimension.
47 ///
48 /// The implementation uses the knowledge of the mapping of enclosing loop to
49 /// vector dimension. `enclosingLoopToVectorDim` carries this information as a
50 /// map with:
51 /// - keys representing "vectorized enclosing loops";
52 /// - values representing the corresponding vector dimension.
53 /// The algorithm traverses "vectorized enclosing loops" and extracts the
54 /// at-most-one MemRef index that is invariant along said loop. This index is
55 /// guaranteed to be at most one by construction: otherwise the MemRef is not
56 /// vectorizable.
57 /// If this invariant index is found, it is added to the permutation_map at the
58 /// proper vector dimension.
59 /// If no index is found to be invariant, 0 is added to the permutation_map and
60 /// corresponds to a vector broadcast along that dimension.
61 ///
62 /// Returns an empty AffineMap if `enclosingLoopToVectorDim` is empty,
63 /// signalling that no permutation map can be constructed given
64 /// `enclosingLoopToVectorDim`.
65 ///
66 /// Examples can be found in the documentation of `makePermutationMap`, in the
67 /// header file.
69  ArrayRef<Value> indices,
70  const DenseMap<Operation *, unsigned> &enclosingLoopToVectorDim) {
71  if (enclosingLoopToVectorDim.empty())
72  return AffineMap();
73  MLIRContext *context =
74  enclosingLoopToVectorDim.begin()->getFirst()->getContext();
75  SmallVector<AffineExpr> perm(enclosingLoopToVectorDim.size(),
76  getAffineConstantExpr(0, context));
77 
78  for (auto kvp : enclosingLoopToVectorDim) {
79  assert(kvp.second < perm.size());
80  auto invariants = getInvariantAccesses(
81  cast<AffineForOp>(kvp.first).getInductionVar(), indices);
82  unsigned numIndices = indices.size();
83  unsigned countInvariantIndices = 0;
84  for (unsigned dim = 0; dim < numIndices; ++dim) {
85  if (!invariants.count(indices[dim])) {
86  assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
87  "permutationMap already has an entry along dim");
88  perm[kvp.second] = getAffineDimExpr(dim, context);
89  } else {
90  ++countInvariantIndices;
91  }
92  }
93  assert((countInvariantIndices == numIndices ||
94  countInvariantIndices == numIndices - 1) &&
95  "Vectorization prerequisite violated: at most 1 index may be "
96  "invariant wrt a vectorized loop");
97  (void)countInvariantIndices;
98  }
99  return AffineMap::get(indices.size(), 0, perm, context);
100 }
101 
102 /// Implementation detail that walks up the parents and records the ones with
103 /// the specified type.
104 /// TODO: could also be implemented as a collect parents followed by a
105 /// filter and made available outside this file.
106 template <typename T>
109  auto *current = block->getParentOp();
110  while (current) {
111  if (auto typedParent = dyn_cast<T>(current)) {
112  assert(res.count(current) == 0 && "Already inserted");
113  res.insert(current);
114  }
115  current = current->getParentOp();
116  }
117  return res;
118 }
119 
120 /// Returns the enclosing AffineForOp, from closest to farthest.
122  return getParentsOfType<AffineForOp>(block);
123 }
124 
126  Block *insertPoint, ArrayRef<Value> indices,
127  const DenseMap<Operation *, unsigned> &loopToVectorDim) {
128  DenseMap<Operation *, unsigned> enclosingLoopToVectorDim;
129  auto enclosingLoops = getEnclosingforOps(insertPoint);
130  for (auto *forInst : enclosingLoops) {
131  auto it = loopToVectorDim.find(forInst);
132  if (it != loopToVectorDim.end()) {
133  enclosingLoopToVectorDim.insert(*it);
134  }
135  }
136  return ::makePermutationMap(indices, enclosingLoopToVectorDim);
137 }
138 
140  Operation *op, ArrayRef<Value> indices,
141  const DenseMap<Operation *, unsigned> &loopToVectorDim) {
142  return makePermutationMap(op->getBlock(), indices, loopToVectorDim);
143 }
144 
145 bool matcher::operatesOnSuperVectorsOf(Operation &op,
146  VectorType subVectorType) {
147  // First, extract the vector type and distinguish between:
148  // a. ops that *must* lower a super-vector (i.e. vector.transfer_read,
149  // vector.transfer_write); and
150  // b. ops that *may* lower a super-vector (all other ops).
151  // The ops that *may* lower a super-vector only do so if the super-vector to
152  // sub-vector ratio exists. The ops that *must* lower a super-vector are
153  // explicitly checked for this property.
154  /// TODO: there should be a single function for all ops to do this so we
155  /// do not have to special case. Maybe a trait, or just a method, unclear atm.
156  bool mustDivide = false;
157  (void)mustDivide;
158  VectorType superVectorType;
159  if (auto transfer = dyn_cast<VectorTransferOpInterface>(op)) {
160  superVectorType = transfer.getVectorType();
161  mustDivide = true;
162  } else if (op.getNumResults() == 0) {
163  if (!isa<func::ReturnOp>(op)) {
164  op.emitError("NYI: assuming only return operations can have 0 "
165  " results at this point");
166  }
167  return false;
168  } else if (op.getNumResults() == 1) {
169  if (auto v = op.getResult(0).getType().dyn_cast<VectorType>()) {
170  superVectorType = v;
171  } else {
172  // Not a vector type.
173  return false;
174  }
175  } else {
176  // Not a vector.transfer and has more than 1 result, fail hard for now to
177  // wake us up when something changes.
178  op.emitError("NYI: operation has more than 1 result");
179  return false;
180  }
181 
182  // Get the ratio.
183  auto ratio =
184  computeShapeRatio(superVectorType.getShape(), subVectorType.getShape());
185 
186  // Sanity check.
187  assert((ratio || !mustDivide) &&
188  "vector.transfer operation in which super-vector size is not an"
189  " integer multiple of sub-vector size");
190 
191  // This catches cases that are not strictly necessary to have multiplicity but
192  // still aren't divisible by the sub-vector shape.
193  // This could be useful information if we wanted to reshape at the level of
194  // the vector type (but we would have to look at the compute and distinguish
195  // between parallel, reduction and possibly other cases.
196  return ratio.has_value();
197 }
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.
Definition: VectorUtils.cpp:68
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:42
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 defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:64
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:56
This class helps build Operations.
Definition: Builders.h:198
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:472
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:31
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:324
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:225
Block * getBlock()
Returns the operation block that contains this operation.
Definition: Operation.h:144
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:321
U dyn_cast() const
Definition: Types.h:270
bool isa() const
Definition: Types.h:260
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:114
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
Include the generated interface declarations.
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 ...
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:513
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Definition: AffineExpr.cpp:488
Optional< SmallVector< int64_t > > computeShapeRatio(ArrayRef< int64_t > shape, ArrayRef< int64_t > subShape)
Compute and return the multi-dimensional integral ratio of subShape to the trailing dimensions of sha...