MLIR  19.0.0git
VectorDropLeadUnitDim.cpp
Go to the documentation of this file.
1 //===- VectorDropLeadUnitDim.cpp - Conversion within the Vector dialect ---===//
2 //
3 // Part of the LLVM 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 #include <numeric>
10 
17 #include "mlir/IR/Builders.h"
18 #include "mlir/IR/TypeUtilities.h"
19 
20 #define DEBUG_TYPE "vector-drop-unit-dim"
21 
22 using namespace mlir;
23 using namespace mlir::vector;
24 
25 // Trims leading one dimensions from `oldType` and returns the result type.
26 // Returns `vector<1xT>` if `oldType` only has one element.
27 static VectorType trimLeadingOneDims(VectorType oldType) {
28  ArrayRef<int64_t> oldShape = oldType.getShape();
29  ArrayRef<int64_t> newShape = oldShape;
30 
31  ArrayRef<bool> oldScalableDims = oldType.getScalableDims();
32  ArrayRef<bool> newScalableDims = oldScalableDims;
33 
34  while (!newShape.empty() && newShape.front() == 1 &&
35  !newScalableDims.front()) {
36  newShape = newShape.drop_front(1);
37  newScalableDims = newScalableDims.drop_front(1);
38  }
39 
40  // Make sure we have at least 1 dimension per vector type requirements.
41  if (newShape.empty()) {
42  newShape = oldShape.take_back();
43  newScalableDims = oldType.getScalableDims().take_back();
44  }
45  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
46 }
47 
48 /// Return a smallVector of size `rank` containing all zeros.
49 static SmallVector<int64_t> splatZero(int64_t rank) {
50  return SmallVector<int64_t>(rank, 0);
51 }
52 namespace {
53 
54 // Casts away leading one dimensions in vector.extract_strided_slice's vector
55 // input by inserting vector.broadcast.
56 struct CastAwayExtractStridedSliceLeadingOneDim
57  : public OpRewritePattern<vector::ExtractStridedSliceOp> {
59 
60  LogicalResult matchAndRewrite(vector::ExtractStridedSliceOp extractOp,
61  PatternRewriter &rewriter) const override {
62  // vector.extract_strided_slice requires the input and output vector to have
63  // the same rank. Here we drop leading one dimensions from the input vector
64  // type to make sure we don't cause mismatch.
65  VectorType oldSrcType = extractOp.getSourceVectorType();
66  VectorType newSrcType = trimLeadingOneDims(oldSrcType);
67 
68  if (newSrcType.getRank() == oldSrcType.getRank())
69  return failure();
70 
71  int64_t dropCount = oldSrcType.getRank() - newSrcType.getRank();
72 
73  VectorType oldDstType = extractOp.getType();
74  VectorType newDstType =
75  VectorType::get(oldDstType.getShape().drop_front(dropCount),
76  oldDstType.getElementType(),
77  oldDstType.getScalableDims().drop_front(dropCount));
78 
79  Location loc = extractOp.getLoc();
80 
81  Value newSrcVector = rewriter.create<vector::ExtractOp>(
82  loc, extractOp.getVector(), splatZero(dropCount));
83 
84  // The offsets/sizes/strides attribute can have a less number of elements
85  // than the input vector's rank: it is meant for the leading dimensions.
86  auto newOffsets = rewriter.getArrayAttr(
87  extractOp.getOffsets().getValue().drop_front(dropCount));
88  auto newSizes = rewriter.getArrayAttr(
89  extractOp.getSizes().getValue().drop_front(dropCount));
90  auto newStrides = rewriter.getArrayAttr(
91  extractOp.getStrides().getValue().drop_front(dropCount));
92 
93  auto newExtractOp = rewriter.create<vector::ExtractStridedSliceOp>(
94  loc, newDstType, newSrcVector, newOffsets, newSizes, newStrides);
95 
96  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(extractOp, oldDstType,
97  newExtractOp);
98 
99  return success();
100  }
101 };
102 
103 // Casts away leading one dimensions in vector.insert_strided_slice's vector
104 // inputs by inserting vector.broadcast.
105 struct CastAwayInsertStridedSliceLeadingOneDim
106  : public OpRewritePattern<vector::InsertStridedSliceOp> {
108 
109  LogicalResult matchAndRewrite(vector::InsertStridedSliceOp insertOp,
110  PatternRewriter &rewriter) const override {
111  VectorType oldSrcType = insertOp.getSourceVectorType();
112  VectorType newSrcType = trimLeadingOneDims(oldSrcType);
113  VectorType oldDstType = insertOp.getDestVectorType();
114  VectorType newDstType = trimLeadingOneDims(oldDstType);
115 
116  int64_t srcDropCount = oldSrcType.getRank() - newSrcType.getRank();
117  int64_t dstDropCount = oldDstType.getRank() - newDstType.getRank();
118  if (srcDropCount == 0 && dstDropCount == 0)
119  return failure();
120 
121  // Trim leading one dimensions from both operands.
122  Location loc = insertOp.getLoc();
123 
124  Value newSrcVector = rewriter.create<vector::ExtractOp>(
125  loc, insertOp.getSource(), splatZero(srcDropCount));
126  Value newDstVector = rewriter.create<vector::ExtractOp>(
127  loc, insertOp.getDest(), splatZero(dstDropCount));
128 
129  auto newOffsets = rewriter.getArrayAttr(
130  insertOp.getOffsets().getValue().take_back(newDstType.getRank()));
131  auto newStrides = rewriter.getArrayAttr(
132  insertOp.getStrides().getValue().take_back(newSrcType.getRank()));
133 
134  auto newInsertOp = rewriter.create<vector::InsertStridedSliceOp>(
135  loc, newDstType, newSrcVector, newDstVector, newOffsets, newStrides);
136 
137  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(insertOp, oldDstType,
138  newInsertOp);
139 
140  return success();
141  }
142 };
143 
144 // Casts away leading one dimensions in vector.insert's vector inputs by
145 // inserting vector.broadcast.
146 struct CastAwayInsertLeadingOneDim : public OpRewritePattern<vector::InsertOp> {
148 
149  LogicalResult matchAndRewrite(vector::InsertOp insertOp,
150  PatternRewriter &rewriter) const override {
151  Type oldSrcType = insertOp.getSourceType();
152  Type newSrcType = oldSrcType;
153  int64_t oldSrcRank = 0, newSrcRank = 0;
154  if (auto type = dyn_cast<VectorType>(oldSrcType)) {
155  newSrcType = trimLeadingOneDims(type);
156  oldSrcRank = type.getRank();
157  newSrcRank = cast<VectorType>(newSrcType).getRank();
158  }
159 
160  VectorType oldDstType = insertOp.getDestVectorType();
161  VectorType newDstType = trimLeadingOneDims(oldDstType);
162 
163  int64_t srcDropCount = oldSrcRank - newSrcRank;
164  int64_t dstDropCount = oldDstType.getRank() - newDstType.getRank();
165  if (srcDropCount == 0 && dstDropCount == 0)
166  return failure();
167 
168  // Trim leading one dimensions from both operands.
169  Location loc = insertOp.getLoc();
170 
171  Value newSrcVector = insertOp.getSource();
172  if (oldSrcRank != 0) {
173  newSrcVector = rewriter.create<vector::ExtractOp>(
174  loc, insertOp.getSource(), splatZero(srcDropCount));
175  }
176  Value newDstVector = rewriter.create<vector::ExtractOp>(
177  loc, insertOp.getDest(), splatZero(dstDropCount));
178 
179  // New position rank needs to be computed in two steps: (1) if destination
180  // type has leading unit dims, we also trim the position array accordingly,
181  // then (2) if source type also has leading unit dims, we need to append
182  // zeroes to the position array accordingly.
183  unsigned oldPosRank = insertOp.getNumIndices();
184  unsigned newPosRank = std::max<int64_t>(0, oldPosRank - dstDropCount);
185  SmallVector<OpFoldResult> oldPosition = insertOp.getMixedPosition();
186  SmallVector<OpFoldResult> newPosition =
187  llvm::to_vector(ArrayRef(oldPosition).take_back(newPosRank));
188  newPosition.resize(newDstType.getRank() - newSrcRank,
189  rewriter.getI64IntegerAttr(0));
190 
191  auto newInsertOp = rewriter.create<vector::InsertOp>(
192  loc, newSrcVector, newDstVector, newPosition);
193 
194  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(insertOp, oldDstType,
195  newInsertOp);
196 
197  return success();
198  }
199 };
200 
201 static Value dropUnitDimsFromMask(OpBuilder &b, Location loc, Value mask,
202  VectorType newType, AffineMap newMap,
203  VectorType oldMaskType) {
204  // Infer the type of the new mask from the new map.
205  VectorType newMaskType = inferTransferOpMaskType(newType, newMap);
206 
207  // If the new mask is broadcastable to the old result type, we can safely
208  // use a `vector.extract` to get the new mask. Otherwise the best we can
209  // do is shape cast.
210  if (vector::isBroadcastableTo(newMaskType, oldMaskType) ==
212  int64_t dropDim = oldMaskType.getRank() - newMaskType.getRank();
213  return b.create<vector::ExtractOp>(loc, mask, splatZero(dropDim));
214  }
215  return b.create<vector::ShapeCastOp>(loc, newMaskType, mask);
216 }
217 
218 // Turns vector.transfer_read on vector with leading 1 dimensions into
219 // vector.shape_cast followed by vector.transfer_read on vector without leading
220 // 1 dimensions.
221 struct CastAwayTransferReadLeadingOneDim
222  : public OpRewritePattern<vector::TransferReadOp> {
224 
225  LogicalResult matchAndRewrite(vector::TransferReadOp read,
226  PatternRewriter &rewriter) const override {
227  // TODO(#78787): Not supported masked op yet.
228  if (cast<MaskableOpInterface>(read.getOperation()).isMasked())
229  return failure();
230  // TODO: support 0-d corner case.
231  if (read.getTransferRank() == 0)
232  return failure();
233 
234  auto shapedType = cast<ShapedType>(read.getSource().getType());
235  if (shapedType.getElementType() != read.getVectorType().getElementType())
236  return failure();
237 
238  VectorType oldType = read.getVectorType();
239  VectorType newType = trimLeadingOneDims(oldType);
240 
241  if (newType == oldType)
242  return failure();
243 
244  AffineMap oldMap = read.getPermutationMap();
245  ArrayRef<AffineExpr> newResults =
246  oldMap.getResults().take_back(newType.getRank());
247  AffineMap newMap =
248  AffineMap::get(oldMap.getNumDims(), oldMap.getNumSymbols(), newResults,
249  rewriter.getContext());
250 
251  ArrayAttr inBoundsAttr;
252  if (read.getInBounds())
253  inBoundsAttr = rewriter.getArrayAttr(
254  read.getInBoundsAttr().getValue().take_back(newType.getRank()));
255 
256  Value mask = Value();
257  if (read.getMask()) {
258  VectorType maskType = read.getMaskType();
259  mask = dropUnitDimsFromMask(rewriter, read.getLoc(), read.getMask(),
260  newType, newMap, maskType);
261  }
262 
263  auto newRead = rewriter.create<vector::TransferReadOp>(
264  read.getLoc(), newType, read.getSource(), read.getIndices(),
265  AffineMapAttr::get(newMap), read.getPadding(), mask, inBoundsAttr);
266  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(read, oldType, newRead);
267 
268  return success();
269  }
270 };
271 
272 // Turns vector.transfer_write on vector with leading 1 dimensions into
273 // vector.shape_cast followed by vector.transfer_write on vector without leading
274 // 1 dimensions.
275 struct CastAwayTransferWriteLeadingOneDim
276  : public OpRewritePattern<vector::TransferWriteOp> {
278 
279  LogicalResult matchAndRewrite(vector::TransferWriteOp write,
280  PatternRewriter &rewriter) const override {
281  // TODO(#78787): Not supported masked op yet.
282  if (cast<MaskableOpInterface>(write.getOperation()).isMasked())
283  return failure();
284  // TODO: support 0-d corner case.
285  if (write.getTransferRank() == 0)
286  return failure();
287 
288  auto shapedType = dyn_cast<ShapedType>(write.getSource().getType());
289  if (shapedType.getElementType() != write.getVectorType().getElementType())
290  return failure();
291 
292  VectorType oldType = write.getVectorType();
293  VectorType newType = trimLeadingOneDims(oldType);
294  if (newType == oldType)
295  return failure();
296  int64_t dropDim = oldType.getRank() - newType.getRank();
297 
298  AffineMap oldMap = write.getPermutationMap();
299  ArrayRef<AffineExpr> newResults =
300  oldMap.getResults().take_back(newType.getRank());
301  AffineMap newMap =
302  AffineMap::get(oldMap.getNumDims(), oldMap.getNumSymbols(), newResults,
303  rewriter.getContext());
304 
305  ArrayAttr inBoundsAttr;
306  if (write.getInBounds())
307  inBoundsAttr = rewriter.getArrayAttr(
308  write.getInBoundsAttr().getValue().take_back(newType.getRank()));
309 
310  auto newVector = rewriter.create<vector::ExtractOp>(
311  write.getLoc(), write.getVector(), splatZero(dropDim));
312 
313  if (write.getMask()) {
314  VectorType maskType = write.getMaskType();
315  Value newMask = dropUnitDimsFromMask(
316  rewriter, write.getLoc(), write.getMask(), newType, newMap, maskType);
317  rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
318  write, newVector, write.getSource(), write.getIndices(),
319  AffineMapAttr::get(newMap), newMask, inBoundsAttr);
320  return success();
321  }
322 
323  rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
324  write, newVector, write.getSource(), write.getIndices(),
325  AffineMapAttr::get(newMap), inBoundsAttr);
326  return success();
327  }
328 };
329 
330 } // namespace
331 
333 mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
334  MaskingOpInterface maskingOp,
335  RewriterBase &rewriter) {
336  VectorType oldAccType = dyn_cast<VectorType>(contractOp.getAccType());
337  if (oldAccType == nullptr)
338  return failure();
339  if (oldAccType.getRank() < 2)
340  return failure();
341  if (oldAccType.getShape()[0] != 1)
342  return failure();
343  // currently we support only dropping one dim but the pattern can be applied
344  // greedily to drop more.
345  int64_t dropDim = 1;
346 
347  auto oldIndexingMaps = contractOp.getIndexingMapsArray();
348  SmallVector<AffineMap> newIndexingMaps;
349 
350  auto oldIteratorTypes = contractOp.getIteratorTypes();
351  SmallVector<Attribute> newIteratorTypes;
352 
353  int64_t dimToDrop = oldIndexingMaps[2].getDimPosition(0);
354 
355  if (!isParallelIterator(oldIteratorTypes[dimToDrop]))
356  // only parallel type iterators can be dropped.
357  return failure();
358 
359  for (const auto &it : llvm::enumerate(oldIteratorTypes)) {
360  int64_t currDim = it.index();
361  if (currDim == dimToDrop)
362  continue;
363  newIteratorTypes.push_back(it.value());
364  }
365 
366  SmallVector<Value> operands = {contractOp.getLhs(), contractOp.getRhs(),
367  contractOp.getAcc()};
368  SmallVector<Value> newOperands;
369  auto loc = contractOp.getLoc();
370 
371  for (const auto &it : llvm::enumerate(oldIndexingMaps)) {
372  // Check if the dim to be dropped exists as a leading dim in the operand
373  // if it does then we use vector.extract to drop it.
374  bool validExtract = false;
375  SmallVector<AffineExpr> results;
376  auto map = it.value();
377  int64_t orginalZeroDim = it.value().getDimPosition(0);
378  if (orginalZeroDim != dimToDrop) {
379  // There are two reasons to be in this path, 1. We need to
380  // tranpose the operand to make the dim to be dropped
381  // leading. 2. The dim to be dropped does not exist and in
382  // that case we dont want to add a unit tranpose but we must
383  // check all the indices to make sure this is the case.
384  bool tranposeNeeded = false;
386  SmallVector<AffineExpr> transposeResults;
387 
388  for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
389  int64_t currDim = map.getDimPosition(i);
390  if (currDim == dimToDrop) {
391  tranposeNeeded = true;
392  perm.insert(perm.begin(), i);
393  auto targetExpr = rewriter.getAffineDimExpr(currDim);
394  transposeResults.insert(transposeResults.begin(), targetExpr);
395  } else {
396  perm.push_back(i);
397  auto targetExpr = rewriter.getAffineDimExpr(currDim);
398  transposeResults.push_back(targetExpr);
399  }
400  }
401 
402  // Checks if only the outer, unit dimensions (of size 1) are permuted.
403  // Such transposes do not materially effect the underlying vector and can
404  // be omitted. EG: perm [1, 0, 2] applied to vector<1x1x8xi32>
405  bool transposeNonOuterUnitDims = false;
406  auto operandShape = operands[it.index()].getType().cast<ShapedType>();
407  for (auto [index, dim] :
408  llvm::enumerate(ArrayRef<int64_t>(perm).drop_back(1))) {
409  if (dim != static_cast<int64_t>(index) &&
410  operandShape.getDimSize(index) != 1) {
411  transposeNonOuterUnitDims = true;
412  break;
413  }
414  }
415 
416  // Do the tranpose now if needed so that we can drop the
417  // correct dim using extract later.
418  if (tranposeNeeded) {
419  map = AffineMap::get(map.getNumDims(), 0, transposeResults,
420  contractOp.getContext());
421  if (transposeNonOuterUnitDims) {
422  operands[it.index()] = rewriter.createOrFold<vector::TransposeOp>(
423  loc, operands[it.index()], perm);
424  }
425  }
426  }
427  // We have taken care to have the dim to be dropped be
428  // the leading dim. If its still not leading that means it
429  // does not exist in this operand and hence we do not need
430  // an extract.
431  if (map.getDimPosition(0) == dimToDrop)
432  validExtract = true;
433 
434  for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
435  int64_t currDim = map.getDimPosition(i);
436  if (currDim == dimToDrop)
437  // This is the dim we are dropping.
438  continue;
439  auto targetExpr = rewriter.getAffineDimExpr(
440  currDim < dimToDrop ? currDim : currDim - 1);
441  results.push_back(targetExpr);
442  }
443  newIndexingMaps.push_back(AffineMap::get(map.getNumDims() - 1, 0, results,
444  contractOp.getContext()));
445  // Extract if its a valid extraction, otherwise use the operand
446  // without extraction.
447  newOperands.push_back(
448  validExtract ? rewriter.create<vector::ExtractOp>(
449  loc, operands[it.index()], splatZero(dropDim))
450  : operands[it.index()]);
451  }
452 
453  // Depending on whether this vector.contract is masked, the replacing Op
454  // should either be a new vector.contract Op or vector.mask Op.
455  Operation *newOp = rewriter.create<vector::ContractionOp>(
456  loc, newOperands[0], newOperands[1], newOperands[2],
457  rewriter.getAffineMapArrayAttr(newIndexingMaps),
458  rewriter.getArrayAttr(newIteratorTypes), contractOp.getKind());
459 
460  if (maskingOp) {
461  auto newMask = rewriter.create<vector::ExtractOp>(loc, maskingOp.getMask(),
462  splatZero(dropDim));
463 
464  newOp = mlir::vector::maskOperation(rewriter, newOp, newMask);
465  }
466 
467  return rewriter
468  .create<vector::BroadcastOp>(loc, contractOp->getResultTypes()[0],
469  newOp->getResults()[0])
470  .getResult();
471 }
472 
473 namespace {
474 
475 /// Turns vector.contract on vector with leading 1 dimensions into
476 /// vector.extract followed by vector.contract on vector without leading
477 /// 1 dimensions. Also performs tranpose of lhs and rhs operands if required
478 /// prior to extract.
479 struct CastAwayContractionLeadingOneDim
480  : public MaskableOpRewritePattern<vector::ContractionOp> {
481  using MaskableOpRewritePattern::MaskableOpRewritePattern;
482 
484  matchAndRewriteMaskableOp(vector::ContractionOp contractOp,
485  MaskingOpInterface maskingOp,
486  PatternRewriter &rewriter) const override {
487  return castAwayContractionLeadingOneDim(contractOp, maskingOp, rewriter);
488  }
489 };
490 
491 /// Looks at elementwise operations on vectors with at least one leading
492 /// dimension equal 1, e.g. vector<1x[4]x1xf32> (but not vector<2x[4]x1xf32>),
493 /// and cast aways the leading one dimensions (_plural_) and then broadcasts
494 /// the results.
495 ///
496 /// Example before:
497 /// %1 = arith.mulf %arg0, %arg1 : vector<1x4x1xf32>
498 /// Example after:
499 /// %2 = arith.mulf %0, %1 : vector<4x1xf32>
500 /// %3 = vector.broadcast %2 : vector<4x1xf32> to vector<1x4x1xf32>
501 ///
502 /// Does support scalable vectors.
503 class CastAwayElementwiseLeadingOneDim : public RewritePattern {
504 public:
505  CastAwayElementwiseLeadingOneDim(MLIRContext *context,
506  PatternBenefit benefit = 1)
507  : RewritePattern(MatchAnyOpTypeTag(), benefit, context) {}
508 
509  LogicalResult matchAndRewrite(Operation *op,
510  PatternRewriter &rewriter) const override {
512  return failure();
513  auto vecType = dyn_cast<VectorType>(op->getResultTypes()[0]);
514  if (!vecType)
515  return failure();
516  VectorType newVecType = trimLeadingOneDims(vecType);
517  if (newVecType == vecType)
518  return failure();
519  int64_t dropDim = vecType.getRank() - newVecType.getRank();
520  SmallVector<Value, 4> newOperands;
521  for (Value operand : op->getOperands()) {
522  if (auto opVecType = dyn_cast<VectorType>(operand.getType())) {
523  newOperands.push_back(rewriter.create<vector::ExtractOp>(
524  op->getLoc(), operand, splatZero(dropDim)));
525  } else {
526  newOperands.push_back(operand);
527  }
528  }
529  Operation *newOp =
530  rewriter.create(op->getLoc(), op->getName().getIdentifier(),
531  newOperands, newVecType, op->getAttrs());
532  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, vecType,
533  newOp->getResult(0));
534  return success();
535  }
536 };
537 
538 // Drops leading 1 dimensions from vector.constant_mask and inserts a
539 // vector.broadcast back to the original shape.
540 struct CastAwayConstantMaskLeadingOneDim
541  : public OpRewritePattern<vector::ConstantMaskOp> {
543 
544  LogicalResult matchAndRewrite(vector::ConstantMaskOp mask,
545  PatternRewriter &rewriter) const override {
546  VectorType oldType = mask.getType();
547  VectorType newType = trimLeadingOneDims(oldType);
548 
549  if (newType == oldType)
550  return failure();
551 
552  int64_t dropDim = oldType.getRank() - newType.getRank();
553  SmallVector<int64_t> dimSizes;
554  for (auto attr : mask.getMaskDimSizes())
555  dimSizes.push_back(llvm::cast<IntegerAttr>(attr).getInt());
556 
557  // If any of the dropped unit dims has a size of `0`, the entire mask is a
558  // zero mask, else the unit dim has no effect on the mask.
559  int64_t flatLeadingSize =
560  std::accumulate(dimSizes.begin(), dimSizes.begin() + dropDim + 1,
561  static_cast<int64_t>(1), std::multiplies<int64_t>());
562  SmallVector<int64_t> newDimSizes({flatLeadingSize});
563  newDimSizes.append(dimSizes.begin() + dropDim + 1, dimSizes.end());
564 
565  auto newMask = rewriter.create<vector::ConstantMaskOp>(
566  mask.getLoc(), newType, rewriter.getI64ArrayAttr(newDimSizes));
567  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(mask, oldType, newMask);
568  return success();
569  }
570 };
571 
572 } // namespace
573 
575  RewritePatternSet &patterns, PatternBenefit benefit) {
576  patterns
577  .add<CastAwayExtractStridedSliceLeadingOneDim,
578  CastAwayInsertStridedSliceLeadingOneDim, CastAwayInsertLeadingOneDim,
579  CastAwayConstantMaskLeadingOneDim, CastAwayTransferReadLeadingOneDim,
580  CastAwayTransferWriteLeadingOneDim, CastAwayElementwiseLeadingOneDim,
581  CastAwayContractionLeadingOneDim>(patterns.getContext(), benefit);
582  populateShapeCastFoldingPatterns(patterns, benefit);
583 }
static SmallVector< int64_t > splatZero(int64_t rank)
Return a smallVector of size rank containing all zeros.
static VectorType trimLeadingOneDims(VectorType oldType)
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:47
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
unsigned getNumSymbols() const
Definition: AffineMap.cpp:382
unsigned getNumDims() const
Definition: AffineMap.cpp:378
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:391
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
Definition: AffineMap.cpp:248
IntegerAttr getI64IntegerAttr(int64_t value)
Definition: Builders.cpp:128
AffineExpr getAffineDimExpr(unsigned position)
Definition: Builders.cpp:371
MLIRContext * getContext() const
Definition: Builders.h:55
ArrayAttr getArrayAttr(ArrayRef< Attribute > value)
Definition: Builders.cpp:273
ArrayAttr getI64ArrayAttr(ArrayRef< int64_t > values)
Definition: Builders.cpp:288
ArrayAttr getAffineMapArrayAttr(ArrayRef< AffineMap > values)
Definition: Builders.cpp:325
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:209
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:522
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:464
StringAttr getIdentifier() const
Return the name of this operation as a StringAttr.
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
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
ArrayRef< NamedAttribute > getAttrs()
Return all of the attributes on this operation.
Definition: Operation.h:507
OperationName getName()
The name of an operation is the key identifier for it.
Definition: Operation.h:119
result_type_range getResultTypes()
Definition: Operation.h:423
operand_range getOperands()
Returns an iterator on the underlying Value's.
Definition: Operation.h:373
result_range getResults()
Definition: Operation.h:410
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:399
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
Definition: PatternMatch.h:34
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:785
MLIRContext * getContext() const
Definition: PatternMatch.h:822
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
Definition: PatternMatch.h:846
RewritePattern is the common base class for all DAG to DAG replacements.
Definition: PatternMatch.h:246
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:400
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:536
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
bool hasElementwiseMappableTraits(Operation *op)
Together, Elementwise, Scalarizable, Vectorizable, and Tensorizable provide an easy way for scalar op...
Definition: Operation.cpp:1391
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
FailureOr< Value > castAwayContractionLeadingOneDim(vector::ContractionOp contractOp, MaskingOpInterface maskingOp, RewriterBase &rewriter)
Cast away the leading unit dim, if exists, for the given contract op.
void populateShapeCastFoldingPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Collect a set of vector.shape_cast folding patterns.
Operation * maskOperation(OpBuilder &builder, Operation *maskableOp, Value mask, Value passthru=Value())
Creates a vector.mask operation around a maskable operation.
VectorType inferTransferOpMaskType(VectorType vecType, AffineMap permMap)
Infers the mask type for a transfer op given its vector type and permutation map.
Definition: VectorOps.cpp:3847
BroadcastableToResult isBroadcastableTo(Type srcType, VectorType dstVectorType, std::pair< int, int > *mismatchingDims=nullptr)
Definition: VectorOps.cpp:2236
bool isParallelIterator(Attribute attr)
Returns true if attr has "parallel" iterator type semantics.
Definition: VectorOps.h:135
void populateCastAwayVectorLeadingOneDimPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Collect a set of leading one dimension removal patterns.
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:358
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...
Definition: PatternMatch.h:362
A pattern for ops that implement MaskableOpInterface and that might be masked (i.e.
Definition: VectorUtils.h:133