30 #include "llvm/ADT/DenseSet.h" 31 #include "llvm/ADT/SetVector.h" 41 if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
43 llvm_unreachable(
"Expected MemRefType or TensorType");
50 return std::accumulate(basis.begin(), basis.end(), 1,
51 std::multiplies<int64_t>());
56 int64_t rank = shape.size();
59 for (int64_t r = 0; r < rank; ++r)
60 sliceDimCounts[r] =
ceilDiv(shape[r], sizes[r]);
63 sliceStrides[rank - 1] = 1;
64 for (int64_t r = rank - 2; r >= 0; --r)
65 sliceStrides[r] = sliceStrides[r + 1] * sliceDimCounts[r + 1];
72 for (
auto it : llvm::zip(vectorOffsets, sizes))
73 result.push_back(std::get<0>(it) * std::get<1>(it));
79 if (superShape.size() < subShape.size()) {
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");
92 if (superSize % subSize != 0)
94 result.push_back(superSize / subSize);
100 int commonSize = subShape.size();
101 std::copy(superShape.rbegin() + commonSize, superShape.rend(),
102 std::back_inserter(result));
104 assert(result.size() == superShape.size() &&
105 "super to sub shape ratio is not of the same size as the super rank");
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());
143 if (enclosingLoopToVectorDim.empty())
146 enclosingLoopToVectorDim.begin()->getFirst()->getContext();
147 SmallVector<AffineExpr, 4>
perm(enclosingLoopToVectorDim.size(),
150 for (
auto kvp : enclosingLoopToVectorDim) {
151 assert(kvp.second <
perm.size());
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])) {
159 "permutationMap already has an entry along dim");
162 ++countInvariantIndices;
165 assert((countInvariantIndices == numIndices ||
166 countInvariantIndices == numIndices - 1) &&
167 "Vectorization prerequisite violated: at most 1 index may be " 168 "invariant wrt a vectorized loop");
177 template <
typename T>
182 if (
auto typedParent = dyn_cast<T>(current)) {
183 assert(res.count(current) == 0 &&
"Already inserted");
186 current = current->getParentOp();
193 return getParentsOfType<AffineForOp>(block);
201 for (
auto *forInst : enclosingLoops) {
202 auto it = loopToVectorDim.find(forInst);
203 if (it != loopToVectorDim.end()) {
204 enclosingLoopToVectorDim.insert(*it);
216 bool matcher::operatesOnSuperVectorsOf(
Operation &op,
217 VectorType subVectorType) {
227 bool mustDivide =
false;
229 VectorType superVectorType;
230 if (
auto transfer = dyn_cast<VectorTransferOpInterface>(op)) {
231 superVectorType = transfer.getVectorType();
234 if (!isa<func::ReturnOp>(op)) {
235 op.
emitError(
"NYI: assuming only return operations can have 0 " 236 " results at this point");
249 op.
emitError(
"NYI: operation has more than 1 result");
254 auto ratio =
shapeRatio(superVectorType, subVectorType);
257 assert((ratio || !mustDivide) &&
258 "vector.transfer operation in which super-vector size is not an" 259 " integer multiple of sub-vector size");
266 return ratio.has_value();
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...
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.
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...
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...
Operation is a basic unit of execution within MLIR.
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Block represents an ordered list of Operations.
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.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
int64_t ceilDiv(int64_t lhs, int64_t rhs)
Returns the result of MLIR's ceildiv operation on constants.
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
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.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued...
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
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...
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
MLIRContext is the top-level object for a collection of MLIR operations.
unsigned getNumResults()
Return the number of results held by this operation.
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...
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
This class helps build Operations.