MLIR  21.0.0git
VectorTransforms.cpp
Go to the documentation of this file.
1 //===- VectorTransforms.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 // This file implements target-independent rewrites as 1->N patterns.
10 //
11 //===----------------------------------------------------------------------===//
12 
14 
15 #include <cassert>
16 #include <cstdint>
17 #include <functional>
18 #include <optional>
19 #include <type_traits>
20 
34 #include "mlir/IR/BuiltinTypes.h"
36 #include "mlir/IR/Location.h"
37 #include "mlir/IR/Matchers.h"
38 #include "mlir/IR/PatternMatch.h"
39 #include "mlir/IR/TypeUtilities.h"
41 
42 #include "llvm/ADT/DenseSet.h"
43 #include "llvm/ADT/MapVector.h"
44 #include "llvm/ADT/STLExtras.h"
45 #include "llvm/Support/CommandLine.h"
46 #include "llvm/Support/Debug.h"
47 #include "llvm/Support/FormatVariadic.h"
48 #include "llvm/Support/raw_ostream.h"
49 
50 #define DEBUG_TYPE "vector-to-vector"
51 
52 using namespace mlir;
53 using namespace mlir::vector;
54 
55 template <typename IntType>
56 static SmallVector<IntType> extractVector(ArrayAttr arrayAttr) {
57  return llvm::to_vector<4>(llvm::map_range(
58  arrayAttr.getAsRange<IntegerAttr>(),
59  [](IntegerAttr attr) { return static_cast<IntType>(attr.getInt()); }));
60 }
61 
62 // Helper to find an index in an affine map.
63 static std::optional<int64_t> getResultIndex(AffineMap map, int64_t index) {
64  for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
65  int64_t idx = map.getDimPosition(i);
66  if (idx == index)
67  return i;
68  }
69  return std::nullopt;
70 }
71 
72 namespace {
73 
74 /// ShapeCastOpFolder folds cancelling ShapeCastOps away.
75 //
76 // Example:
77 //
78 // The following MLIR with cancelling ShapeCastOps:
79 //
80 // %0 = source : vector<5x4x2xf32>
81 // %1 = shape_cast %0 : vector<5x4x2xf32> to vector<20x2xf32>
82 // %2 = shape_cast %1 : vector<20x2xf32> to vector<5x4x2xf32>
83 // %3 = user %2 : vector<5x4x2xf32>
84 //
85 // Should canonicalize to the following:
86 //
87 // %0 = source : vector<5x4x2xf32>
88 // %1 = user %0 : vector<5x4x2xf32>
89 //
90 struct ShapeCastOpFolder : public OpRewritePattern<vector::ShapeCastOp> {
92 
93  LogicalResult matchAndRewrite(vector::ShapeCastOp shapeCastOp,
94  PatternRewriter &rewriter) const override {
95  // Check if 'shapeCastOp' has vector source/result type.
96  auto sourceVectorType =
97  dyn_cast_or_null<VectorType>(shapeCastOp.getSource().getType());
98  auto resultVectorType =
99  dyn_cast_or_null<VectorType>(shapeCastOp.getResult().getType());
100  if (!sourceVectorType || !resultVectorType)
101  return failure();
102 
103  // Check if shape cast op source operand is also a shape cast op.
104  auto sourceShapeCastOp = dyn_cast_or_null<vector::ShapeCastOp>(
105  shapeCastOp.getSource().getDefiningOp());
106  if (!sourceShapeCastOp)
107  return failure();
108  auto operandSourceVectorType =
109  cast<VectorType>(sourceShapeCastOp.getSource().getType());
110  auto operandResultVectorType = sourceShapeCastOp.getType();
111 
112  // Check if shape cast operations invert each other.
113  if (operandSourceVectorType != resultVectorType ||
114  operandResultVectorType != sourceVectorType)
115  return failure();
116 
117  rewriter.replaceOp(shapeCastOp, sourceShapeCastOp.getSource());
118  return success();
119  }
120 };
121 
122 /// Convert MulIOp/MulFOp + MultiDimReductionOp<add> into ContractionOp.
123 /// Ex:
124 /// ```
125 /// %0 = arith.mulf %arg0, %arg1 : vector<8x32x16xf32>
126 /// %1 = vector.multi_reduction add, %0 [1]
127 /// : vector<8x32x16xf32> to vector<8x16xf32>
128 /// ```
129 /// Gets converted to:
130 /// ```
131 /// %1 = vector.contract {indexing_maps = [
132 /// affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
133 /// affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
134 /// affine_map<(d0, d1, d2) -> (d0, d1)>],
135 /// iterator_types = ["parallel", "parallel", "reduction"],
136 /// kind = add} %0, %arg1, %cst_f0
137 /// : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
138 /// ```
139 struct MultiReduceToContract
140  : public OpRewritePattern<vector::MultiDimReductionOp> {
142 
143  LogicalResult matchAndRewrite(vector::MultiDimReductionOp reduceOp,
144  PatternRewriter &rewriter) const override {
145  if (reduceOp.getKind() != vector::CombiningKind::ADD)
146  return failure();
147  Operation *mulOp = reduceOp.getSource().getDefiningOp();
148  if (!mulOp || !isa<arith::MulIOp, arith::MulFOp>(mulOp))
149  return failure();
150  SmallVector<bool> reductionMask = reduceOp.getReductionMask();
151  auto srcMap = rewriter.getMultiDimIdentityMap(reductionMask.size());
153  SmallVector<vector::IteratorType> iteratorTypes;
154  for (const auto &isReduceDim : llvm::enumerate(reductionMask)) {
155  if (!isReduceDim.value()) {
156  iteratorTypes.push_back(vector::IteratorType::parallel);
157  exprs.push_back(rewriter.getAffineDimExpr(isReduceDim.index()));
158  } else {
159  iteratorTypes.push_back(vector::IteratorType::reduction);
160  }
161  }
162  auto dstMap =
163  AffineMap::get(/*dimCount=*/reductionMask.size(),
164  /*symbolCount=*/0, exprs, reduceOp.getContext());
165  rewriter.replaceOpWithNewOp<mlir::vector::ContractionOp>(
166  reduceOp, mulOp->getOperand(0), mulOp->getOperand(1), reduceOp.getAcc(),
167  rewriter.getAffineMapArrayAttr({srcMap, srcMap, dstMap}),
168  rewriter.getArrayAttr(llvm::to_vector(llvm::map_range(
169  iteratorTypes, [&](IteratorType t) -> mlir::Attribute {
170  return IteratorTypeAttr::get(rewriter.getContext(), t);
171  }))));
172  return success();
173  }
174 };
175 
176 /// Merge LHS/RHS (A/B) TransposeOp into ContractionOp user.
177 /// Ex:
178 /// ```
179 /// %0 = vector.transpose %arg0, [2, 0, 1]
180 /// : vector<32x16x8xf32> to vector<8x32x16xf32>
181 /// %1 = vector.contract {indexing_maps = [
182 /// affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
183 /// affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
184 /// affine_map<(d0, d1, d2) -> (d0, d1)>],
185 /// iterator_types = ["parallel", "parallel", "reduction"],
186 /// kind = add} %0, %arg1, %cst_f0
187 /// : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
188 /// ```
189 /// Gets converted to:
190 /// ```
191 /// %1 = vector.contract {indexing_maps = [
192 /// affine_map<(d0, d1, d2) -> (d1, d2, d0)>,
193 /// affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
194 /// affine_map<(d0, d1, d2) -> (d0, d1)>],
195 /// iterator_types = ["parallel", "parallel", "reduction"],
196 /// kind = add} %arg0, %arg1, %cst_f0
197 /// : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
198 /// ```
199 struct CombineContractABTranspose final
200  : public OpRewritePattern<vector::ContractionOp> {
202 
203  LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
204  PatternRewriter &rewriter) const override {
206  llvm::to_vector<4>(contractOp.getIndexingMapsArray());
207  Value lhs = contractOp.getLhs();
208  Value rhs = contractOp.getRhs();
209  size_t index = 0;
210  bool changed = false;
211  for (Value *operand : {&lhs, &rhs}) {
212  AffineMap &map = maps[index++];
213  auto transposeOp = operand->getDefiningOp<vector::TransposeOp>();
214  if (!transposeOp)
215  continue;
216  AffineMap permutationMap = AffineMap::getPermutationMap(
217  transposeOp.getPermutation(), contractOp.getContext());
218  map = inversePermutation(permutationMap).compose(map);
219  *operand = transposeOp.getVector();
220  changed = true;
221  }
222  if (!changed)
223  return failure();
224  rewriter.replaceOpWithNewOp<vector::ContractionOp>(
225  contractOp, lhs, rhs, contractOp.getAcc(),
226  rewriter.getAffineMapArrayAttr(maps), contractOp.getIteratorTypes());
227  return success();
228  }
229 };
230 
231 /// Merges accumulator and result transposes into contract.
232 ///
233 /// For example:
234 /// ```mlir
235 /// %accT = vector.transpose %acc, [0, 2, 1]
236 /// : vector<2x8x4xf32> to vector<2x4x8xf32>
237 /// %contract = vector.contract {
238 /// indexing_maps = [
239 /// affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
240 /// affine_map<(d0, d1, d2, d3) -> (d3, d2)>,
241 /// affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
242 /// ],
243 /// iterator_types = ["parallel", "parallel", "parallel", "reduction"],
244 /// kind = #vector.kind<add>
245 /// } %lhs, %rhs, %accT
246 /// : vector<2x4x4xf32>, vector<4x8xf32> into vector<2x4x8xf32>
247 /// %0 = vector.transpose %contract, [0, 2, 1]
248 /// : vector<2x4x8xf32> to vector<2x8x4>
249 /// ```
250 /// Becomes:
251 /// ```mlir
252 /// %0 = vector.contract {
253 /// indexing_maps = [
254 /// affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,
255 /// affine_map<(d0, d1, d2, d3) -> (d3, d2)>,
256 /// affine_map<(d0, d1, d2, d3) -> (d0, d2, d1)>
257 /// ],
258 /// iterator_types = ["parallel", "parallel", "parallel", "reduction"],
259 /// kind = #vector.kind<add>
260 /// } %lhs, %rhs, %acc
261 /// : vector<2x4x4xf32>, vector<4x8xf32> into vector<2x8x4xf32>
262 /// ```
263 struct CombineContractResultTranspose final
264  : public OpRewritePattern<vector::TransposeOp> {
266 
267  LogicalResult matchAndRewrite(vector::TransposeOp resTOp,
268  PatternRewriter &rewriter) const override {
269  auto contractOp = resTOp.getVector().getDefiningOp<vector::ContractionOp>();
270  if (!contractOp || !contractOp->hasOneUse())
271  return failure();
272 
273  auto accTOp = contractOp.getAcc().getDefiningOp<vector::TransposeOp>();
274  if (!accTOp)
275  return failure();
276 
277  MLIRContext *context = contractOp.getContext();
278  auto maps = llvm::to_vector<3>(contractOp.getIndexingMapsArray());
279  AffineMap contractMap = maps.back();
280 
281  // Accumulator transpose performs f(A) -> B. Contract performs g(C) -> B.
282  // To index into A in contract, we need revert(f)(g(C)) -> A.
283  auto accTMap =
284  AffineMap::getPermutationMap(accTOp.getPermutation(), context);
285 
286  // Contract performs g(C) -> D. Result transpose performs h(D) -> E.
287  // To index into E in contract, we need h(g(C)) -> E.
288  auto resTMap =
289  AffineMap::getPermutationMap(resTOp.getPermutation(), context);
290  auto combinedResMap = resTMap.compose(contractMap);
291 
292  // The accumulator and result share the same indexing map. So they should be
293  // the same to be able to merge. This means combinedResMap is the same as
294  // inversePermutation(accTMap).compose(contractMap), which means
295  if (inversePermutation(accTMap) != resTMap)
296  return failure();
297  maps.back() = combinedResMap;
298 
299  rewriter.replaceOpWithNewOp<vector::ContractionOp>(
300  resTOp, contractOp.getLhs(), contractOp.getRhs(), accTOp.getVector(),
301  rewriter.getAffineMapArrayAttr(maps), contractOp.getIteratorTypes());
302  return success();
303  }
304 };
305 
306 /// Merge BroadcastOp into ContractionOp user.
307 /// Ex:
308 /// ```
309 /// %0 = vector.broadcast %arg0 : vector<32x16xf32> to vector<8x32x16xf32>
310 /// %1 = vector.contract {indexing_maps = [
311 /// affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
312 /// affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
313 /// affine_map<(d0, d1, d2) -> (d0, d1)>],
314 /// iterator_types = ["parallel", "parallel", "reduction"],
315 /// kind = add} %0, %arg1, %cst_f0
316 /// : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
317 /// ```
318 /// Gets converted to:
319 /// ```
320 /// %1 = vector.contract {indexing_maps = [
321 /// affine_map<(d0, d1, d2) -> (d1, d2)>,
322 /// affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
323 /// affine_map<(d0, d1, d2) -> (d0, d1)>],
324 /// iterator_types = ["parallel", "parallel", "reduction"],
325 /// kind = add} %arg0, %arg1, %cst_f0
326 /// : vector<32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>
327 /// ```
328 struct CombineContractBroadcast
329  : public OpRewritePattern<vector::ContractionOp> {
331 
332  LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
333  PatternRewriter &rewriter) const override {
335  llvm::to_vector<4>(contractOp.getIndexingMapsArray());
336  Value lhs = contractOp.getLhs();
337  Value rhs = contractOp.getRhs();
338  size_t index = 0;
339  bool changed = false;
340  for (Value *operand : {&lhs, &rhs}) {
341  AffineMap &map = maps[index++];
342  auto broadcast = operand->getDefiningOp<vector::BroadcastOp>();
343  if (!broadcast)
344  continue;
345  // contractionOp can only take vector as operands.
346  auto srcType = dyn_cast<VectorType>(broadcast.getSourceType());
347  if (!srcType ||
348  srcType.getRank() == broadcast.getResultVectorType().getRank())
349  continue;
350  int64_t rankDiff =
351  broadcast.getResultVectorType().getRank() - srcType.getRank();
352  bool innerDimBroadcast = false;
353  SmallVector<AffineExpr> originalDims;
354  for (const auto &dim : llvm::enumerate(srcType.getShape())) {
355  if (dim.value() != broadcast.getResultVectorType().getDimSize(
356  rankDiff + dim.index())) {
357  innerDimBroadcast = true;
358  break;
359  }
360  originalDims.push_back(
361  rewriter.getAffineDimExpr(dim.index() + rankDiff));
362  }
363  // Contract doesn't support inner dimension broadcast. Once this is
364  // relaxed we can remove this case.
365  if (innerDimBroadcast)
366  continue;
367 
368  // It would be incorrect to fold a broadcast onto a reduction dimension
369  // of non-unit size.
370  bool nonUnitDimReductionBroadcast = false;
371  for (int64_t i = 0; i < rankDiff; ++i) {
372  if (broadcast.getResultVectorType().getDimSize(i) != 1 &&
373  isReductionIterator(contractOp.getIteratorTypes()
374  .getValue()[map.getDimPosition(i)])) {
375  nonUnitDimReductionBroadcast = true;
376  break;
377  }
378  }
379  if (nonUnitDimReductionBroadcast)
380  continue;
381 
382  AffineMap broadcastMap =
383  AffineMap::get(broadcast.getResultVectorType().getRank(), 0,
384  originalDims, contractOp.getContext());
385  map = broadcastMap.compose(map);
386  *operand = broadcast.getSource();
387  changed = true;
388  }
389 
390  if (!changed)
391  return failure();
392 
393  // Determine which dims are usused, now that the maps have been composed
394  // with the broadcast maps.
395  llvm::SmallBitVector unusedDimsBitVector = getUnusedDimsBitVector(maps);
396  // Compress unused dims.
397  for (auto &m : maps)
398  m = compressDims(m, unusedDimsBitVector);
399  // Compute the combined iterators.
400  SmallVector<Attribute> iterators;
401  for (unsigned i = 0; i < unusedDimsBitVector.size(); ++i) {
402  if (!unusedDimsBitVector.test(i))
403  iterators.push_back(contractOp.getIteratorTypes().getValue()[i]);
404  }
405  // Check that compressing unused dims isn't removing all reduction dimension
406  // pairs. For example, if the vector.contract had only one reduction
407  // iterator and that was a unit-dimension created by a broadcast,
408  // then we should bail here, otherwise we would create a contract without
409  // a reduction dimension pair.
410  bool hasReductionIteratorApplyingOnBothSides = false;
411  for (unsigned i = 0; i < iterators.size(); ++i) {
412  if (!isReductionIterator(iterators[i]))
413  continue;
414  if (getResultIndex(maps[0], i) && getResultIndex(maps[1], i)) {
415  hasReductionIteratorApplyingOnBothSides = true;
416  break;
417  }
418  }
419  if (!hasReductionIteratorApplyingOnBothSides)
420  return failure();
421 
422  // If the compressed maps have a dimension that is not used by either LHS or
423  // RHS then the ContractionOp verifier would fail.
424  if (getUnusedDimsBitVector({maps[0], maps[1]}).any())
425  return failure();
426  rewriter.replaceOpWithNewOp<vector::ContractionOp>(
427  contractOp, lhs, rhs, contractOp.getAcc(),
428  rewriter.getAffineMapArrayAttr(maps), rewriter.getArrayAttr(iterators));
429  return success();
430  }
431 };
432 
433 /// Reorders cast(broadcast) to broadcast(cast). This makes broadcast ops and
434 /// contraction ops closer, which kicks in CombineContractBroadcast pattern when
435 /// casting ops are around these operations.
436 /// Ex:
437 /// ```
438 /// %0 = vector.broadcast %arg0 : vector<32x16xi8> to vector<8x32x16xi8>
439 /// %1 = arith.extsi %0 : vector<8x32x16xi8> to vector<8x32x16xi32>
440 /// ```
441 /// Gets converted to:
442 /// ```
443 /// %0 = arith.extsi %0 : vector<32x16xi8> to vector<32x16xi32>
444 /// %1 = vector.broadcast %arg0 : vector<32x16xi32> to vector<8x32x16xi32>
445 /// ```
446 struct ReorderCastOpsOnBroadcast
447  : public OpInterfaceRewritePattern<CastOpInterface> {
449 
450  LogicalResult matchAndRewrite(CastOpInterface op,
451  PatternRewriter &rewriter) const override {
452  if (op->getNumOperands() != 1)
453  return failure();
454  auto bcastOp = op->getOperand(0).getDefiningOp<vector::BroadcastOp>();
455  if (!bcastOp)
456  return failure();
457 
458  Type castResTy = getElementTypeOrSelf(op->getResult(0));
459  if (auto vecTy = dyn_cast<VectorType>(bcastOp.getSourceType()))
460  castResTy = vecTy.clone(castResTy);
461  auto *castOp =
462  rewriter.create(op->getLoc(), op->getName().getIdentifier(),
463  bcastOp.getSource(), castResTy, op->getAttrs());
464  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
465  op, op->getResult(0).getType(), castOp->getResult(0));
466  return success();
467  }
468 };
469 
470 /// Reorders elementwise(transpose) to transpose(elementwise). This makes
471 /// transpose ops and contraction ops closer, which kicks in
472 /// CombineContractABTranspose pattern when elementwise ops are between these
473 /// operations. Ex:
474 /// ```
475 /// %at = vector.transpose %a, [1, 0]: vector<4x2xf32> to vector<2x4xf32>
476 /// %bt = vector.transpose %b, [1, 0]: vector<4x2xf32> to vector<2x4xf32>
477 /// %r = arith.addf %at, %bt : vector<2x4xf32>
478 /// ```
479 /// Gets converted to:
480 /// ```
481 /// %0 = arith.addf %a, %b : vector<4x2xf32>
482 /// %r = vector.transpose %0, [1, 0] : vector<2x4xf32>
483 /// ```
484 struct ReorderElementwiseOpsOnTranspose final
485  : public OpTraitRewritePattern<OpTrait::Elementwise> {
487  LogicalResult matchAndRewrite(Operation *op,
488  PatternRewriter &rewriter) const override {
489  if (op->getNumResults() != 1 || op->getNumRegions() != 0)
490  return failure();
491 
492  // Make sure all operands are transpose/constant ops and collect their
493  // transposition maps.
494  SmallVector<ArrayRef<int64_t>> transposeMaps;
495  transposeMaps.reserve(op->getNumOperands());
496  // Record the initial type before transposition. We'll use its shape later.
497  // Any type will do here as we will check all transpose maps are the same.
498  VectorType srcType;
499  for (Value operand : op->getOperands()) {
500  auto transposeOp = operand.getDefiningOp<vector::TransposeOp>();
501  if (transposeOp) {
502  transposeMaps.push_back(transposeOp.getPermutation());
503  srcType = transposeOp.getSourceVectorType();
504  } else if (!matchPattern(operand, m_Constant())) {
505  return failure();
506  }
507  }
508  if (transposeMaps.empty())
509  return failure();
510  // This is an elementwise op, so all transposed operands should have the
511  // same type. We need to additionally check that all transposes uses the
512  // same map.
513  if (!llvm::all_equal(transposeMaps))
514  return rewriter.notifyMatchFailure(op, "different transpose map");
515 
516  SmallVector<Value> srcValues;
517  srcValues.reserve(op->getNumOperands());
518 
519  // If there are constant operands, we need to insert inverse transposes for
520  // them. Calculate the inverse order first.
521  auto order = transposeMaps.front();
522  SmallVector<int64_t> invOrder(order.size());
523  for (int i = 0, e = order.size(); i < e; ++i)
524  invOrder[order[i]] = i;
525 
526  for (Value operand : op->getOperands()) {
527  auto transposeOp = operand.getDefiningOp<vector::TransposeOp>();
528  if (transposeOp) {
529  srcValues.push_back(transposeOp.getVector());
530  } else {
531  // This is a constant. Create a reverse transpose op for it.
532  auto vectorType =
533  srcType.clone(cast<VectorType>(operand.getType()).getElementType());
534  srcValues.push_back(rewriter.create<vector::TransposeOp>(
535  operand.getLoc(), vectorType, operand, invOrder));
536  }
537  }
538 
539  auto vectorType = srcType.clone(
540  cast<VectorType>(op->getResultTypes()[0]).getElementType());
541  Operation *elementwiseOp =
542  rewriter.create(op->getLoc(), op->getName().getIdentifier(), srcValues,
543  vectorType, op->getAttrs());
544  rewriter.replaceOpWithNewOp<vector::TransposeOp>(
545  op, op->getResultTypes()[0], elementwiseOp->getResult(0),
546  transposeMaps.front());
547  return success();
548  }
549 };
550 
551 // Returns the values in `arrayAttr` as an integer vector.
552 static SmallVector<int64_t> getIntValueVector(ArrayAttr arrayAttr) {
553  return llvm::to_vector<4>(
554  llvm::map_range(arrayAttr.getAsRange<IntegerAttr>(),
555  [](IntegerAttr attr) { return attr.getInt(); }));
556 }
557 
558 // Shuffles vector.bitcast op after vector.extract op.
559 //
560 // This transforms IR like:
561 // %0 = vector.bitcast %src : vector<4xf32> to vector<8xf16>
562 // %1 = vector.extract %0[3] : f16 from vector<8xf16>
563 // Into:
564 // %0 = vector.extract %src[1] : f32 from vector<4xf32>
565 // %1 = vector.bitcast %0: vector<1xf32> to vector<2xf16>
566 // %2 = vector.extract %1[1] : f16 from vector<2xf16>
567 struct BubbleDownVectorBitCastForExtract
568  : public OpRewritePattern<vector::ExtractOp> {
570 
571  LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
572  PatternRewriter &rewriter) const override {
573  // Only support extracting scalars for now.
574  if (extractOp.getSourceVectorType().getRank() != 1)
575  return failure();
576 
577  auto castOp = extractOp.getVector().getDefiningOp<vector::BitCastOp>();
578  if (!castOp)
579  return failure();
580 
581  VectorType castSrcType = castOp.getSourceVectorType();
582  VectorType castDstType = castOp.getResultVectorType();
583  assert(castSrcType.getRank() == castDstType.getRank());
584 
585  // Fail to match if we only have one element in the cast op source.
586  // This is to avoid infinite loop given that this pattern can generate
587  // such cases.
588  if (castSrcType.getNumElements() == 1)
589  return failure();
590 
591  // Only support casting to a larger number of elements or now.
592  // E.g., vector<4xf32> -> vector<8xf16>.
593  if (castSrcType.getNumElements() > castDstType.getNumElements())
594  return failure();
595 
596  unsigned expandRatio =
597  castDstType.getNumElements() / castSrcType.getNumElements();
598 
599  // Get the first element of the mixed position as integer.
600  auto mixedPos = extractOp.getMixedPosition();
601  if (mixedPos.size() > 0 && !isa<Attribute>(mixedPos[0]))
602  return failure();
603  uint64_t index = cast<IntegerAttr>(cast<Attribute>(mixedPos[0])).getInt();
604 
605  // Get the single scalar (as a vector) in the source value that packs the
606  // desired scalar. E.g. extract vector<1xf32> from vector<4xf32>
607  Location loc = extractOp.getLoc();
608  Value packedValue = rewriter.create<vector::ExtractOp>(
609  loc, castOp.getSource(), index / expandRatio);
610  Type packedVecType = VectorType::get(/*shape=*/{1}, packedValue.getType());
611  Value zero = rewriter.create<arith::ConstantOp>(
612  loc, packedVecType, rewriter.getZeroAttr(packedVecType));
613  packedValue = rewriter.create<vector::InsertOp>(loc, packedValue, zero,
614  /*position=*/0);
615 
616  // Cast it to a vector with the desired scalar's type.
617  // E.g. f32 -> vector<2xf16>
618  VectorType packedType =
619  VectorType::get({expandRatio}, castDstType.getElementType());
620  Value castedValue =
621  rewriter.create<vector::BitCastOp>(loc, packedType, packedValue);
622 
623  // Finally extract the desired scalar.
624  rewriter.replaceOpWithNewOp<vector::ExtractOp>(extractOp, castedValue,
625  index % expandRatio);
626  return success();
627  }
628 };
629 
630 // Shuffles vector.bitcast op after vector.extract_strided_slice op.
631 //
632 // This transforms IR like:
633 // %cast = vector.bitcast %arg0: vector<4xf32> to vector<8xf16>
634 // %0 = vector.extract_strided_slice %cast {
635 // offsets = [4], sizes = [4], strides = [1]
636 // } : vector<8xf16> to vector<4xf16>
637 // Into:
638 // %0 = vector.extract_strided_slice %src {
639 // offsets = [2], sizes = [2], strides = [1]
640 // } : vector<4xf32> to vector<2xf32>
641 // %1 = vector.bitcast %0 : vector<2xf32> to vector<4xf16>
642 struct BubbleDownBitCastForStridedSliceExtract
643  : public OpRewritePattern<vector::ExtractStridedSliceOp> {
645 
646  LogicalResult matchAndRewrite(vector::ExtractStridedSliceOp extractOp,
647  PatternRewriter &rewriter) const override {
648  auto castOp = extractOp.getVector().getDefiningOp<vector::BitCastOp>();
649  if (!castOp)
650  return failure();
651 
652  VectorType castSrcType = castOp.getSourceVectorType();
653  VectorType castDstType = castOp.getResultVectorType();
654  assert(castSrcType.getRank() == castDstType.getRank());
655 
656  int64_t castSrcLastDim = castSrcType.getShape().back();
657  int64_t castDstLastDim = castDstType.getShape().back();
658  // Require casting to more elements for now; other cases to be implemented.
659  if (castSrcLastDim > castDstLastDim)
660  return failure();
661 
662  // Only accept all one strides for now.
663  if (llvm::any_of(extractOp.getStrides().getAsValueRange<IntegerAttr>(),
664  [](const APInt &val) { return !val.isOne(); }))
665  return failure();
666 
667  unsigned rank = extractOp.getSourceVectorType().getRank();
668  assert(castDstLastDim % castSrcLastDim == 0);
669  int64_t expandRatio = castDstLastDim / castSrcLastDim;
670 
671  // If we have a less number of offsets than the rank, then implicitly we
672  // are selecting the full range for the last bitcasted dimension; other
673  // dimensions aren't affected. Otherwise, we need to scale down the last
674  // dimension's offset given we are extracting from less elements now.
675  ArrayAttr newOffsets = extractOp.getOffsets();
676  if (newOffsets.size() == rank) {
677  SmallVector<int64_t> offsets = getIntValueVector(newOffsets);
678  if (offsets.back() % expandRatio != 0)
679  return failure();
680  offsets.back() = offsets.back() / expandRatio;
681  newOffsets = rewriter.getI64ArrayAttr(offsets);
682  }
683 
684  // Similarly for sizes.
685  ArrayAttr newSizes = extractOp.getSizes();
686  if (newSizes.size() == rank) {
687  SmallVector<int64_t> sizes = getIntValueVector(newSizes);
688  if (sizes.back() % expandRatio != 0)
689  return failure();
690  sizes.back() = sizes.back() / expandRatio;
691  newSizes = rewriter.getI64ArrayAttr(sizes);
692  }
693 
694  SmallVector<int64_t> dims =
695  llvm::to_vector<4>(cast<VectorType>(extractOp.getType()).getShape());
696  dims.back() = dims.back() / expandRatio;
697  VectorType newExtractType =
698  VectorType::get(dims, castSrcType.getElementType());
699 
700  auto newExtractOp = rewriter.create<vector::ExtractStridedSliceOp>(
701  extractOp.getLoc(), newExtractType, castOp.getSource(), newOffsets,
702  newSizes, extractOp.getStrides());
703 
704  rewriter.replaceOpWithNewOp<vector::BitCastOp>(
705  extractOp, extractOp.getType(), newExtractOp);
706 
707  return success();
708  }
709 };
710 
711 // Shuffles vector.bitcast op before vector.insert_strided_slice op.
712 //
713 // This transforms IR like:
714 // %0 = vector.insert %val, %dst[4] : vector<32xi4> into vector<8x32xi4>
715 // %1 = vector.bitcast %0 : vector<8x32xi4> to vector<8x16xi8>
716 // Into:
717 // %0 = vector.bitcast %val : vector<32xi4> to vector<16xi8>
718 // %1 = vector.bitcast %dst : vector<8x32xi4> to vector<8x16xi8>
719 // %2 = vector.insert %0, %1 [4] : vector<16xi8> into vector<8x16xi8>
720 //
721 struct BubbleUpBitCastForInsert : public OpRewritePattern<vector::BitCastOp> {
723 
724  LogicalResult matchAndRewrite(vector::BitCastOp bitcastOp,
725  PatternRewriter &rewriter) const override {
726  VectorType castSrcType = bitcastOp.getSourceVectorType();
727  VectorType castDstType = bitcastOp.getResultVectorType();
728 
729  // 0-D and scalable vectors are not supported yet.
730  if (castSrcType.getRank() == 0 || castSrcType.isScalable() ||
731  castDstType.isScalable())
732  return failure();
733 
734  int64_t castSrcLastDim = castSrcType.getShape().back();
735  int64_t castDstLastDim = castDstType.getShape().back();
736  bool isNumElemsShrink = castSrcLastDim >= castDstLastDim;
737  int64_t ratio;
738  if (isNumElemsShrink) {
739  assert(castSrcLastDim % castDstLastDim == 0);
740  ratio = castSrcLastDim / castDstLastDim;
741  } else {
742  assert(castDstLastDim % castSrcLastDim == 0);
743  ratio = castDstLastDim / castSrcLastDim;
744  }
745 
746  auto insertOp = bitcastOp.getSource().getDefiningOp<vector::InsertOp>();
747  if (!insertOp)
748  return failure();
749 
750  // Only vector sources are supported for now.
751  auto insertSrcType = dyn_cast<VectorType>(insertOp.getSourceType());
752  if (!insertSrcType)
753  return failure();
754 
755  // Bitcast the source.
756  SmallVector<int64_t> srcDims(insertSrcType.getShape());
757  srcDims.back() =
758  isNumElemsShrink ? srcDims.back() / ratio : srcDims.back() * ratio;
759  VectorType newCastSrcType =
760  VectorType::get(srcDims, castDstType.getElementType());
761  auto newCastSrcOp = rewriter.create<vector::BitCastOp>(
762  bitcastOp.getLoc(), newCastSrcType, insertOp.getSource());
763 
764  SmallVector<int64_t> dstDims(insertOp.getDestVectorType().getShape());
765  dstDims.back() =
766  isNumElemsShrink ? dstDims.back() / ratio : dstDims.back() * ratio;
767  VectorType newCastDstType =
768  VectorType::get(dstDims, castDstType.getElementType());
769 
770  // Bitcast the destination.
771  auto newCastDstOp = rewriter.create<vector::BitCastOp>(
772  bitcastOp.getLoc(), newCastDstType, insertOp.getDest());
773 
774  // Generate new insert.
775  rewriter.replaceOpWithNewOp<vector::InsertOp>(
776  bitcastOp, newCastSrcOp, newCastDstOp, insertOp.getMixedPosition());
777  return success();
778  }
779 };
780 
781 // Shuffles vector.bitcast op before vector.insert_strided_slice op.
782 //
783 // This transforms IR like:
784 // %0 = vector.insert_strided_slice %src, %dst {
785 // offsets = [0], strides = [1]} : vector<4xf16> into vector<8xf16>
786 // %1 = vector.bitcast %0: vector<8xf16> to vector<4xf32>
787 // Into:
788 // %0 = vector.bitcast %src : vector<4xf16> to vector<2xf32>
789 // %1 = vector.bitcast %dst : vector<8xf16> to vector<4xf32>
790 // %2 = vector.insert_strided_slice %src, %dst {
791 // offsets = [0], strides = [1]} : vector<2xf32> into vector<4xf32>
792 struct BubbleUpBitCastForStridedSliceInsert
793  : public OpRewritePattern<vector::BitCastOp> {
795 
796  LogicalResult matchAndRewrite(vector::BitCastOp bitcastOp,
797  PatternRewriter &rewriter) const override {
798  VectorType castSrcType = bitcastOp.getSourceVectorType();
799  VectorType castDstType = bitcastOp.getResultVectorType();
800  assert(castSrcType.getRank() == castDstType.getRank());
801  // Skip 0-D vector which will not from InsertStridedSliceOp.
802  if (castSrcType.getRank() == 0)
803  return failure();
804 
805  int64_t castSrcLastDim = castSrcType.getShape().back();
806  int64_t castDstLastDim = castDstType.getShape().back();
807  // Require casting to less elements for now; other cases to be implemented.
808  if (castSrcLastDim < castDstLastDim)
809  return failure();
810 
811  assert(castSrcLastDim % castDstLastDim == 0);
812  int64_t shrinkRatio = castSrcLastDim / castDstLastDim;
813 
814  auto insertOp =
815  bitcastOp.getSource().getDefiningOp<vector::InsertStridedSliceOp>();
816  if (!insertOp)
817  return failure();
818 
819  // Only accept all one strides for now.
820  if (llvm::any_of(insertOp.getStrides().getAsValueRange<IntegerAttr>(),
821  [](const APInt &val) { return !val.isOne(); }))
822  return failure();
823 
824  unsigned rank = insertOp.getSourceVectorType().getRank();
825  // Require insert op to have the same rank for the source and destination
826  // vector; other cases to be implemented.
827  if (rank != insertOp.getDestVectorType().getRank())
828  return failure();
829 
830  // Requires that shape of insert op src is castable to dstType.
831  unsigned sourceWidth = castSrcType.getElementType().getIntOrFloatBitWidth();
832  unsigned destinationWidth =
833  castDstType.getElementType().getIntOrFloatBitWidth();
834  unsigned numElements = destinationWidth / sourceWidth;
835  if (insertOp.getSourceVectorType().getNumElements() % numElements != 0)
836  return failure();
837 
838  ArrayAttr newOffsets = insertOp.getOffsets();
839  assert(newOffsets.size() == rank);
840  SmallVector<int64_t> offsets = getIntValueVector(newOffsets);
841  if (offsets.back() % shrinkRatio != 0)
842  return failure();
843  offsets.back() = offsets.back() / shrinkRatio;
844  newOffsets = rewriter.getI64ArrayAttr(offsets);
845 
846  SmallVector<int64_t> srcDims =
847  llvm::to_vector<4>(insertOp.getSourceVectorType().getShape());
848  srcDims.back() = srcDims.back() / shrinkRatio;
849  VectorType newCastSrcType =
850  VectorType::get(srcDims, castDstType.getElementType());
851 
852  auto newCastSrcOp = rewriter.create<vector::BitCastOp>(
853  bitcastOp.getLoc(), newCastSrcType, insertOp.getSource());
854 
855  SmallVector<int64_t> dstDims =
856  llvm::to_vector<4>(insertOp.getDestVectorType().getShape());
857  dstDims.back() = dstDims.back() / shrinkRatio;
858  VectorType newCastDstType =
859  VectorType::get(dstDims, castDstType.getElementType());
860 
861  auto newCastDstOp = rewriter.create<vector::BitCastOp>(
862  bitcastOp.getLoc(), newCastDstType, insertOp.getDest());
863 
864  rewriter.replaceOpWithNewOp<vector::InsertStridedSliceOp>(
865  bitcastOp, bitcastOp.getType(), newCastSrcOp, newCastDstOp, newOffsets,
866  insertOp.getStrides());
867 
868  return success();
869  }
870 };
871 
872 // Breaks down vector.bitcast op
873 //
874 // This transforms IR like:
875 // %1 = vector.bitcast %0: vector<8xf16> to vector<4xf32>
876 // Into:
877 // %cst = vector.splat %c0_f32 : vector<4xf32>
878 // %1 = vector.extract_strided_slice %0 {
879 // offsets = [0], sizes = [4], strides = [1]
880 // } : vector<8xf16> to vector<4xf16>
881 // %2 = vector.bitcast %1 : vector<4xf16> to vector<2xf32>
882 // %4 = vector.insert_strided_slice %2, %cst {
883 // offsets = [0], strides = [1]} : vector<2xf32> into vector<4xf32>
884 // %5 = vector.extract_strided_slice %0 {
885 // offsets = [4], sizes = [4], strides = [1]
886 // } : vector<8xf16> to vector<4xf16>
887 // %6 = vector.bitcast %5 : vector<4xf16> to vector<2xf32>
888 // %7 = vector.insert_strided_slice %6, %cst {
889 // offsets = [2], strides = [1]} : vector<2xf32> into vector<4xf32>
890 struct BreakDownVectorBitCast : public OpRewritePattern<vector::BitCastOp> {
892 
893 public:
894  BreakDownVectorBitCast(MLIRContext *context,
895  std::function<bool(vector::BitCastOp)> controlFn,
896  PatternBenefit benefit)
897  : OpRewritePattern(context, benefit), controlFn(std::move(controlFn)) {}
898 
899  LogicalResult matchAndRewrite(vector::BitCastOp bitcastOp,
900  PatternRewriter &rewriter) const override {
901 
902  if (controlFn && !controlFn(bitcastOp))
903  return failure();
904 
905  VectorType castSrcType = bitcastOp.getSourceVectorType();
906  VectorType castDstType = bitcastOp.getResultVectorType();
907  assert(castSrcType.getRank() == castDstType.getRank());
908 
909  // This transformation builds on top of
910  // vector.{extract|insert}_strided_slice, which do not support
911  // extracting/inserting "scallable sub-vectors". Bail out.
912  if (castSrcType.isScalable())
913  return rewriter.notifyMatchFailure(bitcastOp,
914  "Scalable vectors are not supported");
915 
916  // Only support rank 1 case for now.
917  if (castSrcType.getRank() != 1)
918  return failure();
919 
920  int64_t castSrcLastDim = castSrcType.getShape().back();
921  int64_t castDstLastDim = castDstType.getShape().back();
922  // Require casting to less elements for now; other cases to be implemented.
923  if (castSrcLastDim < castDstLastDim)
924  return failure();
925 
926  assert(castSrcLastDim % castDstLastDim == 0);
927  int64_t shrinkRatio = castSrcLastDim / castDstLastDim;
928  // Nothing to do if it is already bitcasting to a single element.
929  if (castSrcLastDim == shrinkRatio)
930  return failure();
931 
932  Location loc = bitcastOp.getLoc();
933  Type elemType = castDstType.getElementType();
934  assert(elemType.isSignlessIntOrIndexOrFloat());
935 
936  Value zero = rewriter.create<arith::ConstantOp>(
937  loc, elemType, rewriter.getZeroAttr(elemType));
938  Value res = rewriter.create<SplatOp>(loc, castDstType, zero);
939 
940  SmallVector<int64_t> sliceShape = {castDstLastDim};
941  SmallVector<int64_t> strides = {1};
942  VectorType newCastDstType =
943  VectorType::get(SmallVector<int64_t>{castDstLastDim / shrinkRatio},
944  castDstType.getElementType());
945 
946  for (int i = 0, e = shrinkRatio; i < e; ++i) {
947  Value extracted = rewriter.create<ExtractStridedSliceOp>(
948  loc, bitcastOp.getSource(), ArrayRef<int64_t>{i * castDstLastDim},
949  sliceShape, strides);
950  Value bitcast =
951  rewriter.create<BitCastOp>(loc, newCastDstType, extracted);
952  res = rewriter.create<InsertStridedSliceOp>(
953  loc, bitcast, res,
954  ArrayRef<int64_t>{i * castDstLastDim / shrinkRatio}, strides);
955  }
956  rewriter.replaceOp(bitcastOp, res);
957  return success();
958  }
959 
960 private:
961  std::function<bool(BitCastOp)> controlFn;
962 };
963 
964 /// Reorders elementwise(broadcast/splat) to broadcast(elementwise). Ex:
965 /// ```
966 /// %a = vector.broadcast %arg1 : index to vector<1x4xindex>
967 /// %b = vector.broadcast %arg2 : index to vector<1x4xindex>
968 /// %r = arith.addi %a, %b : vector<1x4xindex>
969 /// ```
970 /// Gets converted to:
971 /// ```
972 /// %r = arith.addi %arg0, %arg1 : index
973 /// %b = vector.broadcast %r : index to vector<1x4xindex>
974 /// ```
975 ///
976 /// Both `vector.broadcast` and `vector.splat` are supported as broadcasting
977 /// ops.
978 struct ReorderElementwiseOpsOnBroadcast final
979  : public OpTraitRewritePattern<OpTrait::Elementwise> {
981  LogicalResult matchAndRewrite(Operation *op,
982  PatternRewriter &rewriter) const override {
983  if (op->getNumResults() != 1)
984  return failure();
985  if (!llvm::isa<ShapedType>(op->getResults()[0].getType()))
986  return failure();
988  return rewriter.notifyMatchFailure(
989  op, "Op doesn't have ElementwiseMappableTraits");
990  if (op->getNumOperands() == 0)
991  return failure();
992  if (op->getResults()[0].getType() != op->getOperand(0).getType())
993  return rewriter.notifyMatchFailure(op,
994  "result and operand type mismatch");
995  if (isa<vector::FMAOp>(op)) {
996  return rewriter.notifyMatchFailure(
997  op,
998  "Op only accepts vector types - not supported as broadcast source "
999  "might be a scalar");
1000  }
1001 
1002  // Get the type of the lhs operand
1003  auto *lhsBcastOrSplat = op->getOperand(0).getDefiningOp();
1004  if (!lhsBcastOrSplat ||
1005  !isa<vector::BroadcastOp, vector::SplatOp>(*lhsBcastOrSplat))
1006  return failure();
1007  auto lhsBcastOrSplatType = lhsBcastOrSplat->getOperand(0).getType();
1008 
1009  // Make sure that all operands are broadcast from identical types:
1010  // * scalar (`vector.broadcast` + `vector.splat`), or
1011  // * vector (`vector.broadcast`).
1012  // Otherwise the re-ordering wouldn't be safe.
1013  if (!llvm::all_of(op->getOperands(), [&lhsBcastOrSplatType](Value val) {
1014  auto bcast = val.getDefiningOp<vector::BroadcastOp>();
1015  if (bcast)
1016  return (bcast.getOperand().getType() == lhsBcastOrSplatType);
1017  auto splat = val.getDefiningOp<vector::SplatOp>();
1018  if (splat)
1019  return (splat.getOperand().getType() == lhsBcastOrSplatType);
1020  return false;
1021  })) {
1022  return failure();
1023  }
1024 
1025  // Collect the source values before broadcasting
1026  SmallVector<Value> srcValues;
1027  srcValues.reserve(op->getNumOperands());
1028  for (Value operand : op->getOperands()) {
1029  srcValues.push_back(operand.getDefiningOp()->getOperand(0));
1030  }
1031 
1032  // Create the "elementwise" Op
1033  Operation *elementwiseOp =
1034  rewriter.create(op->getLoc(), op->getName().getIdentifier(), srcValues,
1035  lhsBcastOrSplatType, op->getAttrs());
1036 
1037  // Replace the original Op with the elementwise Op
1038  auto vectorType = op->getResultTypes()[0];
1039  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
1040  op, vectorType, elementwiseOp->getResults());
1041 
1042  return success();
1043  }
1044 };
1045 
1046 // Helper that returns a vector comparison that constructs a mask:
1047 // mask = [0,1,..,n-1] + [o,o,..,o] < [b,b,..,b]
1048 //
1049 // If `dim == 0` then the result will be a 0-D vector.
1050 //
1051 // NOTE: The LLVM::GetActiveLaneMaskOp intrinsic would provide an alternative,
1052 // much more compact, IR for this operation, but LLVM eventually
1053 // generates more elaborate instructions for this intrinsic since it
1054 // is very conservative on the boundary conditions.
1055 static Value buildVectorComparison(PatternRewriter &rewriter, Operation *op,
1056  bool force32BitVectorIndices, int64_t dim,
1057  Value b, Value *off = nullptr) {
1058  auto loc = op->getLoc();
1059  // If we can assume all indices fit in 32-bit, we perform the vector
1060  // comparison in 32-bit to get a higher degree of SIMD parallelism.
1061  // Otherwise we perform the vector comparison using 64-bit indices.
1062  Type idxType =
1063  force32BitVectorIndices ? rewriter.getI32Type() : rewriter.getI64Type();
1064  DenseIntElementsAttr indicesAttr;
1065  if (dim == 0 && force32BitVectorIndices) {
1066  indicesAttr = DenseIntElementsAttr::get(
1068  } else if (dim == 0) {
1069  indicesAttr = DenseIntElementsAttr::get(
1071  } else if (force32BitVectorIndices) {
1072  indicesAttr = rewriter.getI32VectorAttr(
1073  llvm::to_vector<4>(llvm::seq<int32_t>(0, dim)));
1074  } else {
1075  indicesAttr = rewriter.getI64VectorAttr(
1076  llvm::to_vector<4>(llvm::seq<int64_t>(0, dim)));
1077  }
1078  Value indices = rewriter.create<arith::ConstantOp>(loc, indicesAttr);
1079  // Add in an offset if requested.
1080  if (off) {
1081  Value o = getValueOrCreateCastToIndexLike(rewriter, loc, idxType, *off);
1082  Value ov = rewriter.create<vector::SplatOp>(loc, indices.getType(), o);
1083  indices = rewriter.create<arith::AddIOp>(loc, ov, indices);
1084  }
1085  // Construct the vector comparison.
1086  Value bound = getValueOrCreateCastToIndexLike(rewriter, loc, idxType, b);
1087  Value bounds =
1088  rewriter.create<vector::SplatOp>(loc, indices.getType(), bound);
1089  return rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, indices,
1090  bounds);
1091 }
1092 
1093 template <typename ConcreteOp>
1094 struct MaterializeTransferMask : public OpRewritePattern<ConcreteOp> {
1095 public:
1096  explicit MaterializeTransferMask(MLIRContext *context, bool enableIndexOpt,
1097  PatternBenefit benefit = 1)
1098  : mlir::OpRewritePattern<ConcreteOp>(context, benefit),
1099  force32BitVectorIndices(enableIndexOpt) {}
1100 
1101  LogicalResult matchAndRewrite(ConcreteOp xferOp,
1102  PatternRewriter &rewriter) const override {
1103  if (!xferOp.hasOutOfBoundsDim())
1104  return failure();
1105 
1106  if (xferOp.getVectorType().getRank() > 1 || xferOp.getIndices().empty())
1107  return failure();
1108 
1109  Location loc = xferOp->getLoc();
1110  VectorType vtp = xferOp.getVectorType();
1111 
1112  // Create the in-bounds mask with all elements between [0 .. dim - offset)
1113  // set and [dim - offset .. vector_length) unset.
1114  //
1115  // TODO: when the leaf transfer rank is k > 1, we need the last `k`
1116  // dimensions here.
1117  unsigned lastIndex = llvm::size(xferOp.getIndices()) - 1;
1118  Value off = xferOp.getIndices()[lastIndex];
1119  Value dim =
1120  vector::createOrFoldDimOp(rewriter, loc, xferOp.getSource(), lastIndex);
1121  Value b = rewriter.create<arith::SubIOp>(loc, dim.getType(), dim, off);
1122  Value mask = rewriter.create<vector::CreateMaskOp>(
1123  loc,
1124  VectorType::get(vtp.getShape(), rewriter.getI1Type(),
1125  vtp.getScalableDims()),
1126  b);
1127  if (xferOp.getMask()) {
1128  // Intersect the in-bounds with the mask specified as an op parameter.
1129  mask = rewriter.create<arith::AndIOp>(loc, mask, xferOp.getMask());
1130  }
1131 
1132  rewriter.modifyOpInPlace(xferOp, [&]() {
1133  xferOp.getMaskMutable().assign(mask);
1134  xferOp.setInBoundsAttr(rewriter.getBoolArrayAttr({true}));
1135  });
1136 
1137  return success();
1138  }
1139 
1140 private:
1141  const bool force32BitVectorIndices;
1142 };
1143 
1144 /// Conversion pattern for a `vector.create_mask` (0-D and 1-D only).
1145 class VectorCreateMaskOpConversion
1146  : public OpRewritePattern<vector::CreateMaskOp> {
1147 public:
1148  explicit VectorCreateMaskOpConversion(MLIRContext *context,
1149  bool enableIndexOpt,
1150  PatternBenefit benefit = 1)
1151  : mlir::OpRewritePattern<vector::CreateMaskOp>(context, benefit),
1152  force32BitVectorIndices(enableIndexOpt) {}
1153 
1154  LogicalResult matchAndRewrite(vector::CreateMaskOp op,
1155  PatternRewriter &rewriter) const override {
1156  auto dstType = op.getType();
1157  if (cast<VectorType>(dstType).isScalable())
1158  return failure();
1159  int64_t rank = dstType.getRank();
1160  if (rank > 1)
1161  return failure();
1162  rewriter.replaceOp(
1163  op, buildVectorComparison(rewriter, op, force32BitVectorIndices,
1164  rank == 0 ? 0 : dstType.getDimSize(0),
1165  op.getOperand(0)));
1166  return success();
1167  }
1168 
1169 private:
1170  const bool force32BitVectorIndices;
1171 };
1172 
1173 /// Returns true if all the `i1` elements of `constantOp` are set to `value`.
1174 static bool allI1ConstantValuesSetTo(arith::ConstantOp constantOp, bool value) {
1175  auto denseAttr = dyn_cast<DenseIntElementsAttr>(constantOp.getValue());
1176  // TODO: Support non-dense constant.
1177  if (!denseAttr)
1178  return false;
1179 
1180  assert(denseAttr.getElementType().isInteger(1) && "Unexpected type");
1181  return denseAttr.isSplat() && denseAttr.getSplatValue<bool>() == value;
1182 }
1183 
1184 /// Folds a select operation between an all-true and all-false vector. For now,
1185 /// only single element vectors (i.e., vector<1xi1>) are supported. That is:
1186 ///
1187 /// %true = arith.constant dense<true> : vector<1xi1>
1188 /// %false = arith.constant dense<false> : vector<1xi1>
1189 /// %result = arith.select %cond, %true, %false : i1, vector<1xi1>
1190 /// =>
1191 /// %result = vector.broadcast %cond : i1 to vector<1xi1>
1192 ///
1193 /// InstCombine seems to handle vectors with multiple elements but not the
1194 /// single element ones.
1195 struct FoldI1Select : public OpRewritePattern<arith::SelectOp> {
1197 
1198  LogicalResult matchAndRewrite(arith::SelectOp selectOp,
1199  PatternRewriter &rewriter) const override {
1200  auto vecType = dyn_cast<VectorType>(selectOp.getType());
1201  if (!vecType || !vecType.getElementType().isInteger(1))
1202  return failure();
1203 
1204  // Only scalar conditions can be folded.
1205  Value cond = selectOp.getCondition();
1206  if (isa<VectorType>(cond.getType()))
1207  return failure();
1208 
1209  // TODO: Support n-D and scalable vectors.
1210  if (vecType.getRank() != 1 || vecType.isScalable())
1211  return failure();
1212 
1213  // TODO: Support vectors with multiple elements.
1214  if (vecType.getShape()[0] != 1)
1215  return failure();
1216 
1217  auto trueConst = selectOp.getTrueValue().getDefiningOp<arith::ConstantOp>();
1218  if (!trueConst || !allI1ConstantValuesSetTo(trueConst, true))
1219  return failure();
1220 
1221  auto falseConst =
1222  selectOp.getFalseValue().getDefiningOp<arith::ConstantOp>();
1223  if (!falseConst || !allI1ConstantValuesSetTo(falseConst, false))
1224  return failure();
1225 
1226  // Replace select with its condition broadcasted to single element vector.
1227  auto elemType = rewriter.getIntegerType(vecType.getNumElements());
1228  auto bcastType = VectorType::get(/*shape=*/{1}, elemType);
1229  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(selectOp, bcastType, cond);
1230  return success();
1231  }
1232 };
1233 
1234 /// Returns the number of dims can be folded away from transfer ops. It returns
1235 /// a failure if it can not determine the number of dims to be folded.
1236 ///
1237 /// Ex 1: returns "2" if `srcType` is memref<512x16x1x1xf32> and
1238 /// `vectorType` is vector<16x16x1x1xf32>
1239 /// (there two inner most dims can be dropped by memref.subview ops)
1240 ///
1241 /// Ex 2: returns "1" if `srcType` is memref<512x16x1x1xf32> with
1242 /// [8192, 16, 8, 1] strides and `vectorType` is vector<16x16x1x1xf32>
1243 /// (only the inner most unit dim of `srcType` can be dropped)
1244 ///
1245 /// Ex 3: return "0" if `srcType` is memref<512x16x1x1xf32> and
1246 /// `vectorType` is vector<16x16x1x[1]xf32>
1247 /// (the most inner dim in `vectorType` is not a unit dim (it's a "scalable
1248 /// unit")
1249 static FailureOr<size_t>
1250 getTransferFoldableInnerUnitDims(MemRefType srcType, VectorType vectorType) {
1251  SmallVector<int64_t> srcStrides;
1252  int64_t srcOffset;
1253  if (failed(srcType.getStridesAndOffset(srcStrides, srcOffset)))
1254  return failure();
1255 
1256  auto isUnitDim = [](VectorType type, int dim) {
1257  return type.getDimSize(dim) == 1 && !type.getScalableDims()[dim];
1258  };
1259 
1260  // According to vector.transfer_read/write semantics, the vector can be a
1261  // slice. Thus, we have to offset the check index with `rankDiff` in
1262  // `srcStrides` and source dim sizes.
1263  size_t result = 0;
1264  int rankDiff = srcType.getRank() - vectorType.getRank();
1265  for (int64_t i = 0, e = vectorType.getRank(); i < e; ++i) {
1266  // Check that the inner dim size is 1 for both memref type and vector slice.
1267  // It can be folded only if they are 1 and the stride is 1.
1268  int dim = vectorType.getRank() - i - 1;
1269  if (srcStrides[dim + rankDiff] != 1 ||
1270  srcType.getDimSize(dim + rankDiff) != 1 || !isUnitDim(vectorType, dim))
1271  break;
1272  result++;
1273  }
1274  return result;
1275 }
1276 
1277 /// Drop inner most contiguous unit dimensions from transfer_read operand.
1278 class DropInnerMostUnitDimsTransferRead
1279  : public OpRewritePattern<vector::TransferReadOp> {
1281 
1282  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
1283  PatternRewriter &rewriter) const override {
1284  // TODO: support 0-d corner case.
1285  if (readOp.getTransferRank() == 0)
1286  return failure();
1287 
1288  // TODO: support mask.
1289  if (readOp.getMask())
1290  return failure();
1291 
1292  auto srcType = dyn_cast<MemRefType>(readOp.getSource().getType());
1293  if (!srcType)
1294  return failure();
1295 
1296  if (!readOp.getPermutationMap().isMinorIdentity())
1297  return failure();
1298 
1299  auto targetType = readOp.getVectorType();
1300  if (targetType.getRank() <= 1)
1301  return failure();
1302 
1303  FailureOr<size_t> maybeDimsToDrop =
1304  getTransferFoldableInnerUnitDims(srcType, targetType);
1305  if (failed(maybeDimsToDrop))
1306  return failure();
1307 
1308  size_t dimsToDrop = maybeDimsToDrop.value();
1309  if (dimsToDrop == 0)
1310  return failure();
1311 
1312  auto inBounds = readOp.getInBoundsValues();
1313  auto droppedInBounds = ArrayRef<bool>(inBounds).take_back(dimsToDrop);
1314  if (llvm::is_contained(droppedInBounds, false))
1315  return failure();
1316 
1317  auto resultTargetVecType =
1318  VectorType::get(targetType.getShape().drop_back(dimsToDrop),
1319  targetType.getElementType(),
1320  targetType.getScalableDims().drop_back(dimsToDrop));
1321 
1322  auto loc = readOp.getLoc();
1324  memref::getMixedSizes(rewriter, loc, readOp.getSource());
1325  SmallVector<OpFoldResult> offsets(srcType.getRank(),
1326  rewriter.getIndexAttr(0));
1327  SmallVector<OpFoldResult> strides(srcType.getRank(),
1328  rewriter.getIndexAttr(1));
1329  MemRefType resultMemrefType = memref::SubViewOp::inferRankReducedResultType(
1330  srcType.getShape().drop_back(dimsToDrop), srcType, offsets, sizes,
1331  strides);
1332  ArrayAttr inBoundsAttr = rewriter.getArrayAttr(
1333  readOp.getInBoundsAttr().getValue().drop_back(dimsToDrop));
1334  Value rankedReducedView = rewriter.create<memref::SubViewOp>(
1335  loc, resultMemrefType, readOp.getSource(), offsets, sizes, strides);
1336  auto permMap = getTransferMinorIdentityMap(
1337  cast<ShapedType>(rankedReducedView.getType()), resultTargetVecType);
1338  Value result = rewriter.create<vector::TransferReadOp>(
1339  loc, resultTargetVecType, rankedReducedView,
1340  readOp.getIndices().drop_back(dimsToDrop), AffineMapAttr::get(permMap),
1341  readOp.getPadding(),
1342  // TODO: support mask.
1343  /*mask=*/Value(), inBoundsAttr);
1344  rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(readOp, targetType,
1345  result);
1346  return success();
1347  }
1348 };
1349 
1350 /// Drop inner most contiguous unit dimensions from transfer_write operand.
1351 /// E.g.,
1352 /// vector.transfer_write %arg1, %arg0[%c0, %arg2, %c0, %c0, %c0]
1353 /// {in_bounds = [true, true, true, true, true]}
1354 /// : vector<1x16x16x1x1xf32>, memref<1x512x16x1x1xf32>
1355 ///
1356 /// will be replaced with
1357 ///
1358 /// %subview = memref.subview %arg0
1359 /// [0, 0, 0, 0, 0] [1, 512, 16, 1, 1] [1, 1, 1, 1, 1]
1360 /// : memref<1x512x16x1x1xf32> to memref<1x512x16xf32>
1361 /// %0 = vector.shape_cast %arg1 : vector<1x16x16x1x1xf32>
1362 /// to vector<1x16x16xf32>
1363 /// vector.transfer_write %0, %subview[%c0, %arg2, %c0]
1364 /// {in_bounds = [true, true, true]}
1365 /// : vector<1x16x16xf32>, memref<1x512x16xf32>
1366 ///
1367 /// Note, this pattern will not collapse "scalable unit" dims (i.e. `[1]`).
1368 class DropInnerMostUnitDimsTransferWrite
1369  : public OpRewritePattern<vector::TransferWriteOp> {
1371 
1372  LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
1373  PatternRewriter &rewriter) const override {
1374  // TODO: support 0-d corner case.
1375  if (writeOp.getTransferRank() == 0)
1376  return failure();
1377 
1378  // TODO: support mask.
1379  if (writeOp.getMask())
1380  return failure();
1381 
1382  auto srcType = dyn_cast<MemRefType>(writeOp.getSource().getType());
1383  if (!srcType)
1384  return failure();
1385 
1386  if (!writeOp.getPermutationMap().isMinorIdentity())
1387  return failure();
1388 
1389  auto targetType = writeOp.getVectorType();
1390  if (targetType.getRank() <= 1)
1391  return failure();
1392 
1393  FailureOr<size_t> maybeDimsToDrop =
1394  getTransferFoldableInnerUnitDims(srcType, targetType);
1395  if (failed(maybeDimsToDrop))
1396  return failure();
1397 
1398  size_t dimsToDrop = maybeDimsToDrop.value();
1399  if (dimsToDrop == 0)
1400  return failure();
1401 
1402  auto inBounds = writeOp.getInBoundsValues();
1403  auto droppedInBounds = ArrayRef<bool>(inBounds).take_back(dimsToDrop);
1404  if (llvm::is_contained(droppedInBounds, false))
1405  return failure();
1406 
1407  auto resultTargetVecType =
1408  VectorType::get(targetType.getShape().drop_back(dimsToDrop),
1409  targetType.getElementType(),
1410  targetType.getScalableDims().drop_back(dimsToDrop));
1411 
1412  Location loc = writeOp.getLoc();
1414  memref::getMixedSizes(rewriter, loc, writeOp.getSource());
1415  SmallVector<OpFoldResult> offsets(srcType.getRank(),
1416  rewriter.getIndexAttr(0));
1417  SmallVector<OpFoldResult> strides(srcType.getRank(),
1418  rewriter.getIndexAttr(1));
1419  MemRefType resultMemrefType = memref::SubViewOp::inferRankReducedResultType(
1420  srcType.getShape().drop_back(dimsToDrop), srcType, offsets, sizes,
1421  strides);
1422  ArrayAttr inBoundsAttr = rewriter.getArrayAttr(
1423  writeOp.getInBoundsAttr().getValue().drop_back(dimsToDrop));
1424 
1425  Value rankedReducedView = rewriter.create<memref::SubViewOp>(
1426  loc, resultMemrefType, writeOp.getSource(), offsets, sizes, strides);
1427  auto permMap = getTransferMinorIdentityMap(
1428  cast<ShapedType>(rankedReducedView.getType()), resultTargetVecType);
1429 
1430  auto shapeCast = rewriter.createOrFold<vector::ShapeCastOp>(
1431  loc, resultTargetVecType, writeOp.getVector());
1432  rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
1433  writeOp, shapeCast, rankedReducedView,
1434  writeOp.getIndices().drop_back(dimsToDrop), AffineMapAttr::get(permMap),
1435  // TODO: support mask.
1436  /*mask=*/Value(), inBoundsAttr);
1437  return success();
1438  }
1439 };
1440 
1441 /// Canonicalization of a `vector.contraction %a, %b, %c` with row-major matmul
1442 /// semantics to a contraction suitable for MMT (matrix matrix multiplication
1443 /// with the RHS transposed) lowering.
1444 struct CanonicalizeContractMatmulToMMT final
1445  : OpRewritePattern<vector::ContractionOp> {
1447 
1448  using FilterConstraintType =
1449  std::function<LogicalResult(vector::ContractionOp op)>;
1450 
1451  CanonicalizeContractMatmulToMMT(MLIRContext *context, PatternBenefit benefit,
1452  FilterConstraintType constraint)
1453  : OpRewritePattern<vector::ContractionOp>(context, benefit),
1454  filter(std::move(constraint)) {}
1455 
1456  LogicalResult matchAndRewrite(vector::ContractionOp op,
1457  PatternRewriter &rewriter) const override {
1458  if (failed(filter(op)))
1459  return failure();
1460 
1461  Location loc = op.getLoc();
1462  Value lhs = op.getLhs();
1463  Value rhs = op.getRhs();
1464  Value res = op.getAcc();
1465 
1466  // Set up the parallel/reduction structure in right form.
1467  using MapList = ArrayRef<ArrayRef<AffineExpr>>;
1468  auto infer = [&](MapList m) {
1469  return AffineMap::inferFromExprList(m, op.getContext());
1470  };
1471  AffineExpr m;
1472  AffineExpr n;
1473  AffineExpr k;
1474  bindDims(rewriter.getContext(), m, n, k);
1475  static constexpr std::array<int64_t, 2> perm = {1, 0};
1476  auto iteratorTypes = op.getIteratorTypes().getValue();
1477  SmallVector<AffineMap, 4> maps = op.getIndexingMapsArray();
1478  if (iteratorTypes.size() != 3 ||
1479  !vector::isParallelIterator(iteratorTypes[0]) ||
1480  !vector::isParallelIterator(iteratorTypes[1]) ||
1481  !vector::isReductionIterator(iteratorTypes[2]))
1482  return rewriter.notifyMatchFailure(op, "contraction is not a gemm");
1483 
1484  // The canonical form is "TNT" = A row-major, B col-major, C row-major.
1485  const auto canonicalForm = infer({{m, k}, {n, k}, {m, n}});
1486  if (maps == canonicalForm)
1487  return rewriter.notifyMatchFailure(op, "already in the canonical form");
1488 
1489  // Create a vector transpose making sure to emit zero/sign-extend at the
1490  // end.
1491  auto createTranspose = [&rewriter, loc](Value mat) -> Value {
1492  if (auto sext = mat.getDefiningOp<arith::ExtSIOp>()) {
1493  Value trans =
1494  rewriter.create<vector::TransposeOp>(loc, sext.getIn(), perm);
1495  VectorType newType =
1496  cast<VectorType>(trans.getType())
1497  .clone(cast<VectorType>(mat.getType()).getElementType());
1498  return rewriter.create<arith::ExtSIOp>(loc, newType, trans);
1499  }
1500  if (auto zext = mat.getDefiningOp<arith::ExtUIOp>()) {
1501  Value trans =
1502  rewriter.create<vector::TransposeOp>(loc, zext.getIn(), perm);
1503  VectorType newType =
1504  VectorType::get(cast<VectorType>(trans.getType()).getShape(),
1505  cast<VectorType>(mat.getType()).getElementType());
1506  return rewriter.create<arith::ExtUIOp>(loc, newType, trans);
1507  }
1508  return rewriter.create<vector::TransposeOp>(loc, mat, perm);
1509  };
1510 
1511  if (maps == infer({{m, k}, {k, n}, {m, n}})) {
1512  rhs = createTranspose(rhs);
1513  } else if (maps == infer({{k, m}, {n, k}, {m, n}})) {
1514  lhs = createTranspose(lhs);
1515  } else if (maps == infer({{k, m}, {k, n}, {m, n}})) {
1516  rhs = createTranspose(rhs);
1517  lhs = createTranspose(lhs);
1518  } else if (maps == infer({{k, m}, {k, n}, {n, m}})) {
1519  std::swap(rhs, lhs);
1520  rhs = createTranspose(rhs);
1521  lhs = createTranspose(lhs);
1522  } else if (maps == infer({{k, m}, {n, k}, {n, m}})) {
1523  std::swap(rhs, lhs);
1524  rhs = createTranspose(rhs);
1525  } else if (maps == infer({{m, k}, {k, n}, {n, m}})) {
1526  std::swap(lhs, rhs);
1527  lhs = createTranspose(lhs);
1528  } else if (maps == infer({{m, k}, {n, k}, {n, m}})) {
1529  std::swap(lhs, rhs);
1530  } else {
1531  return rewriter.notifyMatchFailure(op, "unhandled contraction form");
1532  }
1533  rewriter.replaceOpWithNewOp<vector::ContractionOp>(
1534  op, lhs, rhs, res, rewriter.getAffineMapArrayAttr(canonicalForm),
1535  op.getIteratorTypes());
1536  return success();
1537  };
1538 
1539 private:
1540  FilterConstraintType filter;
1541 };
1542 
1543 /// Pattern to fold arithmetic extensions on floating point data types into
1544 /// vector contraction operations. linalg.matmul introduces arithmetic
1545 /// extensions on its operands. Please mlir snippets below for more details.
1546 /// ```mlir
1547 /// "linalg.matmul"(%lhs, %rhs, %acc) ({
1548 /// ^bb0(%arg1: f16, %arg2: f16, %arg3: f32):
1549 /// %lhs_f32 = "arith.extf"(%arg1) : (f16) -> f32
1550 /// %rhs_f32 = "arith.extf"(%arg2) : (f16) -> f32
1551 /// %mul = "arith.mulf"(%lhs_f32, %rhs_f32) : (f32, f32) -> f32
1552 /// %acc = "arith.addf"(%arg3, %mul) : (f32, f32) -> f32
1553 /// "linalg.yield"(%acc) : (f32) -> ()
1554 /// })
1555 /// ```
1556 /// This restricts the native usage of mixed precision NVIDIA Ampere Tensor
1557 /// Cores, i.e, `mma.sync.*.f32.f16.f16.f32` and `mma.sync.*.f32.bf16.bf16.f32`.
1558 /// This pattern folds the arithmetic extensions into the vector contraction and
1559 /// enables the usage of native mixed precision Tensor Core instructions.
1560 template <typename ExtOp>
1561 struct FoldArithExtIntoContractionOp
1562  : public OpRewritePattern<vector::ContractionOp> {
1564 
1565  LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
1566  PatternRewriter &rewriter) const override {
1567 
1568  auto lhsDefOp = contractOp.getLhs().getDefiningOp<ExtOp>();
1569  auto rhsDefOp = contractOp.getRhs().getDefiningOp<ExtOp>();
1570 
1571  if (!lhsDefOp || !rhsDefOp) {
1572  return rewriter.notifyMatchFailure(contractOp,
1573  "no defining op on contract operands");
1574  }
1575 
1576  rewriter.replaceOpWithNewOp<vector::ContractionOp>(
1577  contractOp, lhsDefOp->getOperand(0), rhsDefOp->getOperand(0),
1578  contractOp.getAcc(), contractOp.getIndexingMapsAttr(),
1579  contractOp.getIteratorTypesAttr());
1580 
1581  return success();
1582  }
1583 };
1584 
1585 /// Pattern to fold chained reduction to a series of vector additions and a
1586 /// final reduction. This form should require fewer subgroup operations.
1587 ///
1588 /// ```mlir
1589 /// %a = vector.reduction <add> %x, %acc
1590 /// %b = vector.reduction <add> %y, %a
1591 /// ==>
1592 /// %a = arith.addf %x, %y
1593 /// %b = vector.reduction <add> %a, %acc
1594 /// ```
1595 struct ChainedReduction final : OpRewritePattern<vector::ReductionOp> {
1597 
1598  LogicalResult matchAndRewrite(vector::ReductionOp op,
1599  PatternRewriter &rewriter) const override {
1600  // TODO: Handle other combining kinds.
1601  if (op.getKind() != vector::CombiningKind::ADD)
1602  return failure();
1603 
1604  // Accumulator is optional.
1605  Value acc = op.getAcc();
1606  if (!acc)
1607  return failure();
1608 
1609  if (!acc.getType().isIntOrFloat())
1610  return failure();
1611 
1612  auto parentReduction = acc.getDefiningOp<vector::ReductionOp>();
1613  if (!parentReduction)
1614  return failure();
1615 
1616  Location loc = op.getLoc();
1617  Value vAdd;
1618  if (isa<IntegerType>(acc.getType())) {
1619  vAdd = rewriter.createOrFold<arith::AddIOp>(
1620  loc, parentReduction.getVector(), op.getVector());
1621  } else {
1622  vAdd = rewriter.create<arith::AddFOp>(loc, parentReduction.getVector(),
1623  op.getVector());
1624  }
1625  rewriter.replaceOpWithNewOp<vector::ReductionOp>(op, op.getKind(), vAdd,
1626  parentReduction.getAcc());
1627  return success();
1628  }
1629 };
1630 
1631 // Helper function dropping unit non-scalable dimension from a VectorType
1632 // keeping at least 1 dimension to avoid generating 0-D vectors. Scalable unit
1633 // dimensions are not dropped. Folding such dimensions would require "shifting"
1634 // the scalable flag onto some other fixed-width dim (e.g. vector<[1]x4xf32> ->
1635 // vector<[4]xf32>). This could be implemented in the future.
1636 static VectorType dropNonScalableUnitDimFromType(VectorType inVecTy) {
1637  auto inVecShape = inVecTy.getShape();
1638  SmallVector<int64_t> newShape;
1639  SmallVector<bool> newScalableDims;
1640  for (auto [dim, isScalable] :
1641  llvm::zip_equal(inVecShape, inVecTy.getScalableDims())) {
1642  if (dim == 1 && !isScalable)
1643  continue;
1644 
1645  newShape.push_back(dim);
1646  newScalableDims.push_back(isScalable);
1647  }
1648  // All dims have been dropped, return vector<1xeType>.
1649  if (newShape.empty()) {
1650  newShape.push_back(1);
1651  newScalableDims.push_back(false);
1652  }
1653 
1654  return VectorType::get(newShape, inVecTy.getElementType(), newScalableDims);
1655 }
1656 
1657 /// For vectors with at least one unit dim, replaces:
1658 /// elementwise(a, b)
1659 /// with:
1660 /// sc_a = shape_cast(a)
1661 /// sc_b = shape_cast(b)
1662 /// res = elementwise(sc_a, sc_b)
1663 /// return shape_cast(res)
1664 /// The newly inserted shape_cast Ops fold (before elementwise Op) and then
1665 /// restore (after elementwise Op) the unit dim. Vectors `a` and `b` are
1666 /// required to be rank > 1.
1667 ///
1668 /// Ex:
1669 /// %mul = arith.mulf %B_row, %A_row : vector<1x[4]xf32>
1670 /// %cast = vector.shape_cast %mul : vector<1x[4]xf32> to vector<[4]xf32>
1671 ///
1672 /// gets converted to:
1673 ///
1674 /// %B_row_sc = vector.shape_cast %B_row : vector<1x[4]xf32> to vector<[4]xf32>
1675 /// %A_row_sc = vector.shape_cast %A_row : vector<1x[4]xf32> to vector<[4]xf32>
1676 /// %mul = arith.mulf %B_row_sc, %A_row_sc : vector<[4]xf32>
1677 /// %cast_new = vector.shape_cast %mul : vector<[4]xf32> to vector<1x[4]xf32>
1678 /// %cast = vector.shape_cast %cast_new : vector<1x[4]xf32> to vector<[4]xf32>
1679 ///
1680 /// Patterns for folding shape_casts should instantly eliminate `%cast_new` and
1681 /// `%cast`.
1682 struct DropUnitDimFromElementwiseOps final
1683  : public OpTraitRewritePattern<OpTrait::Elementwise> {
1685  LogicalResult matchAndRewrite(Operation *op,
1686  PatternRewriter &rewriter) const override {
1687  if (op->getNumResults() != 1 || op->getNumRegions() != 0)
1688  return failure();
1689 
1690  auto resultVectorType = dyn_cast<VectorType>(op->getResult(0).getType());
1691  if (!resultVectorType)
1692  return failure();
1693 
1694  // Check the operand pre-conditions. For `Elementwise` ops all operands are
1695  // guaranteed to have identical shapes (with some exceptions such as
1696  // `arith.select`) and it suffices to only check one of them.
1697  auto sourceVectorType = dyn_cast<VectorType>(op->getOperand(0).getType());
1698  if (!sourceVectorType)
1699  return failure();
1700  if (sourceVectorType.getRank() < 2)
1701  return failure();
1702 
1703  SmallVector<Value> newOperands;
1704  auto loc = op->getLoc();
1705  for (auto operand : op->getOperands()) {
1706  auto opVectorType = cast<VectorType>(operand.getType());
1707  auto newVType = dropNonScalableUnitDimFromType(opVectorType);
1708  if (newVType == opVectorType)
1709  return rewriter.notifyMatchFailure(op, "No unit dimension to remove.");
1710 
1711  auto opSC = rewriter.create<vector::ShapeCastOp>(loc, newVType, operand);
1712  newOperands.push_back(opSC);
1713  }
1714 
1715  VectorType newResultVectorType =
1716  dropNonScalableUnitDimFromType(resultVectorType);
1717  // Create an updated elementwise Op without unit dim.
1718  Operation *elementwiseOp =
1719  rewriter.create(loc, op->getName().getIdentifier(), newOperands,
1720  newResultVectorType, op->getAttrs());
1721 
1722  // Restore the unit dim by applying vector.shape_cast to the result.
1723  rewriter.replaceOpWithNewOp<ShapeCastOp>(op, resultVectorType,
1724  elementwiseOp->getResult(0));
1725 
1726  return success();
1727  }
1728 };
1729 
1730 /// A pattern to drop unit dims from vector.transpose.
1731 ///
1732 /// Example:
1733 ///
1734 /// BEFORE:
1735 /// ```mlir
1736 /// %transpose = vector.transpose %vector, [3, 0, 1, 2]
1737 /// : vector<1x1x4x[4]xf32> to vector<[4]x1x1x4xf32>
1738 /// ```
1739 ///
1740 /// AFTER:
1741 /// ```mlir
1742 /// %dropDims = vector.shape_cast %vector
1743 /// : vector<1x1x4x[4]xf32> to vector<4x[4]xf32>
1744 /// %transpose = vector.transpose %0, [1, 0]
1745 /// : vector<4x[4]xf32> to vector<[4]x4xf32>
1746 /// %restoreDims = vector.shape_cast %transpose
1747 /// : vector<[4]x4xf32> to vector<[4]x1x1x4xf32>
1748 /// ```
1749 struct DropUnitDimsFromTransposeOp final
1750  : OpRewritePattern<vector::TransposeOp> {
1752 
1753  LogicalResult matchAndRewrite(vector::TransposeOp op,
1754  PatternRewriter &rewriter) const override {
1755  VectorType sourceType = op.getSourceVectorType();
1756  VectorType sourceTypeWithoutUnitDims =
1757  dropNonScalableUnitDimFromType(sourceType);
1758 
1759  if (sourceType == sourceTypeWithoutUnitDims)
1760  return failure();
1761 
1762  // Construct a map from dimIdx -> number of dims dropped before dimIdx.
1763  auto sourceDims = llvm::to_vector(vector::getDims(sourceType));
1764  SmallVector<int64_t> droppedDimsBefore(sourceType.getRank());
1765  int64_t droppedDims = 0;
1766  for (auto [i, dim] : llvm::enumerate(sourceDims)) {
1767  droppedDimsBefore[i] = droppedDims;
1768  if (dim == std::make_tuple(1, false))
1769  ++droppedDims;
1770  }
1771 
1772  // Drop unit dims from transpose permutation.
1773  ArrayRef<int64_t> perm = op.getPermutation();
1774  SmallVector<int64_t> newPerm;
1775  for (int64_t idx : perm) {
1776  if (sourceDims[idx] == std::make_tuple(1, false))
1777  continue;
1778  newPerm.push_back(idx - droppedDimsBefore[idx]);
1779  }
1780 
1781  // Fixup for `newPerm`. The `sourceTypeWithoutUnitDims` could be vector<1xT>
1782  // type when the dimensions are unit dimensions. In this case, the newPerm
1783  // should be [0].
1784  if (newPerm.empty()) {
1785  newPerm.push_back(0);
1786  }
1787 
1788  Location loc = op.getLoc();
1789  // Drop the unit dims via shape_cast.
1790  auto dropDimsShapeCast = rewriter.create<vector::ShapeCastOp>(
1791  loc, sourceTypeWithoutUnitDims, op.getVector());
1792  // Create the new transpose.
1793  auto transposeWithoutUnitDims =
1794  rewriter.create<vector::TransposeOp>(loc, dropDimsShapeCast, newPerm);
1795  // Restore the unit dims via shape cast.
1796  rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
1797  op, op.getResultVectorType(), transposeWithoutUnitDims);
1798 
1799  return success();
1800  }
1801 };
1802 
1803 /// A pattern to drop unit dims from the iter_args of an scf.for.
1804 ///
1805 /// Example:
1806 ///
1807 /// BEFORE:
1808 /// ```mlir
1809 /// %res = scf.for ... iter_args(%iter = %init) -> vector<[4]x1x1x4xf32> {
1810 /// ...
1811 /// scf.yield %
1812 /// }
1813 /// ```
1814 ///
1815 /// AFTER:
1816 /// ```mlir
1817 /// %drop = vector.shape_cast %init
1818 /// : vector<4x1x1x[4]xf32> to vector<4x[4]xf32>
1819 /// %new_loop = scf.for ... iter_args(%iter = %drop) -> vector<[4]x4xf32> {
1820 /// %new_iter = vector.shape_cast %iter
1821 /// : vector<[4]x4xf32> to vector<[4]x1x1x4xf32>
1822 /// ...
1823 /// }
1824 /// %res = vector.shape_cast %new_loop
1825 /// : vector<[4]x4xf32> to vector<[4]x1x1x4xf32>
1826 /// ```
1827 struct DropUnitDimsFromScfForOp final : OpRewritePattern<scf::ForOp> {
1829 
1830  LogicalResult matchAndRewrite(scf::ForOp forOp,
1831  PatternRewriter &rewriter) const override {
1832  /// Find the first iter_arg with droppable unit dims. Further applications
1833  /// of this pattern will apply to later arguments.
1834  for (OpOperand &operand : forOp.getInitArgsMutable()) {
1835  auto vectorType = dyn_cast<VectorType>(operand.get().getType());
1836  if (!vectorType)
1837  continue;
1838 
1839  VectorType newVectorType = dropNonScalableUnitDimFromType(vectorType);
1840  if (vectorType == newVectorType)
1841  continue;
1842 
1843  // Create a new ForOp with that iter operand replaced.
1844  auto castFn = [](OpBuilder &b, Location loc, Type type, Value source) {
1845  return b.create<vector::ShapeCastOp>(loc, type, source);
1846  };
1847 
1848  Value replacement =
1849  castFn(rewriter, forOp.getLoc(), newVectorType, operand.get());
1850  rewriter.replaceOp(forOp,
1851  replaceAndCastForOpIterArg(rewriter, forOp, operand,
1852  replacement, castFn));
1853  return success();
1854  }
1855  return failure();
1856  }
1857 };
1858 
1859 /// Pattern to eliminate redundant zero-constants added to reduction operands.
1860 /// It's enough for there to be one initial zero value, so we can eliminate the
1861 /// extra ones that feed into `vector.reduction <add>`. These get created by the
1862 /// `ChainedReduction` pattern.
1863 ///
1864 /// ```mlir
1865 /// %a = arith.addf %x, %zero
1866 /// %b = arith.addf %a, %y
1867 /// %c = vector.reduction <add> %b, %acc
1868 /// ==>
1869 /// %b = arith.addf %a, %y
1870 /// %c = vector.reduction <add> %b, %acc
1871 /// ```
1872 struct ReduceRedundantZero final : OpRewritePattern<vector::ReductionOp> {
1874 
1875  LogicalResult matchAndRewrite(vector::ReductionOp op,
1876  PatternRewriter &rewriter) const override {
1877  // TODO: Handle other reduction kinds and their identity values.
1878  if (op.getKind() != vector::CombiningKind::ADD)
1879  return failure();
1880 
1881  Type elemType = op.getSourceVectorType().getElementType();
1882  // The integer case should be handled by `arith.addi` folders, only check
1883  // for floats here.
1884  if (!isa<FloatType>(elemType))
1885  return failure();
1886 
1887  auto vAdd = op.getVector().getDefiningOp<arith::AddFOp>();
1888  if (!vAdd)
1889  return failure();
1890  auto addLhs = vAdd.getLhs().getDefiningOp<arith::AddFOp>();
1891  if (!addLhs)
1892  return failure();
1893 
1894  if (!matchPattern(addLhs.getRhs(), m_AnyZeroFloat()))
1895  return failure();
1896 
1897  auto newAdd = rewriter.create<arith::AddFOp>(vAdd.getLoc(), addLhs.getLhs(),
1898  vAdd.getRhs());
1899  rewriter.replaceOpWithNewOp<vector::ReductionOp>(op, op.getKind(), newAdd,
1900  op.getAcc());
1901  return success();
1902  }
1903 };
1904 
1905 /// Example:
1906 /// ```
1907 /// %a = vector.reduction <add> %x : vector<2xf32> into f32
1908 /// ```
1909 /// is transformed into:
1910 /// ```
1911 /// %y = vector.extract %x[0] : f32 from vector<2xf32>
1912 /// %z = vector.extract %x[1] : f32 from vector<2xf32>
1913 /// %a = arith.addf %y, %z : f32
1914 /// ```
1915 struct BreakDownVectorReduction final : OpRewritePattern<vector::ReductionOp> {
1916  BreakDownVectorReduction(MLIRContext *context,
1917  unsigned maxNumElementsToExtract,
1918  PatternBenefit benefit)
1919  : OpRewritePattern(context, benefit),
1920  maxNumElementsToExtract(maxNumElementsToExtract) {}
1921 
1922  LogicalResult matchAndRewrite(vector::ReductionOp op,
1923  PatternRewriter &rewriter) const override {
1924  VectorType type = op.getSourceVectorType();
1925  if (type.isScalable() || op.isMasked())
1926  return failure();
1927  assert(type.getRank() == 1 && "Expected a 1-d vector");
1928 
1929  int64_t numElems = type.getNumElements();
1930  if (numElems > maxNumElementsToExtract) {
1931  return rewriter.notifyMatchFailure(
1932  op, llvm::formatv("has too many vector elements ({0}) to break down "
1933  "(max allowed: {1})",
1934  numElems, maxNumElementsToExtract));
1935  }
1936 
1937  Location loc = op.getLoc();
1938  SmallVector<Value> extracted(numElems, nullptr);
1939  for (auto [idx, extractedElem] : llvm::enumerate(extracted))
1940  extractedElem = rewriter.create<vector::ExtractOp>(
1941  loc, op.getVector(), static_cast<int64_t>(idx));
1942 
1943  Value res = extracted.front();
1944  for (auto extractedElem : llvm::drop_begin(extracted))
1945  res = vector::makeArithReduction(rewriter, loc, op.getKind(), res,
1946  extractedElem, op.getFastmathAttr());
1947  if (Value acc = op.getAcc())
1948  res = vector::makeArithReduction(rewriter, loc, op.getKind(), res, acc,
1949  op.getFastmathAttr());
1950 
1951  rewriter.replaceOp(op, res);
1952  return success();
1953  }
1954 
1955 private:
1956  unsigned maxNumElementsToExtract = 0;
1957 };
1958 
1959 /// Fold `mulf(tr(broadcast(A)), broadcast(B))` into `vector.outerproduct(A,
1960 /// B)`.
1961 /// Example:
1962 /// %lhsBcast = vector.broadcast %lhs : vector<4xi32> to vector<4x4xi32>
1963 /// %lhsT = vector.transpose %lhsBcast, [1, 0] : vector<4x4xi32> to
1964 /// vector<4x4xi32> %rhsBcast = vector.broadcast %rhs : vector<4xi32> to
1965 /// vector<4x4xi32> %mul = arith.muli %lhsT, %rhsBcast : vector<4x4xi32>
1966 ///
1967 /// Becomes :
1968 ///
1969 /// %res = vector.outerproduct %lhs, %rhs : vector<4xi32>, vector<4xi32>
1970 ///
1971 /// Supports only 1D-to-2D broadcasts. The following cases are not supported.
1972 /// %ex1 = vector.broadcast %lhsCast : vector<1x4xf32> to vector<4x4xf32>
1973 /// %ex2 = vector.broadcast %lhsCast : f32 to vector<4x4xf32>
1974 /// %ex3 = vector.broadcast %lhsCast : vector<1x1xf32> to vector<4x4xf32>
1975 template <typename MulOpType>
1976 struct FoldArithToVectorOuterProduct : public OpRewritePattern<MulOpType> {
1978  // Returns whether a vector.broadcast matches requirements for an outerproduct
1979  // pattern. aka a 1D-to-2D broadcastOp without broadcasted unit dimension.
1980  bool isValidBroadcastSource(vector::BroadcastOp broadcastOp) const {
1981  // Fail if it is not a 1-to-2 dimension to broadcast to avoid generating
1982  // shape_casts/broadcasts which does not belong in this pattern.
1983  if (!broadcastOp.computeBroadcastedUnitDims().empty())
1984  return false;
1985  // Avoid broadcast like f32 or vector<f32> -> ResType
1986  auto srcType = dyn_cast<VectorType>(broadcastOp.getSourceType());
1987  return srcType && srcType.getRank() != 2;
1988  }
1989 
1990  LogicalResult matchAndRewrite(MulOpType mulOp,
1991  PatternRewriter &rewriter) const override {
1992  auto resType = llvm::cast<VectorType>(mulOp.getResult().getType());
1993  if (!resType)
1994  return failure();
1995  if (resType.getRank() != 2)
1996  return failure();
1997  /// If operandA can be written as tr(broadcast(A)) and operandB as
1998  /// broadcast(B) where broadcasts are 1D-to-2D, create and return
1999  /// vector.outerproduct(A, B). Returns failure() otherwise.
2000  auto matchOuterProduct =
2001  [&](Value operandA,
2002  Value operandB) -> FailureOr<vector::OuterProductOp> {
2003  auto transposedLhs = operandA.getDefiningOp<vector::TransposeOp>();
2004  if (!transposedLhs)
2005  return failure();
2006  // Fail unless this is a true 2-D matrix transpose.
2007  ArrayRef<int64_t> permutation = transposedLhs.getPermutation();
2008  if (permutation.size() != 2 || permutation[0] != 1 || permutation[1] != 0)
2009  return failure();
2010 
2011  auto broadcastedLhs =
2012  transposedLhs.getVector().getDefiningOp<vector::BroadcastOp>();
2013  if (!broadcastedLhs || !isValidBroadcastSource(broadcastedLhs))
2014  return failure();
2015 
2016  auto broadcastedRhs = operandB.getDefiningOp<vector::BroadcastOp>();
2017  if (!broadcastedRhs || !isValidBroadcastSource(broadcastedRhs))
2018  return failure();
2019 
2020  return rewriter.create<vector::OuterProductOp>(
2021  mulOp->getLoc(), resType, broadcastedLhs.getSource(),
2022  broadcastedRhs.getSource(), Value(), vector::CombiningKind::ADD);
2023  };
2024 
2025  Value lhs = mulOp->getOperand(0), rhs = mulOp->getOperand(1);
2026  auto maybeOuterP = matchOuterProduct(lhs, rhs);
2027  // Handle commutativity, the transposed op is the outerproduct LHS.
2028  if (failed(maybeOuterP))
2029  maybeOuterP = matchOuterProduct(rhs, lhs);
2030  if (failed(maybeOuterP))
2031  return failure();
2032  rewriter.replaceOp(mulOp, maybeOuterP->getResult());
2033  return success();
2034  }
2035 };
2036 
2037 } // namespace
2038 
2041  patterns.add<FoldArithExtIntoContractionOp<arith::ExtFOp>,
2042  FoldArithExtIntoContractionOp<arith::ExtSIOp>>(
2043  patterns.getContext());
2044 }
2045 
2047  RewritePatternSet &patterns, bool force32BitVectorIndices,
2048  PatternBenefit benefit) {
2049  patterns.add<VectorCreateMaskOpConversion,
2050  MaterializeTransferMask<vector::TransferReadOp>,
2051  MaterializeTransferMask<vector::TransferWriteOp>>(
2052  patterns.getContext(), force32BitVectorIndices, benefit);
2053  patterns.add<FoldI1Select>(patterns.getContext(), benefit);
2054 }
2055 
2057  PatternBenefit benefit) {
2058  patterns.add<ShapeCastOpFolder>(patterns.getContext(), benefit);
2059 }
2060 
2063  // TODO: Consider either:
2064  // * including DropInnerMostUnitDimsTransferRead and
2065  // DropInnerMostUnitDimsTransferWrite, or
2066  // * better naming to distinguish this and
2067  // populateVectorTransferCollapseInnerMostContiguousDimsPatterns.
2068  patterns.add<DropUnitDimFromElementwiseOps, DropUnitDimsFromScfForOp,
2069  DropUnitDimsFromTransposeOp, ShapeCastOpFolder>(
2070  patterns.getContext(), benefit);
2071 }
2072 
2075  patterns.add<BubbleDownVectorBitCastForExtract,
2076  BubbleDownBitCastForStridedSliceExtract,
2077  BubbleUpBitCastForInsert, BubbleUpBitCastForStridedSliceInsert>(
2078  patterns.getContext(), benefit);
2079 }
2080 
2083  std::function<bool(vector::BitCastOp)> controlFn, PatternBenefit benefit) {
2084  patterns.add<BreakDownVectorBitCast>(patterns.getContext(),
2085  std::move(controlFn), benefit);
2086 }
2087 
2090  std::function<LogicalResult(vector::ContractionOp)> constraint,
2091  PatternBenefit benefit) {
2092  patterns.add<CanonicalizeContractMatmulToMMT>(patterns.getContext(), benefit,
2093  std::move(constraint));
2094 }
2095 
2098  patterns.add<MultiReduceToContract, CombineContractBroadcast,
2099  CombineContractABTranspose, CombineContractResultTranspose>(
2100  patterns.getContext(), benefit);
2101 }
2102 
2106  patterns.add<DropInnerMostUnitDimsTransferRead,
2107  DropInnerMostUnitDimsTransferWrite>(patterns.getContext(),
2108  benefit);
2109 }
2110 
2112  PatternBenefit benefit) {
2113  patterns.add<ReorderElementwiseOpsOnTranspose, ReorderCastOpsOnBroadcast,
2114  ReorderElementwiseOpsOnBroadcast>(patterns.getContext(),
2115  benefit);
2116 }
2117 
2120  patterns.add<ChainedReduction>(patterns.getContext(), benefit);
2121  patterns.add<ReduceRedundantZero>(patterns.getContext(),
2122  PatternBenefit(benefit.getBenefit() + 1));
2123 }
2124 
2126  RewritePatternSet &patterns, unsigned maxNumElementsToExtract,
2127  PatternBenefit benefit) {
2128  patterns.add<BreakDownVectorReduction>(patterns.getContext(),
2129  maxNumElementsToExtract, benefit);
2130 }
2131 
2134  patterns.add<FoldArithToVectorOuterProduct<arith::MulFOp>,
2135  FoldArithToVectorOuterProduct<arith::MulIOp>>(
2136  patterns.getContext());
2137 }
2138 
2139 //===----------------------------------------------------------------------===//
2140 // TableGen'd enum attribute definitions
2141 //===----------------------------------------------------------------------===//
2142 
2143 #include "mlir/Dialect/Vector/Transforms/VectorTransformsEnums.cpp.inc"
static uint64_t zext(uint32_t arg)
static Value broadcast(Location loc, Value toBroadcast, unsigned numElements, const TypeConverter &typeConverter, ConversionPatternRewriter &rewriter)
Broadcasts the value to vector with numElements number of elements.
static std::optional< int64_t > getResultIndex(AffineMap map, int64_t index)
static SmallVector< IntType > extractVector(ArrayAttr arrayAttr)
Base type for affine expression.
Definition: AffineExpr.h:68
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
unsigned getDimPosition(unsigned idx) const
Extracts the position of the dimensional expression at the given result, when the caller knows it is ...
Definition: AffineMap.cpp:415
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
unsigned getNumResults() const
Definition: AffineMap.cpp:402
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
Definition: AffineMap.cpp:264
AffineMap compose(AffineMap map) const
Returns the AffineMap resulting from composing this with map.
Definition: AffineMap.cpp:556
static SmallVector< AffineMap, 4 > inferFromExprList(ArrayRef< ArrayRef< AffineExpr >> exprsList, MLIRContext *context)
Returns a vector of AffineMaps; each with as many results as exprs.size(), as many dims as the larges...
Definition: AffineMap.cpp:312
Attributes are known-constant values of operations.
Definition: Attributes.h:25
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:104
AffineMap getMultiDimIdentityMap(unsigned rank)
Definition: Builders.cpp:383
IntegerType getI64Type()
Definition: Builders.cpp:65
IntegerType getI32Type()
Definition: Builders.cpp:63
IntegerType getIntegerType(unsigned width)
Definition: Builders.cpp:67
TypedAttr getZeroAttr(Type type)
Definition: Builders.cpp:320
AffineExpr getAffineDimExpr(unsigned position)
Definition: Builders.cpp:360
MLIRContext * getContext() const
Definition: Builders.h:56
DenseIntElementsAttr getI32VectorAttr(ArrayRef< int32_t > values)
Definition: Builders.cpp:118
DenseIntElementsAttr getI64VectorAttr(ArrayRef< int64_t > values)
Definition: Builders.cpp:124
IntegerType getI1Type()
Definition: Builders.cpp:53
ArrayAttr getArrayAttr(ArrayRef< Attribute > value)
Definition: Builders.cpp:262
ArrayAttr getI64ArrayAttr(ArrayRef< int64_t > values)
Definition: Builders.cpp:277
ArrayAttr getBoolArrayAttr(ArrayRef< bool > values)
Definition: Builders.cpp:266
ArrayAttr getAffineMapArrayAttr(ArrayRef< AffineMap > values)
Definition: Builders.cpp:314
An attribute that represents a reference to a dense integer vector or tensor object.
static DenseIntElementsAttr get(const ShapedType &type, Arg &&arg)
Get an instance of a DenseIntElementsAttr with the given arguments.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:66
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class helps build Operations.
Definition: Builders.h:205
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:518
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:453
This class represents an operand of an operation.
Definition: Value.h:267
OpTraitRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting again...
Definition: PatternMatch.h:396
OpTraitRewritePattern(MLIRContext *context, PatternBenefit benefit=1)
Definition: PatternMatch.h:398
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
Value getOperand(unsigned idx)
Definition: Operation.h:350
Operation * clone(IRMapping &mapper, CloneOptions options=CloneOptions::all())
Create a deep copy of this operation, remapping any operands that use values outside of the operation...
Definition: Operation.cpp:717
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:407
unsigned getNumRegions()
Returns the number of regions held by this operation.
Definition: Operation.h:674
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
unsigned getNumOperands()
Definition: Operation.h:346
ArrayRef< NamedAttribute > getAttrs()
Return all of the attributes on this operation.
Definition: Operation.h:512
OperationName getName()
The name of an operation is the key identifier for it.
Definition: Operation.h:119
result_type_range getResultTypes()
Definition: Operation.h:428
operand_range getOperands()
Returns an iterator on the underlying Value's.
Definition: Operation.h:378
result_range getResults()
Definition: Operation.h:415
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:404
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
Definition: PatternMatch.h:34
unsigned short getBenefit() const
If the corresponding pattern can match, return its benefit. If the.
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:803
type_range getType() const
Definition: ValueRange.cpp:39
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
Definition: PatternMatch.h:736
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
void modifyOpInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around an in-place modification of an operation.
Definition: PatternMatch.h:648
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:554
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
bool isIntOrFloat() const
Return true if this is an integer (of any signedness) or a float type.
Definition: Types.cpp:116
bool isSignlessIntOrIndexOrFloat() const
Return true if this is a signless integer, index, or float type.
Definition: Types.cpp:104
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
MLIRContext * getContext() const
Utility to get the associated MLIRContext that this value is defined in.
Definition: Value.h:132
Type getType() const
Return the type of this value.
Definition: Value.h:129
Location getLoc() const
Return the location of this value.
Definition: Value.cpp:26
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
bool hasElementwiseMappableTraits(Operation *op)
Together, Elementwise, Scalarizable, Vectorizable, and Tensorizable provide an easy way for scalar op...
Definition: Operation.cpp:1393
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:344
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given memref value.
Definition: MemRefOps.cpp:78
SmallVector< Value > replaceAndCastForOpIterArg(RewriterBase &rewriter, scf::ForOp forOp, OpOperand &operand, Value replacement, const ValueTypeCastFnTy &castFn)
Definition: SCF.cpp:781
Value makeArithReduction(OpBuilder &b, Location loc, CombiningKind kind, Value v1, Value acc, arith::FastMathFlagsAttr fastmath=nullptr, Value mask=nullptr)
Returns the result value of reducing two scalar/vector values with the corresponding arith operation.
void populateDropUnitDimWithShapeCastPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Collect a set of patterns that use vector.shape_cast to help fold unit dims.
bool isReductionIterator(Attribute attr)
Returns true if attr has "reduction" iterator type semantics.
Definition: VectorOps.h:152
void populateBreakDownVectorBitCastOpPatterns(RewritePatternSet &patterns, std::function< bool(BitCastOp)> controlFn=nullptr, PatternBenefit benefit=1)
Populate patterns with a pattern to break down 1-D vector.bitcast ops based on the destination vector...
auto getDims(VectorType vType)
Returns a range over the dims (size and scalability) of a VectorType.
Definition: VectorUtils.h:124
void populateElementwiseToVectorOpsPatterns(RewritePatternSet &patterns)
Collect a set of patterns that fold elementwise op on vectors to the vector dialect.
AffineMap getTransferMinorIdentityMap(ShapedType shapedType, VectorType vectorType)
Build the default minor identity map suitable for a vector transfer.
Definition: VectorOps.cpp:153
void populateShapeCastFoldingPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Collect a set of vector.shape_cast folding patterns.
void populateChainedVectorReductionFoldingPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Patterns that fold chained vector reductions.
void populateBubbleVectorBitCastOpPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Collect a set of patterns that bubble up/down bitcast ops.
bool isParallelIterator(Attribute attr)
Returns true if attr has "parallel" iterator type semantics.
Definition: VectorOps.h:147
void populateFoldArithExtensionPatterns(RewritePatternSet &patterns)
Collect a set of patterns that fold arithmetic extension on floating point into vector contract for t...
void populateVectorContractCanonicalizeMatmulToMMT(RewritePatternSet &patterns, std::function< LogicalResult(vector::ContractionOp)> constraint=[](vector::ContractionOp) { return success();}, PatternBenefit=1)
Canonicalization of a vector.contraction a, b, c with row-major matmul semantics to a contraction wit...
void populateSinkVectorOpsPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Patterns that remove redundant Vector Ops by re-ordering them with e.g.
void populateVectorMaskMaterializationPatterns(RewritePatternSet &patterns, bool force32BitVectorIndices, PatternBenefit benefit=1)
These patterns materialize masks for various vector ops such as transfers.
void populateVectorTransferCollapseInnerMostContiguousDimsPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Collect a set of patterns to reduce the rank of the operands of vector transfer ops to operate on the...
void populateVectorReductionToContractPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Collect patterns to convert reduction op to vector.contract and fold transpose/broadcast ops into the...
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:41
void populateBreakDownVectorReductionPatterns(RewritePatternSet &patterns, unsigned maxNumElementsToExtract=2, PatternBenefit benefit=1)
Patterns to break down vector reductions into a series of arith reductions over vector elements.
Include the generated interface declarations.
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
Definition: Matchers.h:490
const FrozenRewritePatternSet GreedyRewriteConfig bool * changed
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
Definition: AffineExpr.h:348
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
Definition: AffineMap.cpp:791
Value getValueOrCreateCastToIndexLike(OpBuilder &b, Location loc, Type targetType, Value value)
Create a cast from an index-like value (index or integer) to another index-like value.
Definition: Utils.cpp:120
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
const FrozenRewritePatternSet & patterns
detail::constant_float_predicate_matcher m_AnyZeroFloat()
Matches a constant scalar / vector splat / tensor splat float (both positive and negative) zero.
Definition: Matchers.h:399
AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims)
Drop the dims that are listed in unusedDims.
Definition: AffineMap.cpp:717
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
llvm::SmallBitVector getUnusedDimsBitVector(ArrayRef< AffineMap > maps)
Definition: AffineMap.cpp:930
detail::constant_op_matcher m_Constant()
Matches a constant foldable operation.
Definition: Matchers.h:369
OpInterfaceRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting a...
Definition: PatternMatch.h:379
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:368