MLIR  19.0.0git
VectorOps.cpp
Go to the documentation of this file.
1 //===- VectorOps.cpp - MLIR Vector Dialect Operations ---------------------===//
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 convenience types for working with super-vectorization
10 // operations, in particular super-vector loads and stores.
11 //
12 //===----------------------------------------------------------------------===//
13 
15 
24 #include "mlir/IR/AffineExpr.h"
25 #include "mlir/IR/AffineMap.h"
26 #include "mlir/IR/Builders.h"
28 #include "mlir/IR/BuiltinOps.h"
29 #include "mlir/IR/BuiltinTypes.h"
31 #include "mlir/IR/IRMapping.h"
33 #include "mlir/IR/PatternMatch.h"
34 #include "mlir/IR/TypeUtilities.h"
37 #include "mlir/Support/LLVM.h"
39 #include "llvm/ADT/ArrayRef.h"
40 #include "llvm/ADT/STLExtras.h"
41 #include "llvm/ADT/SmallVector.h"
42 #include "llvm/ADT/StringSet.h"
43 #include "llvm/ADT/TypeSwitch.h"
44 #include "llvm/ADT/bit.h"
45 
46 #include <cassert>
47 #include <cstdint>
48 #include <numeric>
49 
50 #include "mlir/Dialect/Vector/IR/VectorDialect.cpp.inc"
51 // Pull in all enum type and utility function definitions.
52 #include "mlir/Dialect/Vector/IR/VectorEnums.cpp.inc"
53 
54 using namespace mlir;
55 using namespace mlir::vector;
56 
57 /// Helper enum to classify mask value.
58 enum class MaskFormat {
59  AllTrue = 0,
60  AllFalse = 1,
61  Unknown = 2,
62 };
63 
64 /// Helper method to classify a mask value. Currently, the method
65 /// looks "under the hood" of a constant value with dense attributes
66 /// and a constant mask operation (since the client may be called at
67 /// various stages during progressive lowering).
69  if (auto c = mask.getDefiningOp<arith::ConstantOp>()) {
70  // Inspect constant dense values. We count up for bits that
71  // are set, count down for bits that are cleared, and bail
72  // when a mix is detected.
73  if (auto denseElts = llvm::dyn_cast<DenseIntElementsAttr>(c.getValue())) {
74  int64_t val = 0;
75  for (bool b : denseElts.getValues<bool>())
76  if (b && val >= 0)
77  val++;
78  else if (!b && val <= 0)
79  val--;
80  else
81  return MaskFormat::Unknown;
82  if (val > 0)
83  return MaskFormat::AllTrue;
84  if (val < 0)
85  return MaskFormat::AllFalse;
86  }
87  } else if (auto m = mask.getDefiningOp<ConstantMaskOp>()) {
88  // Inspect constant mask index. If the index exceeds the
89  // dimension size, all bits are set. If the index is zero
90  // or less, no bits are set.
91  ArrayAttr masks = m.getMaskDimSizes();
92  auto shape = m.getType().getShape();
93  bool allTrue = true;
94  bool allFalse = true;
95  for (auto [maskIdx, dimSize] : llvm::zip_equal(masks, shape)) {
96  int64_t i = llvm::cast<IntegerAttr>(maskIdx).getInt();
97  if (i < dimSize)
98  allTrue = false;
99  if (i > 0)
100  allFalse = false;
101  }
102  if (allTrue)
103  return MaskFormat::AllTrue;
104  if (allFalse)
105  return MaskFormat::AllFalse;
106  } else if (auto m = mask.getDefiningOp<CreateMaskOp>()) {
107  // Finds all-false create_masks. An all-true create_mask requires all
108  // dims to be constants, so that'll be folded to a constant_mask, then
109  // detected in the constant_mask case.
110  auto maskOperands = m.getOperands();
111  for (Value operand : maskOperands) {
112  if (auto constantOp = operand.getDefiningOp<arith::ConstantOp>()) {
113  int64_t dimSize =
114  llvm::cast<IntegerAttr>(constantOp.getValue()).getInt();
115  if (dimSize <= 0)
116  return MaskFormat::AllFalse;
117  }
118  }
119  return MaskFormat::Unknown;
120  }
121  return MaskFormat::Unknown;
122 }
123 
124 /// Default callback to build a region with a 'vector.yield' terminator with no
125 /// arguments.
127  builder.create<vector::YieldOp>(loc);
128 }
129 
130 // Helper for verifying combining kinds in contractions and reductions.
131 static bool isSupportedCombiningKind(CombiningKind combiningKind,
132  Type elementType) {
133  switch (combiningKind) {
134  case CombiningKind::ADD:
135  case CombiningKind::MUL:
136  return elementType.isIntOrIndexOrFloat();
138  case CombiningKind::MINSI:
139  case CombiningKind::MAXUI:
140  case CombiningKind::MAXSI:
141  case CombiningKind::AND:
142  case CombiningKind::OR:
143  case CombiningKind::XOR:
144  return elementType.isIntOrIndex();
145  case CombiningKind::MINNUMF:
146  case CombiningKind::MAXNUMF:
147  case CombiningKind::MINIMUMF:
148  case CombiningKind::MAXIMUMF:
149  return llvm::isa<FloatType>(elementType);
150  }
151  return false;
152 }
153 
155  VectorType vectorType) {
156  int64_t elementVectorRank = 0;
157  VectorType elementVectorType =
158  llvm::dyn_cast<VectorType>(shapedType.getElementType());
159  if (elementVectorType)
160  elementVectorRank += elementVectorType.getRank();
161  // 0-d transfers are to/from tensor<t>/memref<t> and vector<1xt>.
162  // TODO: replace once we have 0-d vectors.
163  if (shapedType.getRank() == 0 &&
164  vectorType.getShape() == ArrayRef<int64_t>{1})
165  return AffineMap::get(
166  /*numDims=*/0, /*numSymbols=*/0,
167  getAffineConstantExpr(0, shapedType.getContext()));
169  shapedType.getRank(), vectorType.getRank() - elementVectorRank,
170  shapedType.getContext());
171 }
172 
173 /// Check if `write` is of a constant splat and the masked `read` is padded with
174 /// the same splat value -- meaning it could be the same value as the initial
175 /// constant splat.
176 static bool isSplatWriteConsistentWithMaskedRead(vector::TransferWriteOp write,
177  vector::TransferReadOp read) {
178  auto readMask = read.getMask();
179  auto writeMask = write.getMask();
180  // Check if the masks are consistent. The splat value could be the same if the
181  // read is masked (and padded with the splat value), and the write is unmasked
182  // or has the same mask. Note this does not allow the case where the write is
183  // masked and the read is unmasked, as then the read could be of more elements
184  // than the write (which may not be the same value).
185  bool couldBeSameSplat = readMask && (!writeMask || writeMask == readMask);
186  if (!couldBeSameSplat)
187  return false;
188  // Check for constant splat (as the source of the write).
189  DenseElementsAttr splatAttr;
190  if (!matchPattern(write.getVector(),
191  m_Constant<DenseElementsAttr>(&splatAttr)) ||
192  !splatAttr.isSplat()) {
193  return false;
194  }
195  // The padding of the read and the constant splat value must be the same.
196  Attribute padAttr;
197  if (!matchPattern(read.getPadding(), m_Constant(&padAttr)))
198  return false;
199  return padAttr == splatAttr.getSplatValue<Attribute>();
200 }
201 
202 bool mlir::vector::checkSameValueRAW(vector::TransferWriteOp defWrite,
203  vector::TransferReadOp read) {
204  return !defWrite.hasOutOfBoundsDim() &&
205  defWrite.getIndices() == read.getIndices() &&
206  defWrite.getVectorType() == read.getVectorType() &&
207  defWrite.getPermutationMap() == read.getPermutationMap() &&
208  ((!defWrite.getMask() && !read.getMask()) ||
209  isSplatWriteConsistentWithMaskedRead(defWrite, read));
210 }
211 
212 bool mlir::vector::checkSameValueWAW(vector::TransferWriteOp write,
213  vector::TransferWriteOp priorWrite) {
214  return priorWrite.getIndices() == write.getIndices() &&
215  priorWrite.getMask() == write.getMask() &&
216  priorWrite.getVectorType() == write.getVectorType() &&
217  priorWrite.getPermutationMap() == write.getPermutationMap();
218 }
219 
221  VectorTransferOpInterface transferA, VectorTransferOpInterface transferB,
222  bool testDynamicValueUsingBounds) {
223  // For simplicity only look at transfer of same type.
224  if (transferA.getVectorType() != transferB.getVectorType())
225  return false;
226  unsigned rankOffset = transferA.getLeadingShapedRank();
227  for (unsigned i = 0, e = transferA.getIndices().size(); i < e; i++) {
228  Value indexA = transferA.getIndices()[i];
229  Value indexB = transferB.getIndices()[i];
230  std::optional<int64_t> cstIndexA = getConstantIntValue(indexA);
231  std::optional<int64_t> cstIndexB = getConstantIntValue(indexB);
232 
233  if (i < rankOffset) {
234  // For leading dimensions, if we can prove that index are different we
235  // know we are accessing disjoint slices.
236  if (cstIndexA.has_value() && cstIndexB.has_value()) {
237  if (*cstIndexA != *cstIndexB)
238  return true;
239  continue;
240  }
241  if (testDynamicValueUsingBounds) {
242  // First try to see if we can fully compose and simplify the affine
243  // expression as a fast track.
244  FailureOr<uint64_t> delta =
246  if (succeeded(delta) && *delta != 0)
247  return true;
248 
249  FailureOr<bool> testEqual =
250  ValueBoundsConstraintSet::areEqual(indexA, indexB);
251  if (succeeded(testEqual) && !testEqual.value())
252  return true;
253  }
254  } else {
255  // For this dimension, we slice a part of the memref we need to make sure
256  // the intervals accessed don't overlap.
257  int64_t vectorDim = transferA.getVectorType().getDimSize(i - rankOffset);
258  if (cstIndexA.has_value() && cstIndexB.has_value()) {
259  int64_t distance = std::abs(*cstIndexA - *cstIndexB);
260  if (distance >= vectorDim)
261  return true;
262  continue;
263  }
264  if (testDynamicValueUsingBounds) {
265  // First try to see if we can fully compose and simplify the affine
266  // expression as a fast track.
267  FailureOr<int64_t> delta =
269  if (succeeded(delta) && std::abs(*delta) >= vectorDim)
270  return true;
271 
272  FailureOr<int64_t> computeDelta =
274  if (succeeded(computeDelta)) {
275  if (std::abs(computeDelta.value()) >= vectorDim)
276  return true;
277  }
278  }
279  }
280  }
281  return false;
282 }
283 
284 bool mlir::vector::isDisjointTransferSet(VectorTransferOpInterface transferA,
285  VectorTransferOpInterface transferB,
286  bool testDynamicValueUsingBounds) {
287  if (transferA.getSource() != transferB.getSource())
288  return false;
289  return isDisjointTransferIndices(transferA, transferB,
290  testDynamicValueUsingBounds);
291 }
292 
293 // Helper to iterate over n-D vector slice elements. Calculate the next
294 // `position` in the n-D vector of size `shape`, applying an offset `offsets`.
295 // Modifies the `position` in place. Returns a failure when `position` becomes
296 // the end position.
298  ArrayRef<int64_t> shape,
299  ArrayRef<int64_t> offsets) {
300  for (auto [posInDim, dimSize, offsetInDim] :
301  llvm::reverse(llvm::zip_equal(position, shape, offsets))) {
302  ++posInDim;
303  if (posInDim < dimSize + offsetInDim)
304  return success();
305 
306  // Carry the overflow to the next loop iteration.
307  posInDim = offsetInDim;
308  }
309 
310  return failure();
311 }
312 
313 /// Returns the integer numbers in `values`. `values` are expected to be
314 /// constant operations.
317  llvm::transform(values, std::back_inserter(ints), [](Value value) {
318  auto constOp = value.getDefiningOp<arith::ConstantIndexOp>();
319  assert(constOp && "Unexpected non-constant index");
320  return constOp.value();
321  });
322  return ints;
323 }
324 
325 /// Returns the integer numbers in `foldResults`. `foldResults` are expected to
326 /// be constant operations.
329  llvm::transform(
330  foldResults, std::back_inserter(ints), [](OpFoldResult foldResult) {
331  assert(foldResult.is<Attribute>() && "Unexpected non-constant index");
332  return cast<IntegerAttr>(foldResult.get<Attribute>()).getInt();
333  });
334  return ints;
335 }
336 
337 /// Convert `foldResults` into Values. Integer attributes are converted to
338 /// constant op.
340  ArrayRef<OpFoldResult> foldResults) {
341  SmallVector<Value> values;
342  llvm::transform(foldResults, std::back_inserter(values),
343  [&](OpFoldResult foldResult) {
344  if (auto attr = foldResult.dyn_cast<Attribute>())
345  return builder
347  loc, cast<IntegerAttr>(attr).getInt())
348  .getResult();
349 
350  return foldResult.get<Value>();
351  });
352  return values;
353 }
354 
355 //===----------------------------------------------------------------------===//
356 // CombiningKindAttr
357 //===----------------------------------------------------------------------===//
358 
359 namespace mlir {
360 namespace vector {
361 namespace detail {
363  using KeyTy = uint64_t;
364 
365  BitmaskEnumStorage(KeyTy val) : value(val) {}
366 
367  bool operator==(const KeyTy &key) const { return value == key; }
368 
370  const KeyTy &key) {
371  return new (allocator.allocate<BitmaskEnumStorage>())
372  BitmaskEnumStorage(key);
373  }
374 
375  KeyTy value = 0;
376 };
377 } // namespace detail
378 } // namespace vector
379 } // namespace mlir
380 
381 //===----------------------------------------------------------------------===//
382 // VectorDialect
383 //===----------------------------------------------------------------------===//
384 
385 namespace {
386 /// This class defines the interface for handling inlining with vector dialect
387 /// operations.
388 struct VectorInlinerInterface : public DialectInlinerInterface {
390 
391  /// All vector dialect ops can be inlined.
392  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
393  return true;
394  }
395 };
396 } // namespace
397 
398 void VectorDialect::initialize() {
399  addAttributes<
400 #define GET_ATTRDEF_LIST
401 #include "mlir/Dialect/Vector/IR/VectorAttributes.cpp.inc"
402  >();
403 
404  addOperations<
405 #define GET_OP_LIST
406 #include "mlir/Dialect/Vector/IR/VectorOps.cpp.inc"
407  >();
408 
409  addInterfaces<VectorInlinerInterface>();
410 
411  declarePromisedInterfaces<bufferization::BufferizableOpInterface,
412  TransferReadOp, TransferWriteOp, GatherOp, MaskOp,
413  YieldOp>();
414  declarePromisedInterfaces<SubsetOpInterface, TransferReadOp,
415  TransferWriteOp>();
416  declarePromisedInterface<SubsetExtractionOpInterface, TransferReadOp>();
417  declarePromisedInterface<SubsetInsertionOpInterface, TransferWriteOp>();
418 }
419 
420 /// Materialize a single constant operation from a given attribute value with
421 /// the desired resultant type.
423  Attribute value, Type type,
424  Location loc) {
425  return arith::ConstantOp::materialize(builder, value, type, loc);
426 }
427 
429  return builder.getIntegerType(64);
430 }
431 
433  ArrayRef<int64_t> values) {
434  return builder.getI64ArrayAttr(values);
435 }
436 
437 //===----------------------------------------------------------------------===//
438 // MultiDimReductionOp
439 //===----------------------------------------------------------------------===//
440 
441 void vector::MultiDimReductionOp::build(OpBuilder &builder,
442  OperationState &result, Value source,
443  Value acc, ArrayRef<bool> reductionMask,
444  CombiningKind kind) {
445  SmallVector<int64_t> reductionDims;
446  for (const auto &en : llvm::enumerate(reductionMask))
447  if (en.value())
448  reductionDims.push_back(en.index());
449  build(builder, result, kind, source, acc,
450  builder.getI64ArrayAttr(reductionDims));
451 }
452 
453 OpFoldResult MultiDimReductionOp::fold(FoldAdaptor adaptor) {
454  // Single parallel dim, this is a noop.
455  if (getSourceVectorType().getRank() == 1 && !isReducedDim(0))
456  return getSource();
457  return {};
458 }
459 
460 std::optional<SmallVector<int64_t, 4>>
461 MultiDimReductionOp::getShapeForUnroll() {
462  return llvm::to_vector<4>(getSourceVectorType().getShape());
463 }
464 
466  SmallVector<int64_t> targetShape;
467  SmallVector<bool> scalableDims;
468  Type inferredReturnType;
469  auto sourceScalableDims = getSourceVectorType().getScalableDims();
470  for (auto it : llvm::enumerate(getSourceVectorType().getShape()))
471  if (!llvm::any_of(getReductionDims().getValue(), [&](Attribute attr) {
472  return llvm::cast<IntegerAttr>(attr).getValue() == it.index();
473  })) {
474  targetShape.push_back(it.value());
475  scalableDims.push_back(sourceScalableDims[it.index()]);
476  }
477  // TODO: update to also allow 0-d vectors when available.
478  if (targetShape.empty())
479  inferredReturnType = getSourceVectorType().getElementType();
480  else
481  inferredReturnType = VectorType::get(
482  targetShape, getSourceVectorType().getElementType(), scalableDims);
483  if (getType() != inferredReturnType)
484  return emitOpError() << "destination type " << getType()
485  << " is incompatible with source type "
486  << getSourceVectorType();
487 
488  return success();
489 }
490 
491 /// Returns the mask type expected by this operation.
492 Type MultiDimReductionOp::getExpectedMaskType() {
493  auto vecType = getSourceVectorType();
494  return VectorType::get(vecType.getShape(),
495  IntegerType::get(vecType.getContext(), /*width=*/1),
496  vecType.getScalableDims());
497 }
498 
499 namespace {
500 // Only unit dimensions that are being reduced are folded. If the dimension is
501 // unit, but not reduced, it is not folded, thereby keeping the output type the
502 // same. If not all dimensions which are reduced are of unit dimension, this
503 // transformation does nothing. This is just a generalization of
504 // ElideSingleElementReduction for ReduceOp.
505 struct ElideUnitDimsInMultiDimReduction
506  : public OpRewritePattern<MultiDimReductionOp> {
508 
509  LogicalResult matchAndRewrite(MultiDimReductionOp reductionOp,
510  PatternRewriter &rewriter) const override {
511  ArrayRef<int64_t> shape = reductionOp.getSourceVectorType().getShape();
512  for (const auto &dim : enumerate(shape)) {
513  if (reductionOp.isReducedDim(dim.index()) && dim.value() != 1)
514  return failure();
515  }
516 
517  // Vector mask setup.
518  OpBuilder::InsertionGuard guard(rewriter);
519  Operation *rootOp;
520  Value mask;
521  if (reductionOp.isMasked()) {
522  rewriter.setInsertionPoint(reductionOp.getMaskingOp());
523  rootOp = reductionOp.getMaskingOp();
524  mask = reductionOp.getMaskingOp().getMask();
525  } else {
526  rootOp = reductionOp;
527  }
528 
529  Location loc = reductionOp.getLoc();
530  Value acc = reductionOp.getAcc();
531  Value cast;
532  if (auto dstVecType = dyn_cast<VectorType>(reductionOp.getDestType())) {
533  if (mask) {
534  VectorType newMaskType =
535  VectorType::get(dstVecType.getShape(), rewriter.getI1Type(),
536  dstVecType.getScalableDims());
537  mask = rewriter.create<vector::ShapeCastOp>(loc, newMaskType, mask);
538  }
539  cast = rewriter.create<vector::ShapeCastOp>(
540  loc, reductionOp.getDestType(), reductionOp.getSource());
541  } else {
542  // This means we are reducing all the dimensions, and all reduction
543  // dimensions are of size 1. So a simple extraction would do.
544  SmallVector<int64_t> zeroIdx(shape.size(), 0);
545  if (mask)
546  mask = rewriter.create<vector::ExtractOp>(loc, mask, zeroIdx);
547  cast = rewriter.create<vector::ExtractOp>(loc, reductionOp.getSource(),
548  zeroIdx);
549  }
550 
551  Value result =
552  vector::makeArithReduction(rewriter, loc, reductionOp.getKind(), acc,
553  cast, /*fastmath=*/nullptr, mask);
554  rewriter.replaceOp(rootOp, result);
555  return success();
556  }
557 };
558 } // namespace
559 
560 void MultiDimReductionOp::getCanonicalizationPatterns(
561  RewritePatternSet &results, MLIRContext *context) {
562  results.add<ElideUnitDimsInMultiDimReduction>(context);
563 }
564 
565 //===----------------------------------------------------------------------===//
566 // ReductionOp
567 //===----------------------------------------------------------------------===//
568 
569 void vector::ReductionOp::build(OpBuilder &builder, OperationState &result,
570  CombiningKind kind, Value vector,
571  arith::FastMathFlags fastMathFlags) {
572  build(builder, result, kind, vector, /*acc=*/Value(), fastMathFlags);
573 }
574 
575 void vector::ReductionOp::build(OpBuilder &builder, OperationState &result,
576  CombiningKind kind, Value vector, Value acc,
577  arith::FastMathFlags fastMathFlags) {
578  build(builder, result,
579  llvm::cast<VectorType>(vector.getType()).getElementType(), kind, vector,
580  acc, fastMathFlags);
581 }
582 
584  // Verify for 0-D and 1-D vector.
585  int64_t rank = getSourceVectorType().getRank();
586  if (rank > 1)
587  return emitOpError("unsupported reduction rank: ") << rank;
588 
589  // Verify supported reduction kind.
590  Type eltType = getDest().getType();
591  if (!isSupportedCombiningKind(getKind(), eltType))
592  return emitOpError("unsupported reduction type '")
593  << eltType << "' for kind '" << stringifyCombiningKind(getKind())
594  << "'";
595 
596  return success();
597 }
598 
599 // MaskableOpInterface methods.
600 
601 /// Returns the mask type expected by this operation.
602 Type ReductionOp::getExpectedMaskType() {
603  auto vecType = getSourceVectorType();
604  return VectorType::get(vecType.getShape(),
605  IntegerType::get(vecType.getContext(), /*width=*/1),
606  vecType.getScalableDims());
607 }
608 
610  OpBuilder &builder, Location loc,
611  Value vector) {
612  switch (op) {
613  case arith::AtomicRMWKind::addf:
614  case arith::AtomicRMWKind::addi:
615  return builder.create<vector::ReductionOp>(vector.getLoc(),
616  CombiningKind::ADD, vector);
617  case arith::AtomicRMWKind::mulf:
618  case arith::AtomicRMWKind::muli:
619  return builder.create<vector::ReductionOp>(vector.getLoc(),
620  CombiningKind::MUL, vector);
621  case arith::AtomicRMWKind::minimumf:
622  return builder.create<vector::ReductionOp>(vector.getLoc(),
623  CombiningKind::MINIMUMF, vector);
624  case arith::AtomicRMWKind::mins:
625  return builder.create<vector::ReductionOp>(vector.getLoc(),
626  CombiningKind::MINSI, vector);
627  case arith::AtomicRMWKind::minu:
628  return builder.create<vector::ReductionOp>(vector.getLoc(),
629  CombiningKind::MINUI, vector);
630  case arith::AtomicRMWKind::maximumf:
631  return builder.create<vector::ReductionOp>(vector.getLoc(),
632  CombiningKind::MAXIMUMF, vector);
633  case arith::AtomicRMWKind::maxs:
634  return builder.create<vector::ReductionOp>(vector.getLoc(),
635  CombiningKind::MAXSI, vector);
636  case arith::AtomicRMWKind::maxu:
637  return builder.create<vector::ReductionOp>(vector.getLoc(),
638  CombiningKind::MAXUI, vector);
639  case arith::AtomicRMWKind::andi:
640  return builder.create<vector::ReductionOp>(vector.getLoc(),
641  CombiningKind::AND, vector);
642  case arith::AtomicRMWKind::ori:
643  return builder.create<vector::ReductionOp>(vector.getLoc(),
644  CombiningKind::OR, vector);
645  // TODO: Add remaining reduction operations.
646  default:
647  (void)emitOptionalError(loc, "Reduction operation type not supported");
648  break;
649  }
650  return nullptr;
651 }
652 
653 std::optional<SmallVector<int64_t, 4>> ReductionOp::getShapeForUnroll() {
654  return llvm::to_vector<4>(getSourceVectorType().getShape());
655 }
656 
657 namespace {
658 struct ElideSingleElementReduction : public OpRewritePattern<ReductionOp> {
660 
661  LogicalResult matchAndRewrite(ReductionOp reductionOp,
662  PatternRewriter &rewriter) const override {
663  // Vector mask setup.
664  OpBuilder::InsertionGuard guard(rewriter);
665  auto maskableOp =
666  cast<vector::MaskableOpInterface>(reductionOp.getOperation());
667  Operation *rootOp;
668  Value mask;
669  if (maskableOp.isMasked()) {
670  rewriter.setInsertionPoint(maskableOp.getMaskingOp());
671  rootOp = maskableOp.getMaskingOp();
672  mask = maskableOp.getMaskingOp().getMask();
673  } else {
674  rootOp = reductionOp;
675  }
676 
677  auto vectorType = reductionOp.getSourceVectorType();
678  if (vectorType.getRank() != 0 && vectorType.getDimSize(0) != 1)
679  return failure();
680 
681  Location loc = reductionOp.getLoc();
682  Value result;
683  if (vectorType.getRank() == 0) {
684  if (mask)
685  mask = rewriter.create<ExtractElementOp>(loc, mask);
686  result = rewriter.create<ExtractElementOp>(loc, reductionOp.getVector());
687  } else {
688  if (mask)
689  mask = rewriter.create<ExtractOp>(loc, mask, 0);
690  result = rewriter.create<ExtractOp>(loc, reductionOp.getVector(), 0);
691  }
692 
693  if (Value acc = reductionOp.getAcc())
694  result = vector::makeArithReduction(rewriter, loc, reductionOp.getKind(),
695  result, acc,
696  reductionOp.getFastmathAttr(), mask);
697 
698  rewriter.replaceOp(rootOp, result);
699  return success();
700  }
701 };
702 } // namespace
703 
704 void ReductionOp::getCanonicalizationPatterns(RewritePatternSet &results,
705  MLIRContext *context) {
706  results.add<ElideSingleElementReduction>(context);
707 }
708 
709 //===----------------------------------------------------------------------===//
710 // ContractionOp
711 //===----------------------------------------------------------------------===//
712 
713 void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
714  Value lhs, Value rhs, Value acc,
715  ArrayRef<ArrayRef<AffineExpr>> indexingExprs,
716  ArrayRef<IteratorType> iteratorTypes) {
717  result.addOperands({lhs, rhs, acc});
718  result.addTypes(acc.getType());
719  result.addAttribute(
720  getIndexingMapsAttrName(result.name),
721  builder.getAffineMapArrayAttr(
722  AffineMap::inferFromExprList(indexingExprs, builder.getContext())));
723  result.addAttribute(
724  getIteratorTypesAttrName(result.name),
725  builder.getArrayAttr(llvm::to_vector(llvm::map_range(
726  iteratorTypes, [&](IteratorType t) -> mlir::Attribute {
727  return IteratorTypeAttr::get(builder.getContext(), t);
728  }))));
729 }
730 
731 void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
732  Value lhs, Value rhs, Value acc,
733  ArrayAttr indexingMaps,
734  ArrayAttr iteratorTypes) {
735  build(builder, result, lhs, rhs, acc, indexingMaps, iteratorTypes,
736  ContractionOp::getDefaultKind());
737 }
738 
739 void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
740  Value lhs, Value rhs, Value acc,
741  ArrayAttr indexingMaps,
742  ArrayAttr iteratorTypes, CombiningKind kind) {
743  result.addOperands({lhs, rhs, acc});
744  result.addTypes(acc.getType());
745  result.addAttribute(getIndexingMapsAttrName(result.name), indexingMaps);
746  result.addAttribute(getIteratorTypesAttrName(result.name), iteratorTypes);
747  result.addAttribute(getKindAttrName(result.name),
748  CombiningKindAttr::get(builder.getContext(), kind));
749 }
750 
756  SmallVector<Type, 2> types;
757  Type resultType;
758  auto loc = parser.getCurrentLocation();
759  DictionaryAttr dictAttr;
760  // TODO: Unify linalg op attribute parsing.
761  if (parser.parseAttribute(dictAttr) || parser.parseOperand(lhsInfo) ||
762  parser.parseComma() || parser.parseOperand(rhsInfo) ||
763  parser.parseComma() || parser.parseOperand(accInfo) ||
764  parser.parseTrailingOperandList(masksInfo) ||
765  parser.parseOptionalAttrDict(result.attributes) ||
766  parser.parseColonTypeList(types) ||
767  parser.parseKeywordType("into", resultType) ||
768  parser.resolveOperand(lhsInfo, types[0], result.operands) ||
769  parser.resolveOperand(rhsInfo, types[1], result.operands) ||
770  parser.resolveOperand(accInfo, resultType, result.operands) ||
771  parser.addTypeToList(resultType, result.types))
772  return failure();
773  result.attributes.append(dictAttr.getValue().begin(),
774  dictAttr.getValue().end());
775 
776  // Convert array of string into an array of IteratyType enums. This is needed,
777  // because tests still use the old format when 'iterator_types' attribute is
778  // represented as an array of strings.
779  // TODO: Remove this conversion once tests are fixed.
780  ArrayAttr iteratorTypes = llvm::cast<ArrayAttr>(
781  result.attributes.get(getIteratorTypesAttrName(result.name)));
782 
783  SmallVector<Attribute> iteratorTypeAttrs;
784 
785  for (StringRef s : iteratorTypes.getAsValueRange<StringAttr>()) {
786  auto maybeIteratorType = symbolizeIteratorType(s);
787  if (!maybeIteratorType.has_value())
788  return parser.emitError(loc) << "unexpected iterator_type (" << s << ")";
789 
790  iteratorTypeAttrs.push_back(
791  IteratorTypeAttr::get(parser.getContext(), maybeIteratorType.value()));
792  }
793  result.attributes.set(getIteratorTypesAttrName(result.name),
794  parser.getBuilder().getArrayAttr(iteratorTypeAttrs));
795 
796  if (!result.attributes.get(getKindAttrName(result.name))) {
797  result.addAttribute(
798  getKindAttrName(result.name),
800  ContractionOp::getDefaultKind()));
801  }
802  if (masksInfo.empty())
803  return success();
804  if (masksInfo.size() != 2)
805  return parser.emitError(parser.getNameLoc(),
806  "expected zero or exactly 2 vector mask operands");
807  auto lhsType = llvm::cast<VectorType>(types[0]);
808  auto rhsType = llvm::cast<VectorType>(types[1]);
809  auto maskElementType = parser.getBuilder().getI1Type();
810  std::array<VectorType, 2> maskTypes = {
811  VectorType::Builder(lhsType).setElementType(maskElementType),
812  VectorType::Builder(rhsType).setElementType(maskElementType)};
813  if (parser.resolveOperands(masksInfo, maskTypes, loc, result.operands))
814  return failure();
815  return success();
816 }
817 
819  // TODO: Unify printing code with linalg ops.
820  auto attrNames = getTraitAttrNames();
821  llvm::StringSet<> traitAttrsSet;
822  traitAttrsSet.insert(attrNames.begin(), attrNames.end());
824  for (auto attr : (*this)->getAttrs()) {
825  if (attr.getName() == getIteratorTypesAttrName()) {
826  auto iteratorTypes =
827  llvm::cast<ArrayAttr>(attr.getValue())
828  .getAsValueRange<IteratorTypeAttr, IteratorType>();
829  // Convert IteratorType enums into the string representation. This is
830  // needed, because tests still use the old format when 'iterator_types'
831  // attribute is represented as an array of strings.
832  // TODO: Remove this conversion once tests are fixed.
833  SmallVector<Attribute> iteratorTypeNames = llvm::to_vector(
834  llvm::map_range(iteratorTypes, [&](IteratorType t) -> Attribute {
835  return StringAttr::get(getContext(), stringifyIteratorType(t));
836  }));
837 
838  attrs.emplace_back(getIteratorTypesAttrName(),
839  ArrayAttr::get(getContext(), iteratorTypeNames));
840  } else if (traitAttrsSet.count(attr.getName().strref()) > 0)
841  attrs.push_back(attr);
842  }
843 
844  auto dictAttr = DictionaryAttr::get(getContext(), attrs);
845  p << " " << dictAttr << " " << getLhs() << ", ";
846  p << getRhs() << ", " << getAcc();
847 
848  p.printOptionalAttrDict((*this)->getAttrs(), attrNames);
849  p << " : " << getLhs().getType() << ", " << getRhs().getType() << " into "
850  << getResultType();
851 }
852 
853 static bool verifyDimMap(VectorType lhsType, VectorType rhsType,
854  const std::vector<std::pair<int64_t, int64_t>> &map) {
855  for (auto &dimPair : map) {
856  if (dimPair.first < 0 || dimPair.first >= lhsType.getRank() ||
857  dimPair.second < 0 || dimPair.second >= rhsType.getRank() ||
858  lhsType.getDimSize(dimPair.first) != rhsType.getDimSize(dimPair.second))
859  return false;
860  }
861  return true;
862 }
863 
865  ContractionOp op, VectorType lhsType, VectorType rhsType, Type accType,
866  Type resType,
867  const std::vector<std::pair<int64_t, int64_t>> &contractingDimMap,
868  const std::vector<std::pair<int64_t, int64_t>> &batchDimMap) {
869  DenseSet<int64_t> lhsContractingDimSet;
870  DenseSet<int64_t> rhsContractingDimSet;
871  for (auto &dimPair : contractingDimMap) {
872  lhsContractingDimSet.insert(dimPair.first);
873  rhsContractingDimSet.insert(dimPair.second);
874  }
875  DenseSet<int64_t> rhsBatchDimSet;
876  for (auto &dimPair : batchDimMap)
877  rhsBatchDimSet.insert(dimPair.second);
878 
879  // Add free and batch dimensions from 'lhsType' to 'expectedResultDims'.
880  SmallVector<int64_t, 4> expectedResultDims;
881  for (int64_t i = 0, e = lhsType.getRank(); i < e; ++i) {
882  if (lhsContractingDimSet.count(i) > 0)
883  continue;
884  expectedResultDims.push_back(lhsType.getDimSize(i));
885  }
886 
887  // Add free dimensions from 'rhsType' to 'expectedResultDims'.
888  for (int64_t i = 0, e = rhsType.getRank(); i < e; ++i) {
889  if (rhsContractingDimSet.count(i) > 0 || rhsBatchDimSet.count(i) > 0)
890  continue;
891  expectedResultDims.push_back(rhsType.getDimSize(i));
892  }
893 
894  // Verify 'expectedResultDims'.
895  if (expectedResultDims.empty()) {
896  // No batch or free dimension implies a scalar result.
897  if (llvm::isa<VectorType>(resType) || llvm::isa<VectorType>(accType))
898  return op.emitOpError("invalid accumulator/result vector shape");
899  } else {
900  // At least one batch or free dimension implies a vector result.
901  auto resVectorType = llvm::dyn_cast<VectorType>(resType);
902  auto accVectorType = llvm::dyn_cast<VectorType>(accType);
903  if (!resVectorType || !accVectorType)
904  return op.emitOpError("invalid accumulator/result vector shape");
905 
906  // Infer expected result vector type. Lhs + rhs map and lhs + rhs vector
907  // types fully define the result vector type. This assumes the affine maps
908  // are well-formed, which must have been verified already.
909  MLIRContext *ctx = op.getContext();
910  AffineMap lhsMap = op.getIndexingMapsArray()[0];
911  AffineMap rhsMap = op.getIndexingMapsArray()[1];
912  if (getUnusedDimsBitVector({lhsMap, rhsMap}).any())
913  return op.emitOpError(
914  "expected all dimensions to be either a LHS or a RHS dimension");
915  SmallVector<AffineExpr, 4> extents(lhsMap.getNumInputs());
916  for (auto pair :
917  {std::make_pair(lhsType, lhsMap), std::make_pair(rhsType, rhsMap)}) {
918  VectorType v = pair.first;
919  auto map = pair.second;
920  for (unsigned idx = 0, e = v.getRank(); idx < e; ++idx) {
921  unsigned pos = map.getDimPosition(idx);
922  if (!extents[pos])
923  extents[pos] = getAffineConstantExpr(v.getShape()[idx], ctx);
924  }
925  }
926  if (!llvm::all_of(extents, [](AffineExpr e) { return e; }))
927  return op.emitOpError("expected all dimensions to get an extent as "
928  "either a LHS or a RHS dimension");
929 
930  AffineMap resMap = op.getIndexingMapsArray()[2];
931  auto extentsMap = AffineMap::get(/*dimCount=*/extents.size(),
932  /*symbolCount=*/0, extents, ctx);
933  // Compose the resMap with the extentsMap, which is a constant map.
934  AffineMap expectedMap = simplifyAffineMap(resMap.compose(extentsMap));
935  assert(llvm::all_of(expectedMap.getResults(),
936  llvm::IsaPred<AffineConstantExpr>) &&
937  "expected constant extent along all dimensions.");
938  // Extract the expected shape and build the type.
939  auto expectedShape = llvm::to_vector<4>(
940  llvm::map_range(expectedMap.getResults(), [](AffineExpr e) {
941  return cast<AffineConstantExpr>(e).getValue();
942  }));
943  auto expected =
944  VectorType::get(expectedShape, resVectorType.getElementType(),
945  resVectorType.getScalableDims());
946  if (resVectorType != expected || accVectorType != expected)
947  return op.emitOpError(
948  "invalid accumulator/result vector shape, expected: ")
949  << expected;
950  }
951  return success();
952 }
953 
955  VectorType lhsType = getLhsType();
956  VectorType rhsType = getRhsType();
957  Type accType = getAccType();
958  Type resType = getResultType();
959 
960  if (llvm::isa<IntegerType>(lhsType.getElementType())) {
961  if (!lhsType.getElementType().isSignlessInteger())
962  return emitOpError("only supports signless integer types");
963  }
964 
965  // Verify that an indexing map was specified for each vector operand.
966  if (getIndexingMapsArray().size() != 3)
967  return emitOpError("expected an indexing map for each vector operand");
968 
969  // Verify that each index map has 'numIterators' inputs, no symbols, and
970  // that the number of map outputs equals the rank of its associated
971  // vector operand.
972  unsigned numIterators = getIteratorTypes().getValue().size();
973  for (const auto &it : llvm::enumerate(getIndexingMapsArray())) {
974  auto index = it.index();
975  auto map = it.value();
976  if (map.getNumSymbols() != 0)
977  return emitOpError("expected indexing map ")
978  << index << " to have no symbols";
979  auto vectorType = llvm::dyn_cast<VectorType>(getOperand(index).getType());
980  unsigned rank = vectorType ? vectorType.getShape().size() : 0;
981  // Verify that the map has the right number of inputs, outputs, and indices.
982  // This also correctly accounts for (..) -> () for rank-0 results.
983  if (map.getNumDims() != numIterators)
984  return emitOpError("expected indexing map ")
985  << index << " to have " << numIterators << " number of inputs";
986  if (map.getNumResults() != rank)
987  return emitOpError("expected indexing map ")
988  << index << " to have " << rank << " number of outputs";
989  if (!map.isProjectedPermutation())
990  return emitOpError("expected indexing map ")
991  << index << " to be a projected permutation of its inputs";
992  }
993 
994  auto contractingDimMap = getContractingDimMap();
995  auto batchDimMap = getBatchDimMap();
996 
997  // Verify at least one contracting dimension pair was specified.
998  if (contractingDimMap.empty())
999  return emitOpError("expected at least one contracting dimension pair");
1000 
1001  // Verify contracting dimension map was properly constructed.
1002  if (!verifyDimMap(lhsType, rhsType, contractingDimMap))
1003  return emitOpError("invalid contracting dimension map");
1004 
1005  // Verify batch dimension map was properly constructed.
1006  if (!verifyDimMap(lhsType, rhsType, batchDimMap))
1007  return emitOpError("invalid batch dimension map");
1008 
1009  // Verify 'accType' and 'resType' shape.
1010  if (failed(verifyOutputShape(*this, lhsType, rhsType, accType, resType,
1011  contractingDimMap, batchDimMap)))
1012  return failure();
1013 
1014  // Verify supported combining kind.
1015  auto vectorType = llvm::dyn_cast<VectorType>(resType);
1016  auto elementType = vectorType ? vectorType.getElementType() : resType;
1017  if (!isSupportedCombiningKind(getKind(), elementType))
1018  return emitOpError("unsupported contraction type");
1019 
1020  return success();
1021 }
1022 
1023 // MaskableOpInterface methods.
1024 
1025 /// Returns the mask type expected by this operation. Mostly used for
1026 /// verification purposes. It requires the operation to be vectorized."
1027 Type ContractionOp::getExpectedMaskType() {
1028  auto indexingMaps = this->getIndexingMapsArray();
1029  AffineMap lhsIdxMap = indexingMaps[0];
1030  AffineMap rhsIdxMap = indexingMaps[1];
1031  VectorType lhsType = this->getLhsType();
1032  VectorType rhsType = this->getRhsType();
1033 
1034  unsigned numVecDims = lhsIdxMap.getNumDims();
1035  SmallVector<int64_t> maskShape(numVecDims, ShapedType::kDynamic);
1036  SmallVector<bool> maskShapeScalableDims(numVecDims, false);
1037 
1038  // Using the information in the indexing maps, extract the size of each
1039  // dimension in the vector.contract operation from the two input operands.
1040  for (auto [dimIdx, dimSize] : llvm::enumerate(lhsType.getShape())) {
1041  maskShape[lhsIdxMap.getDimPosition(dimIdx)] = dimSize;
1042  maskShapeScalableDims[lhsIdxMap.getDimPosition(dimIdx)] =
1043  lhsType.getScalableDims()[dimIdx];
1044  }
1045  for (auto [dimIdx, dimSize] : llvm::enumerate(rhsType.getShape())) {
1046  maskShape[rhsIdxMap.getDimPosition(dimIdx)] = dimSize;
1047  maskShapeScalableDims[rhsIdxMap.getDimPosition(dimIdx)] =
1048  rhsType.getScalableDims()[dimIdx];
1049  }
1050 
1051  assert(!ShapedType::isDynamicShape(maskShape) &&
1052  "Mask shape couldn't be computed");
1053 
1054  return VectorType::get(maskShape,
1055  IntegerType::get(lhsType.getContext(), /*width=*/1),
1056  maskShapeScalableDims);
1057 }
1058 
1059 SmallVector<StringRef> ContractionOp::getTraitAttrNames() {
1060  return SmallVector<StringRef>{getIndexingMapsAttrName(),
1061  getIteratorTypesAttrName(), getKindAttrName()};
1062 }
1063 
1064 static int64_t getResultIndex(AffineMap map, AffineExpr targetExpr) {
1065  for (int64_t i = 0, e = map.getNumResults(); i < e; ++i)
1066  if (targetExpr == map.getResult(i))
1067  return i;
1068  return -1;
1069 }
1070 
1071 static std::vector<std::pair<int64_t, int64_t>>
1072 getDimMap(ArrayRef<AffineMap> indexingMaps, ArrayAttr iteratorTypes,
1073  IteratorType targetIteratorType, MLIRContext *context) {
1074  std::vector<std::pair<int64_t, int64_t>> dimMap;
1075  for (const auto &it : llvm::enumerate(iteratorTypes)) {
1076  auto iteratorType = llvm::cast<IteratorTypeAttr>(it.value()).getValue();
1077  if (iteratorType != targetIteratorType)
1078  continue;
1079  // Search lhs/rhs map results for 'targetExpr'.
1080  auto targetExpr = getAffineDimExpr(it.index(), context);
1081  int64_t lhsDim = getResultIndex(indexingMaps[0], targetExpr);
1082  int64_t rhsDim = getResultIndex(indexingMaps[1], targetExpr);
1083  if (lhsDim >= 0 && rhsDim >= 0)
1084  dimMap.emplace_back(lhsDim, rhsDim);
1085  }
1086  return dimMap;
1087 }
1088 
1089 void ContractionOp::getIterationBounds(
1090  SmallVectorImpl<int64_t> &iterationBounds) {
1091  auto lhsShape = getLhsType().getShape();
1092  auto resVectorType = llvm::dyn_cast<VectorType>(getResultType());
1093  SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
1094  SmallVector<int64_t, 2> iterationShape;
1095  for (const auto &it : llvm::enumerate(getIteratorTypes())) {
1096  // Search lhs/rhs map results for 'targetExpr'.
1097  auto targetExpr = getAffineDimExpr(it.index(), getContext());
1098  auto iteratorType = llvm::cast<IteratorTypeAttr>(it.value()).getValue();
1099  if (iteratorType == IteratorType::reduction) {
1100  // Get reduction dim size from lhs shape (same size in rhsShape).
1101  int64_t lhsDimIndex = getResultIndex(indexingMaps[0], targetExpr);
1102  assert(lhsDimIndex >= 0);
1103  iterationBounds.push_back(lhsShape[lhsDimIndex]);
1104  continue;
1105  }
1106  // Get parallel dimension size from result shape.
1107  int64_t resDimIndex = getResultIndex(indexingMaps[2], targetExpr);
1108  assert(resDimIndex >= 0);
1109  assert(resVectorType != nullptr);
1110  iterationBounds.push_back(resVectorType.getShape()[resDimIndex]);
1111  }
1112 }
1113 
1114 void ContractionOp::getIterationIndexMap(
1115  std::vector<DenseMap<int64_t, int64_t>> &iterationIndexMap) {
1116  unsigned numMaps = getIndexingMapsArray().size();
1117  iterationIndexMap.resize(numMaps);
1118  for (const auto &it : llvm::enumerate(getIndexingMapsArray())) {
1119  auto index = it.index();
1120  auto map = it.value();
1121  for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
1122  auto dim = cast<AffineDimExpr>(map.getResult(i));
1123  iterationIndexMap[index][dim.getPosition()] = i;
1124  }
1125  }
1126 }
1127 
1128 std::vector<std::pair<int64_t, int64_t>> ContractionOp::getContractingDimMap() {
1129  SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
1130  return getDimMap(indexingMaps, getIteratorTypes(), IteratorType::reduction,
1131  getContext());
1132 }
1133 
1134 std::vector<std::pair<int64_t, int64_t>> ContractionOp::getBatchDimMap() {
1135  SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
1136  return getDimMap(indexingMaps, getIteratorTypes(), IteratorType::parallel,
1137  getContext());
1138 }
1139 
1140 std::optional<SmallVector<int64_t, 4>> ContractionOp::getShapeForUnroll() {
1142  getIterationBounds(shape);
1143  return shape;
1144 }
1145 
1146 /// Return a fused vector::ContractionOp which represents a patterns such as:
1147 ///
1148 /// ```mlir
1149 /// %c0 = vector.constant 0: ...
1150 /// %c = vector.contract %a, %b, %c0: ...
1151 /// %e = add %c, %d: ...
1152 /// ```
1153 ///
1154 /// by:
1155 ///
1156 /// ```mlir
1157 /// %e = vector.contract %a, %b, %d: ...
1158 /// ```
1159 ///
1160 /// Return null if the canonicalization does not apply.
1161 // TODO: This should be a folding of Add into Contract in core but while they
1162 // live in different dialects, it is not possible without unnatural
1163 // dependencies.
1164 template <typename AddOpType>
1165 struct CanonicalizeContractAdd : public OpRewritePattern<AddOpType> {
1167 
1169  PatternRewriter &rewriter) const override {
1170  auto canonicalize = [&](Value maybeContraction,
1171  Value otherOperand) -> vector::ContractionOp {
1172  vector::ContractionOp contractionOp =
1173  dyn_cast_or_null<vector::ContractionOp>(
1174  maybeContraction.getDefiningOp());
1175  if (!contractionOp)
1176  return vector::ContractionOp();
1177  if (auto maybeZero = dyn_cast_or_null<arith::ConstantOp>(
1178  contractionOp.getAcc().getDefiningOp())) {
1179  if (maybeZero.getValue() ==
1180  rewriter.getZeroAttr(contractionOp.getAcc().getType())) {
1181  IRMapping bvm;
1182  bvm.map(contractionOp.getAcc(), otherOperand);
1183  auto newContraction =
1184  cast<vector::ContractionOp>(rewriter.clone(*contractionOp, bvm));
1185  rewriter.replaceOp(addOp, newContraction.getResult());
1186  return newContraction;
1187  }
1188  }
1189  return vector::ContractionOp();
1190  };
1191 
1192  Value a = addOp->getOperand(0), b = addOp->getOperand(1);
1193  vector::ContractionOp contract = canonicalize(a, b);
1194  contract = contract ? contract : canonicalize(b, a);
1195  return contract ? success() : failure();
1196  }
1197 };
1198 
1199 void ContractionOp::getCanonicalizationPatterns(RewritePatternSet &results,
1200  MLIRContext *context) {
1203 }
1204 
1205 //===----------------------------------------------------------------------===//
1206 // ExtractElementOp
1207 //===----------------------------------------------------------------------===//
1208 
1209 void vector::ExtractElementOp::build(OpBuilder &builder, OperationState &result,
1210  Value source) {
1211  result.addOperands({source});
1212  result.addTypes(llvm::cast<VectorType>(source.getType()).getElementType());
1213 }
1214 
1216  VectorType vectorType = getSourceVectorType();
1217  if (vectorType.getRank() == 0) {
1218  if (getPosition())
1219  return emitOpError("expected position to be empty with 0-D vector");
1220  return success();
1221  }
1222  if (vectorType.getRank() != 1)
1223  return emitOpError("unexpected >1 vector rank");
1224  if (!getPosition())
1225  return emitOpError("expected position for 1-D vector");
1226  return success();
1227 }
1228 
1229 OpFoldResult vector::ExtractElementOp::fold(FoldAdaptor adaptor) {
1230  // Skip the 0-D vector here now.
1231  if (!adaptor.getPosition())
1232  return {};
1233 
1234  // Fold extractelement (splat X) -> X.
1235  if (auto splat = getVector().getDefiningOp<vector::SplatOp>())
1236  return splat.getInput();
1237 
1238  // Fold extractelement(broadcast(X)) -> X.
1239  if (auto broadcast = getVector().getDefiningOp<vector::BroadcastOp>())
1240  if (!llvm::isa<VectorType>(broadcast.getSource().getType()))
1241  return broadcast.getSource();
1242 
1243  auto src = dyn_cast_or_null<DenseElementsAttr>(adaptor.getVector());
1244  auto pos = dyn_cast_or_null<IntegerAttr>(adaptor.getPosition());
1245  if (!pos || !src)
1246  return {};
1247 
1248  auto srcElements = src.getValues<Attribute>();
1249 
1250  uint64_t posIdx = pos.getInt();
1251  if (posIdx >= srcElements.size())
1252  return {};
1253 
1254  return srcElements[posIdx];
1255 }
1256 
1257 //===----------------------------------------------------------------------===//
1258 // ExtractOp
1259 //===----------------------------------------------------------------------===//
1260 
1261 void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
1262  Value source, int64_t position) {
1263  build(builder, result, source, ArrayRef<int64_t>{position});
1264 }
1265 
1266 void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
1267  Value source, OpFoldResult position) {
1268  build(builder, result, source, ArrayRef<OpFoldResult>{position});
1269 }
1270 
1271 void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
1272  Value source, ArrayRef<int64_t> position) {
1273  build(builder, result, source, /*dynamic_position=*/ArrayRef<Value>(),
1274  builder.getDenseI64ArrayAttr(position));
1275 }
1276 
1277 void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
1278  Value source, ArrayRef<OpFoldResult> position) {
1279  SmallVector<int64_t> staticPos;
1280  SmallVector<Value> dynamicPos;
1281  dispatchIndexOpFoldResults(position, dynamicPos, staticPos);
1282  build(builder, result, source, dynamicPos,
1283  builder.getDenseI64ArrayAttr(staticPos));
1284 }
1285 
1287 ExtractOp::inferReturnTypes(MLIRContext *, std::optional<Location>,
1288  ExtractOp::Adaptor adaptor,
1289  SmallVectorImpl<Type> &inferredReturnTypes) {
1290  auto vectorType = llvm::cast<VectorType>(adaptor.getVector().getType());
1291  if (static_cast<int64_t>(adaptor.getStaticPosition().size()) ==
1292  vectorType.getRank()) {
1293  inferredReturnTypes.push_back(vectorType.getElementType());
1294  } else {
1295  auto n = std::min<size_t>(adaptor.getStaticPosition().size(),
1296  vectorType.getRank());
1297  inferredReturnTypes.push_back(VectorType::get(
1298  vectorType.getShape().drop_front(n), vectorType.getElementType(),
1299  vectorType.getScalableDims().drop_front(n)));
1300  }
1301  return success();
1302 }
1303 
1304 bool ExtractOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
1305  // Allow extracting 1-element vectors instead of scalars.
1306  auto isCompatible = [](TypeRange l, TypeRange r) {
1307  auto vectorType = llvm::dyn_cast<VectorType>(l.front());
1308  return vectorType && vectorType.getShape().equals({1}) &&
1309  vectorType.getElementType() == r.front();
1310  };
1311  if (l.size() == 1 && r.size() == 1 &&
1312  (isCompatible(l, r) || isCompatible(r, l)))
1313  return true;
1314  return l == r;
1315 }
1316 
1318  // Note: This check must come before getMixedPosition() to prevent a crash.
1319  auto dynamicMarkersCount =
1320  llvm::count_if(getStaticPosition(), ShapedType::isDynamic);
1321  if (static_cast<size_t>(dynamicMarkersCount) != getDynamicPosition().size())
1322  return emitOpError(
1323  "mismatch between dynamic and static positions (kDynamic marker but no "
1324  "corresponding dynamic position) -- this can only happen due to an "
1325  "incorrect fold/rewrite");
1326  auto position = getMixedPosition();
1327  if (position.size() > static_cast<unsigned>(getSourceVectorType().getRank()))
1328  return emitOpError(
1329  "expected position attribute of rank no greater than vector rank");
1330  for (auto [idx, pos] : llvm::enumerate(position)) {
1331  if (pos.is<Attribute>()) {
1332  int64_t constIdx = cast<IntegerAttr>(pos.get<Attribute>()).getInt();
1333  if (constIdx < 0 || constIdx >= getSourceVectorType().getDimSize(idx)) {
1334  return emitOpError("expected position attribute #")
1335  << (idx + 1)
1336  << " to be a non-negative integer smaller than the "
1337  "corresponding vector dimension";
1338  }
1339  }
1340  }
1341  return success();
1342 }
1343 
1344 template <typename IntType>
1345 static SmallVector<IntType> extractVector(ArrayAttr arrayAttr) {
1346  return llvm::to_vector<4>(llvm::map_range(
1347  arrayAttr.getAsRange<IntegerAttr>(),
1348  [](IntegerAttr attr) { return static_cast<IntType>(attr.getInt()); }));
1349 }
1350 
1351 /// Fold the result of chains of ExtractOp in place by simply concatenating the
1352 /// positions.
1353 static LogicalResult foldExtractOpFromExtractChain(ExtractOp extractOp) {
1354  if (!extractOp.getVector().getDefiningOp<ExtractOp>())
1355  return failure();
1356 
1357  // TODO: Canonicalization for dynamic position not implemented yet.
1358  if (extractOp.hasDynamicPosition())
1359  return failure();
1360 
1361  SmallVector<int64_t> globalPosition;
1362  ExtractOp currentOp = extractOp;
1363  ArrayRef<int64_t> extrPos = currentOp.getStaticPosition();
1364  globalPosition.append(extrPos.rbegin(), extrPos.rend());
1365  while (ExtractOp nextOp = currentOp.getVector().getDefiningOp<ExtractOp>()) {
1366  currentOp = nextOp;
1367  // TODO: Canonicalization for dynamic position not implemented yet.
1368  if (currentOp.hasDynamicPosition())
1369  return failure();
1370  ArrayRef<int64_t> extrPos = currentOp.getStaticPosition();
1371  globalPosition.append(extrPos.rbegin(), extrPos.rend());
1372  }
1373  extractOp.setOperand(0, currentOp.getVector());
1374  // OpBuilder is only used as a helper to build an I64ArrayAttr.
1375  OpBuilder b(extractOp.getContext());
1376  std::reverse(globalPosition.begin(), globalPosition.end());
1377  extractOp.setStaticPosition(globalPosition);
1378  return success();
1379 }
1380 
1381 namespace {
1382 /// Fold an ExtractOp that is fed by a chain of InsertOps and TransposeOps.
1383 /// Walk back a chain of InsertOp/TransposeOp until we hit a match.
1384 /// Compose TransposeOp permutations as we walk back.
1385 /// This helper class keeps an updated extraction position `extractPosition`
1386 /// with extra trailing sentinels.
1387 /// The sentinels encode the internal transposition status of the result vector.
1388 /// As we iterate, extractPosition is permuted and updated.
1389 class ExtractFromInsertTransposeChainState {
1390 public:
1391  ExtractFromInsertTransposeChainState(ExtractOp e);
1392 
1393  /// Iterate over producing insert and transpose ops until we find a fold.
1394  Value fold();
1395 
1396 private:
1397  /// Return true if the vector at position `a` is contained within the vector
1398  /// at position `b`. Under insert/extract semantics, this is the same as `a`
1399  /// is a prefix of `b`.
1400  template <typename ContainerA, typename ContainerB>
1401  bool isContainedWithin(const ContainerA &a, const ContainerB &b) {
1402  return a.size() <= b.size() &&
1403  std::equal(a.begin(), a.begin() + a.size(), b.begin());
1404  }
1405 
1406  /// Return true if the vector at position `a` intersects the vector at
1407  /// position `b`. Under insert/extract semantics, this is the same as equality
1408  /// of all entries of `a` that are >=0 with the corresponding entries of b.
1409  /// Comparison is on the common prefix (i.e. zip).
1410  template <typename ContainerA, typename ContainerB>
1411  bool intersectsWhereNonNegative(const ContainerA &a, const ContainerB &b) {
1412  for (auto [elemA, elemB] : llvm::zip(a, b)) {
1413  if (elemA < 0 || elemB < 0)
1414  continue;
1415  if (elemA != elemB)
1416  return false;
1417  }
1418  return true;
1419  }
1420 
1421  /// Folding is only possible in the absence of an internal permutation in the
1422  /// result vector.
1423  bool canFold() {
1424  return (sentinels == ArrayRef(extractPosition).drop_front(extractedRank));
1425  }
1426 
1427  // Helper to get the next defining op of interest.
1428  void updateStateForNextIteration(Value v) {
1429  nextInsertOp = v.getDefiningOp<vector::InsertOp>();
1430  nextTransposeOp = v.getDefiningOp<vector::TransposeOp>();
1431  };
1432 
1433  // Case 1. If we hit a transpose, just compose the map and iterate.
1434  // Invariant: insert + transpose do not change rank, we can always compose.
1435  LogicalResult handleTransposeOp();
1436 
1437  // Case 2: the insert position matches extractPosition exactly, early return.
1438  LogicalResult handleInsertOpWithMatchingPos(Value &res);
1439 
1440  /// Case 3: if the insert position is a prefix of extractPosition, extract a
1441  /// portion of the source of the insert.
1442  /// Example:
1443  /// ```
1444  /// %ins = vector.insert %source, %vest[1]: vector<3x4> into vector<2x3x4x5>
1445  /// // extractPosition == [1, 2, 3]
1446  /// %ext = vector.extract %ins[1, 0]: vector<5> from vector<3x4x5>
1447  /// // can fold to vector.extract %source[0, 3]
1448  /// %ext = vector.extract %source[3]: vector<6> from vector<5x6>
1449  /// ```
1450  /// To traverse through %source, we need to set the leading dims to 0 and
1451  /// drop the extra leading dims.
1452  /// This method updates the internal state.
1453  LogicalResult handleInsertOpWithPrefixPos(Value &res);
1454 
1455  /// Try to fold in place to extract(source, extractPosition) and return the
1456  /// folded result. Return null if folding is not possible (e.g. due to an
1457  /// internal tranposition in the result).
1458  Value tryToFoldExtractOpInPlace(Value source);
1459 
1460  ExtractOp extractOp;
1461  int64_t vectorRank;
1462  int64_t extractedRank;
1463 
1464  InsertOp nextInsertOp;
1465  TransposeOp nextTransposeOp;
1466 
1467  /// Sentinel values that encode the internal permutation status of the result.
1468  /// They are set to (-1, ... , -k) at the beginning and appended to
1469  /// `extractPosition`.
1470  /// In the end, the tail of `extractPosition` must be exactly `sentinels` to
1471  /// ensure that there is no internal transposition.
1472  /// Internal transposition cannot be accounted for with a folding pattern.
1473  // TODO: We could relax the internal transposition with an extra transposition
1474  // operation in a future canonicalizer.
1475  SmallVector<int64_t> sentinels;
1477 };
1478 } // namespace
1479 
1480 ExtractFromInsertTransposeChainState::ExtractFromInsertTransposeChainState(
1481  ExtractOp e)
1482  : extractOp(e), vectorRank(extractOp.getSourceVectorType().getRank()),
1483  extractedRank(extractOp.getNumIndices()) {
1484  assert(vectorRank >= extractedRank && "Extracted position overflow");
1485  sentinels.reserve(vectorRank - extractedRank);
1486  for (int64_t i = 0, e = vectorRank - extractedRank; i < e; ++i)
1487  sentinels.push_back(-(i + 1));
1488  extractPosition.assign(extractOp.getStaticPosition().begin(),
1489  extractOp.getStaticPosition().end());
1490  llvm::append_range(extractPosition, sentinels);
1491 }
1492 
1493 // Case 1. If we hit a transpose, just compose the map and iterate.
1494 // Invariant: insert + transpose do not change rank, we can always compose.
1495 LogicalResult ExtractFromInsertTransposeChainState::handleTransposeOp() {
1496  // TODO: Canonicalization for dynamic position not implemented yet.
1497  if (extractOp.hasDynamicPosition())
1498  return failure();
1499 
1500  if (!nextTransposeOp)
1501  return failure();
1502  AffineMap m = inversePermutation(AffineMap::getPermutationMap(
1503  nextTransposeOp.getPermutation(), extractOp.getContext()));
1505  return success();
1506 }
1507 
1508 // Case 2: the insert position matches extractPosition exactly, early return.
1510 ExtractFromInsertTransposeChainState::handleInsertOpWithMatchingPos(
1511  Value &res) {
1512  // TODO: Canonicalization for dynamic position not implemented yet.
1513  if (extractOp.hasDynamicPosition() || nextInsertOp.hasDynamicPosition())
1514  return failure();
1515 
1516  ArrayRef<int64_t> insertedPos = nextInsertOp.getStaticPosition();
1517  if (insertedPos != llvm::ArrayRef(extractPosition).take_front(extractedRank))
1518  return failure();
1519  // Case 2.a. early-exit fold.
1520  res = nextInsertOp.getSource();
1521  // Case 2.b. if internal transposition is present, canFold will be false.
1522  return success(canFold());
1523 }
1524 
1525 /// Case 3: if inserted position is a prefix of extractPosition,
1526 /// extract a portion of the source of the insertion.
1527 /// This method updates the internal state.
1529 ExtractFromInsertTransposeChainState::handleInsertOpWithPrefixPos(Value &res) {
1530  // TODO: Canonicalization for dynamic position not implemented yet.
1531  if (extractOp.hasDynamicPosition() || nextInsertOp.hasDynamicPosition())
1532  return failure();
1533 
1534  ArrayRef<int64_t> insertedPos = nextInsertOp.getStaticPosition();
1535  if (!isContainedWithin(insertedPos, extractPosition))
1536  return failure();
1537  // Set leading dims to zero.
1538  std::fill_n(extractPosition.begin(), insertedPos.size(), 0);
1539  // Drop extra leading dims.
1540  extractPosition.erase(extractPosition.begin(),
1541  extractPosition.begin() + insertedPos.size());
1542  extractedRank = extractPosition.size() - sentinels.size();
1543  // Case 3.a. early-exit fold (break and delegate to post-while path).
1544  res = nextInsertOp.getSource();
1545  // Case 3.b. if internal transposition is present, canFold will be false.
1546  return success();
1547 }
1548 
1549 /// Try to fold in place to extract(source, extractPosition) and return the
1550 /// folded result. Return null if folding is not possible (e.g. due to an
1551 /// internal tranposition in the result).
1552 Value ExtractFromInsertTransposeChainState::tryToFoldExtractOpInPlace(
1553  Value source) {
1554  // TODO: Canonicalization for dynamic position not implemented yet.
1555  if (extractOp.hasDynamicPosition())
1556  return Value();
1557 
1558  // If we can't fold (either internal transposition, or nothing to fold), bail.
1559  bool nothingToFold = (source == extractOp.getVector());
1560  if (nothingToFold || !canFold())
1561  return Value();
1562 
1563  // Otherwise, fold by updating the op inplace and return its result.
1564  OpBuilder b(extractOp.getContext());
1565  extractOp.setStaticPosition(
1566  ArrayRef(extractPosition).take_front(extractedRank));
1567  extractOp.getVectorMutable().assign(source);
1568  return extractOp.getResult();
1569 }
1570 
1571 /// Iterate over producing insert and transpose ops until we find a fold.
1572 Value ExtractFromInsertTransposeChainState::fold() {
1573  // TODO: Canonicalization for dynamic position not implemented yet.
1574  if (extractOp.hasDynamicPosition())
1575  return Value();
1576 
1577  Value valueToExtractFrom = extractOp.getVector();
1578  updateStateForNextIteration(valueToExtractFrom);
1579  while (nextInsertOp || nextTransposeOp) {
1580  // Case 1. If we hit a transpose, just compose the map and iterate.
1581  // Invariant: insert + transpose do not change rank, we can always compose.
1582  if (succeeded(handleTransposeOp())) {
1583  valueToExtractFrom = nextTransposeOp.getVector();
1584  updateStateForNextIteration(valueToExtractFrom);
1585  continue;
1586  }
1587 
1588  Value result;
1589  // Case 2: the position match exactly.
1590  if (succeeded(handleInsertOpWithMatchingPos(result)))
1591  return result;
1592 
1593  // Case 3: if the inserted position is a prefix of extractPosition, we can
1594  // just extract a portion of the source of the insert.
1595  if (succeeded(handleInsertOpWithPrefixPos(result)))
1596  return tryToFoldExtractOpInPlace(result);
1597 
1598  // Case 4: extractPositionRef intersects insertedPosRef on non-sentinel
1599  // values. This is a more difficult case and we bail.
1600  ArrayRef<int64_t> insertedPos = nextInsertOp.getStaticPosition();
1601  if (isContainedWithin(extractPosition, insertedPos) ||
1602  intersectsWhereNonNegative(extractPosition, insertedPos))
1603  return Value();
1604 
1605  // Case 5: No intersection, we forward the extract to insertOp.dest().
1606  valueToExtractFrom = nextInsertOp.getDest();
1607  updateStateForNextIteration(valueToExtractFrom);
1608  }
1609  // If after all this we can fold, go for it.
1610  return tryToFoldExtractOpInPlace(valueToExtractFrom);
1611 }
1612 
1613 /// Returns true if the operation has a 0-D vector type operand or result.
1614 static bool hasZeroDimVectors(Operation *op) {
1615  auto hasZeroDimVectorType = [](Type type) -> bool {
1616  auto vecType = dyn_cast<VectorType>(type);
1617  return vecType && vecType.getRank() == 0;
1618  };
1619 
1620  return llvm::any_of(op->getOperandTypes(), hasZeroDimVectorType) ||
1621  llvm::any_of(op->getResultTypes(), hasZeroDimVectorType);
1622 }
1623 
1624 /// Fold extractOp with scalar result coming from BroadcastOp or SplatOp.
1625 static Value foldExtractFromBroadcast(ExtractOp extractOp) {
1626  // TODO: Canonicalization for dynamic position not implemented yet.
1627  if (extractOp.hasDynamicPosition())
1628  return Value();
1629 
1630  Operation *defOp = extractOp.getVector().getDefiningOp();
1631  if (!defOp || !isa<vector::BroadcastOp, SplatOp>(defOp))
1632  return Value();
1633 
1634  // 0-D vectors not supported.
1635  assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
1636  if (hasZeroDimVectors(defOp))
1637  return Value();
1638 
1639  Value source = defOp->getOperand(0);
1640  if (extractOp.getType() == source.getType())
1641  return source;
1642  auto getRank = [](Type type) {
1643  return llvm::isa<VectorType>(type) ? llvm::cast<VectorType>(type).getRank()
1644  : 0;
1645  };
1646 
1647  // If splat or broadcast from a scalar, just return the source scalar.
1648  unsigned broadcastSrcRank = getRank(source.getType());
1649  if (broadcastSrcRank == 0 && source.getType() == extractOp.getType())
1650  return source;
1651 
1652  unsigned extractResultRank = getRank(extractOp.getType());
1653  if (extractResultRank >= broadcastSrcRank)
1654  return Value();
1655  // Check that the dimension of the result haven't been broadcasted.
1656  auto extractVecType = llvm::dyn_cast<VectorType>(extractOp.getType());
1657  auto broadcastVecType = llvm::dyn_cast<VectorType>(source.getType());
1658  if (extractVecType && broadcastVecType &&
1659  extractVecType.getShape() !=
1660  broadcastVecType.getShape().take_back(extractResultRank))
1661  return Value();
1662 
1663  auto broadcastOp = cast<vector::BroadcastOp>(defOp);
1664  int64_t broadcastDstRank = broadcastOp.getResultVectorType().getRank();
1665 
1666  // Detect all the positions that come from "dim-1" broadcasting.
1667  // These dimensions correspond to "dim-1" broadcasted dims; set the mathching
1668  // extract position to `0` when extracting from the source operand.
1669  llvm::SetVector<int64_t> broadcastedUnitDims =
1670  broadcastOp.computeBroadcastedUnitDims();
1671  SmallVector<int64_t> extractPos(extractOp.getStaticPosition());
1672  int64_t broadcastRankDiff = broadcastDstRank - broadcastSrcRank;
1673  for (int64_t i = broadcastRankDiff, e = extractPos.size(); i < e; ++i)
1674  if (broadcastedUnitDims.contains(i))
1675  extractPos[i] = 0;
1676  // `rankDiff` leading dimensions correspond to new broadcasted dims, drop the
1677  // matching extract position when extracting from the source operand.
1678  int64_t rankDiff = broadcastSrcRank - extractResultRank;
1679  extractPos.erase(extractPos.begin(),
1680  std::next(extractPos.begin(), extractPos.size() - rankDiff));
1681  // OpBuilder is only used as a helper to build an I64ArrayAttr.
1682  OpBuilder b(extractOp.getContext());
1683  extractOp.setOperand(0, source);
1684  extractOp.setStaticPosition(extractPos);
1685  return extractOp.getResult();
1686 }
1687 
1688 // Fold extractOp with source coming from ShapeCast op.
1689 static Value foldExtractFromShapeCast(ExtractOp extractOp) {
1690  // TODO: Canonicalization for dynamic position not implemented yet.
1691  if (extractOp.hasDynamicPosition())
1692  return Value();
1693 
1694  auto shapeCastOp = extractOp.getVector().getDefiningOp<vector::ShapeCastOp>();
1695  if (!shapeCastOp)
1696  return Value();
1697 
1698  // 0-D vectors not supported.
1699  assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
1700  if (hasZeroDimVectors(shapeCastOp))
1701  return Value();
1702 
1703  // Get the nth dimension size starting from lowest dimension.
1704  auto getDimReverse = [](VectorType type, int64_t n) {
1705  return type.getShape().take_back(n + 1).front();
1706  };
1707  int64_t destinationRank =
1708  llvm::isa<VectorType>(extractOp.getType())
1709  ? llvm::cast<VectorType>(extractOp.getType()).getRank()
1710  : 0;
1711  if (destinationRank > shapeCastOp.getSourceVectorType().getRank())
1712  return Value();
1713  if (destinationRank > 0) {
1714  auto destinationType =
1715  llvm::cast<VectorType>(extractOp.getResult().getType());
1716  for (int64_t i = 0; i < destinationRank; i++) {
1717  // The lowest dimension of the destination must match the lowest
1718  // dimension of the shapecast op source.
1719  // TODO: This case could be support in a canonicalization pattern.
1720  if (getDimReverse(shapeCastOp.getSourceVectorType(), i) !=
1721  getDimReverse(destinationType, i))
1722  return Value();
1723  }
1724  }
1725  // Extract the strides associated with the extract op vector source. Then use
1726  // this to calculate a linearized position for the extract.
1727  SmallVector<int64_t> extractedPos(extractOp.getStaticPosition());
1728  std::reverse(extractedPos.begin(), extractedPos.end());
1729  SmallVector<int64_t, 4> strides;
1730  int64_t stride = 1;
1731  for (int64_t i = 0, e = extractedPos.size(); i < e; i++) {
1732  strides.push_back(stride);
1733  stride *=
1734  getDimReverse(extractOp.getSourceVectorType(), i + destinationRank);
1735  }
1736 
1737  int64_t position = linearize(extractedPos, strides);
1738  // Then extract the strides associated to the shapeCast op vector source and
1739  // delinearize the position using those strides.
1740  SmallVector<int64_t, 4> newStrides;
1741  int64_t numDimension =
1742  shapeCastOp.getSourceVectorType().getRank() - destinationRank;
1743  stride = 1;
1744  for (int64_t i = 0; i < numDimension; i++) {
1745  newStrides.push_back(stride);
1746  stride *=
1747  getDimReverse(shapeCastOp.getSourceVectorType(), i + destinationRank);
1748  }
1749  std::reverse(newStrides.begin(), newStrides.end());
1750  SmallVector<int64_t, 4> newPosition = delinearize(position, newStrides);
1751  // OpBuilder is only used as a helper to build an I64ArrayAttr.
1752  OpBuilder b(extractOp.getContext());
1753  extractOp.setStaticPosition(newPosition);
1754  extractOp.setOperand(0, shapeCastOp.getSource());
1755  return extractOp.getResult();
1756 }
1757 
1758 /// Fold an ExtractOp from ExtractStridedSliceOp.
1759 static Value foldExtractFromExtractStrided(ExtractOp extractOp) {
1760  // TODO: Canonicalization for dynamic position not implemented yet.
1761  if (extractOp.hasDynamicPosition())
1762  return Value();
1763 
1764  auto extractStridedSliceOp =
1765  extractOp.getVector().getDefiningOp<vector::ExtractStridedSliceOp>();
1766  if (!extractStridedSliceOp)
1767  return Value();
1768 
1769  // 0-D vectors not supported.
1770  assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
1771  if (hasZeroDimVectors(extractStridedSliceOp))
1772  return Value();
1773 
1774  // Return if 'extractStridedSliceOp' has non-unit strides.
1775  if (extractStridedSliceOp.hasNonUnitStrides())
1776  return Value();
1777 
1778  // Trim offsets for dimensions fully extracted.
1779  auto sliceOffsets =
1780  extractVector<int64_t>(extractStridedSliceOp.getOffsets());
1781  while (!sliceOffsets.empty()) {
1782  size_t lastOffset = sliceOffsets.size() - 1;
1783  if (sliceOffsets.back() != 0 ||
1784  extractStridedSliceOp.getType().getDimSize(lastOffset) !=
1785  extractStridedSliceOp.getSourceVectorType().getDimSize(lastOffset))
1786  break;
1787  sliceOffsets.pop_back();
1788  }
1789  unsigned destinationRank = 0;
1790  if (auto vecType = llvm::dyn_cast<VectorType>(extractOp.getType()))
1791  destinationRank = vecType.getRank();
1792  // The dimensions of the result need to be untouched by the
1793  // extractStridedSlice op.
1794  if (destinationRank > extractStridedSliceOp.getSourceVectorType().getRank() -
1795  sliceOffsets.size())
1796  return Value();
1797 
1798  SmallVector<int64_t> extractedPos(extractOp.getStaticPosition());
1799  assert(extractedPos.size() >= sliceOffsets.size());
1800  for (size_t i = 0, e = sliceOffsets.size(); i < e; i++)
1801  extractedPos[i] = extractedPos[i] + sliceOffsets[i];
1802  extractOp.getVectorMutable().assign(extractStridedSliceOp.getVector());
1803 
1804  // OpBuilder is only used as a helper to build an I64ArrayAttr.
1805  OpBuilder b(extractOp.getContext());
1806  extractOp.setStaticPosition(extractedPos);
1807  return extractOp.getResult();
1808 }
1809 
1810 /// Fold extract_op fed from a chain of insertStridedSlice ops.
1811 static Value foldExtractStridedOpFromInsertChain(ExtractOp extractOp) {
1812  // TODO: Canonicalization for dynamic position not implemented yet.
1813  if (extractOp.hasDynamicPosition())
1814  return Value();
1815 
1816  int64_t destinationRank =
1817  llvm::isa<VectorType>(extractOp.getType())
1818  ? llvm::cast<VectorType>(extractOp.getType()).getRank()
1819  : 0;
1820  auto insertOp = extractOp.getVector().getDefiningOp<InsertStridedSliceOp>();
1821  if (!insertOp)
1822  return Value();
1823 
1824  // 0-D vectors not supported.
1825  assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
1826  if (hasZeroDimVectors(insertOp))
1827  return Value();
1828 
1829  while (insertOp) {
1830  int64_t insertRankDiff = insertOp.getDestVectorType().getRank() -
1831  insertOp.getSourceVectorType().getRank();
1832  if (destinationRank > insertOp.getSourceVectorType().getRank())
1833  return Value();
1834  auto insertOffsets = extractVector<int64_t>(insertOp.getOffsets());
1835  ArrayRef<int64_t> extractOffsets = extractOp.getStaticPosition();
1836 
1837  if (llvm::any_of(insertOp.getStrides(), [](Attribute attr) {
1838  return llvm::cast<IntegerAttr>(attr).getInt() != 1;
1839  }))
1840  return Value();
1841  bool disjoint = false;
1842  SmallVector<int64_t, 4> offsetDiffs;
1843  for (unsigned dim = 0, e = extractOffsets.size(); dim < e; ++dim) {
1844  int64_t start = insertOffsets[dim];
1845  int64_t size =
1846  (dim < insertRankDiff)
1847  ? 1
1848  : insertOp.getSourceVectorType().getDimSize(dim - insertRankDiff);
1849  int64_t end = start + size;
1850  int64_t offset = extractOffsets[dim];
1851  // Check if the start of the extract offset is in the interval inserted.
1852  if (start <= offset && offset < end) {
1853  if (dim >= insertRankDiff)
1854  offsetDiffs.push_back(offset - start);
1855  continue;
1856  }
1857  disjoint = true;
1858  break;
1859  }
1860  // The extract element chunk overlap with the vector inserted.
1861  if (!disjoint) {
1862  // If any of the inner dimensions are only partially inserted we have a
1863  // partial overlap.
1864  int64_t srcRankDiff =
1865  insertOp.getSourceVectorType().getRank() - destinationRank;
1866  for (int64_t i = 0; i < destinationRank; i++) {
1867  if (insertOp.getSourceVectorType().getDimSize(i + srcRankDiff) !=
1868  insertOp.getDestVectorType().getDimSize(i + srcRankDiff +
1869  insertRankDiff))
1870  return Value();
1871  }
1872  extractOp.getVectorMutable().assign(insertOp.getSource());
1873  // OpBuilder is only used as a helper to build an I64ArrayAttr.
1874  OpBuilder b(extractOp.getContext());
1875  extractOp.setStaticPosition(offsetDiffs);
1876  return extractOp.getResult();
1877  }
1878  // If the chunk extracted is disjoint from the chunk inserted, keep
1879  // looking in the insert chain.
1880  insertOp = insertOp.getDest().getDefiningOp<InsertStridedSliceOp>();
1881  }
1882  return Value();
1883 }
1884 
1885 OpFoldResult ExtractOp::fold(FoldAdaptor) {
1886  if (getNumIndices() == 0)
1887  return getVector();
1889  return getResult();
1890  if (auto res = ExtractFromInsertTransposeChainState(*this).fold())
1891  return res;
1892  if (auto res = foldExtractFromBroadcast(*this))
1893  return res;
1894  if (auto res = foldExtractFromShapeCast(*this))
1895  return res;
1896  if (auto val = foldExtractFromExtractStrided(*this))
1897  return val;
1898  if (auto val = foldExtractStridedOpFromInsertChain(*this))
1899  return val;
1900  return OpFoldResult();
1901 }
1902 
1903 namespace {
1904 
1905 // Pattern to rewrite a ExtractOp(Broadcast) -> Broadcast.
1906 class ExtractOpFromBroadcast final : public OpRewritePattern<ExtractOp> {
1907 public:
1909 
1910  LogicalResult matchAndRewrite(ExtractOp extractOp,
1911  PatternRewriter &rewriter) const override {
1912  Operation *defOp = extractOp.getVector().getDefiningOp();
1913  if (!defOp || !isa<vector::BroadcastOp, SplatOp>(defOp))
1914  return failure();
1915 
1916  Value source = defOp->getOperand(0);
1917  if (extractOp.getType() == source.getType())
1918  return failure();
1919  auto getRank = [](Type type) {
1920  return llvm::isa<VectorType>(type)
1921  ? llvm::cast<VectorType>(type).getRank()
1922  : 0;
1923  };
1924  unsigned broadcastSrcRank = getRank(source.getType());
1925  unsigned extractResultRank = getRank(extractOp.getType());
1926  // We only consider the case where the rank of the source is less than or
1927  // equal to the rank of the extract dst. The other cases are handled in the
1928  // folding patterns.
1929  if (extractResultRank < broadcastSrcRank)
1930  return failure();
1931 
1932  // Special case if broadcast src is a 0D vector.
1933  if (extractResultRank == 0) {
1934  assert(broadcastSrcRank == 0 && llvm::isa<VectorType>(source.getType()));
1935  rewriter.replaceOpWithNewOp<vector::ExtractElementOp>(extractOp, source);
1936  return success();
1937  }
1938  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
1939  extractOp, extractOp.getType(), source);
1940  return success();
1941  }
1942 };
1943 
1944 // Pattern to rewrite a ExtractOp(splat ConstantOp) -> ConstantOp.
1945 class ExtractOpSplatConstantFolder final : public OpRewritePattern<ExtractOp> {
1946 public:
1948 
1949  LogicalResult matchAndRewrite(ExtractOp extractOp,
1950  PatternRewriter &rewriter) const override {
1951  // Return if 'ExtractOp' operand is not defined by a splat vector
1952  // ConstantOp.
1953  Value sourceVector = extractOp.getVector();
1954  Attribute vectorCst;
1955  if (!matchPattern(sourceVector, m_Constant(&vectorCst)))
1956  return failure();
1957  auto splat = llvm::dyn_cast<SplatElementsAttr>(vectorCst);
1958  if (!splat)
1959  return failure();
1960  TypedAttr newAttr = splat.getSplatValue<TypedAttr>();
1961  if (auto vecDstType = llvm::dyn_cast<VectorType>(extractOp.getType()))
1962  newAttr = DenseElementsAttr::get(vecDstType, newAttr);
1963  rewriter.replaceOpWithNewOp<arith::ConstantOp>(extractOp, newAttr);
1964  return success();
1965  }
1966 };
1967 
1968 // Pattern to rewrite a ExtractOp(non-splat ConstantOp)[...] -> ConstantOp.
1969 class ExtractOpNonSplatConstantFolder final
1970  : public OpRewritePattern<ExtractOp> {
1971 public:
1973 
1974  LogicalResult matchAndRewrite(ExtractOp extractOp,
1975  PatternRewriter &rewriter) const override {
1976  // TODO: Canonicalization for dynamic position not implemented yet.
1977  if (extractOp.hasDynamicPosition())
1978  return failure();
1979 
1980  // Return if 'ExtractOp' operand is not defined by a compatible vector
1981  // ConstantOp.
1982  Value sourceVector = extractOp.getVector();
1983  Attribute vectorCst;
1984  if (!matchPattern(sourceVector, m_Constant(&vectorCst)))
1985  return failure();
1986 
1987  auto vecTy = llvm::cast<VectorType>(sourceVector.getType());
1988  if (vecTy.isScalable())
1989  return failure();
1990 
1991  // The splat case is handled by `ExtractOpSplatConstantFolder`.
1992  auto dense = llvm::dyn_cast<DenseElementsAttr>(vectorCst);
1993  if (!dense || dense.isSplat())
1994  return failure();
1995 
1996  // Calculate the linearized position of the continuous chunk of elements to
1997  // extract.
1998  llvm::SmallVector<int64_t> completePositions(vecTy.getRank(), 0);
1999  copy(extractOp.getStaticPosition(), completePositions.begin());
2000  int64_t elemBeginPosition =
2001  linearize(completePositions, computeStrides(vecTy.getShape()));
2002  auto denseValuesBegin = dense.value_begin<TypedAttr>() + elemBeginPosition;
2003 
2004  TypedAttr newAttr;
2005  if (auto resVecTy = llvm::dyn_cast<VectorType>(extractOp.getType())) {
2006  SmallVector<Attribute> elementValues(
2007  denseValuesBegin, denseValuesBegin + resVecTy.getNumElements());
2008  newAttr = DenseElementsAttr::get(resVecTy, elementValues);
2009  } else {
2010  newAttr = *denseValuesBegin;
2011  }
2012 
2013  rewriter.replaceOpWithNewOp<arith::ConstantOp>(extractOp, newAttr);
2014  return success();
2015  }
2016 };
2017 
2018 // Pattern to rewrite a ExtractOp(CreateMask) -> CreateMask.
2019 class ExtractOpFromCreateMask final : public OpRewritePattern<ExtractOp> {
2020 public:
2022 
2023  LogicalResult matchAndRewrite(ExtractOp extractOp,
2024  PatternRewriter &rewriter) const override {
2025  auto createMaskOp =
2026  extractOp.getVector().getDefiningOp<vector::CreateMaskOp>();
2027  if (!createMaskOp)
2028  return failure();
2029 
2030  VectorType extractedMaskType =
2031  llvm::dyn_cast<VectorType>(extractOp.getResult().getType());
2032 
2033  if (!extractedMaskType)
2034  return failure();
2035 
2036  auto maskOperands = createMaskOp.getOperands();
2037  ArrayRef<int64_t> extractOpPos = extractOp.getStaticPosition();
2038  VectorType maskType = createMaskOp.getVectorType();
2039 
2040  bool containsUnknownDims = false;
2041  bool allFalse = getMaskFormat(createMaskOp) == MaskFormat::AllFalse;
2042 
2043  for (size_t dimIdx = 0; !allFalse && dimIdx < extractOpPos.size();
2044  dimIdx++) {
2045  int64_t pos = extractOpPos[dimIdx];
2046  Value operand = maskOperands[dimIdx];
2047  auto constantOp = operand.getDefiningOp<arith::ConstantOp>();
2048  if (!constantOp) {
2049  // Bounds of this dim unknown.
2050  containsUnknownDims = true;
2051  continue;
2052  }
2053 
2054  int64_t createMaskBound =
2055  llvm::cast<IntegerAttr>(constantOp.getValue()).getInt();
2056 
2057  if (pos != ShapedType::kDynamic) {
2058  // If any position is outside the range from the `create_mask`, then the
2059  // extracted mask will be all-false.
2060  allFalse |= pos >= createMaskBound;
2061  } else if (createMaskBound < maskType.getDimSize(dimIdx)) {
2062  // This dim is not all-true and since this is a dynamic index we don't
2063  // know if the extraction is within the true or false region.
2064  // Note: Zero dims have already handled via getMaskFormat().
2065  containsUnknownDims = true;
2066  }
2067  }
2068 
2069  if (allFalse) {
2070  rewriter.replaceOpWithNewOp<arith::ConstantOp>(
2071  extractOp, DenseElementsAttr::get(extractedMaskType, false));
2072  } else if (!containsUnknownDims) {
2073  rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(
2074  extractOp, extractedMaskType,
2075  maskOperands.drop_front(extractOpPos.size()));
2076  } else {
2077  return failure();
2078  }
2079  return success();
2080  }
2081 };
2082 
2083 // Folds extract(shape_cast(..)) into shape_cast when the total element count
2084 // does not change.
2085 LogicalResult foldExtractFromShapeCastToShapeCast(ExtractOp extractOp,
2086  PatternRewriter &rewriter) {
2087  auto castOp = extractOp.getVector().getDefiningOp<ShapeCastOp>();
2088  if (!castOp)
2089  return failure();
2090 
2091  VectorType sourceType = castOp.getSourceVectorType();
2092  auto targetType = dyn_cast<VectorType>(extractOp.getResult().getType());
2093  if (!targetType)
2094  return failure();
2095 
2096  if (sourceType.getNumElements() != targetType.getNumElements())
2097  return failure();
2098 
2099  rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(extractOp, targetType,
2100  castOp.getSource());
2101  return success();
2102 }
2103 
2104 } // namespace
2105 
2106 void ExtractOp::getCanonicalizationPatterns(RewritePatternSet &results,
2107  MLIRContext *context) {
2108  results.add<ExtractOpSplatConstantFolder, ExtractOpNonSplatConstantFolder,
2109  ExtractOpFromBroadcast, ExtractOpFromCreateMask>(context);
2110  results.add(foldExtractFromShapeCastToShapeCast);
2111 }
2112 
2113 static void populateFromInt64AttrArray(ArrayAttr arrayAttr,
2114  SmallVectorImpl<int64_t> &results) {
2115  for (auto attr : arrayAttr)
2116  results.push_back(llvm::cast<IntegerAttr>(attr).getInt());
2117 }
2118 
2119 //===----------------------------------------------------------------------===//
2120 // FmaOp
2121 //===----------------------------------------------------------------------===//
2122 
2123 std::optional<SmallVector<int64_t, 4>> FMAOp::getShapeForUnroll() {
2124  return llvm::to_vector<4>(getVectorType().getShape());
2125 }
2126 
2127 //===----------------------------------------------------------------------===//
2128 // BroadcastOp
2129 //===----------------------------------------------------------------------===//
2130 
2131 /// Return the dimensions of the result vector that were formerly ones in the
2132 /// source tensor and thus correspond to "dim-1" broadcasting.
2135  ArrayRef<int64_t> dstShape) {
2136  int64_t rankDiff = dstShape.size() - srcShape.size();
2137  int64_t dstDim = rankDiff;
2139  for (auto [s1, s2] :
2140  llvm::zip_equal(srcShape, dstShape.drop_front(rankDiff))) {
2141  if (s1 != s2) {
2142  assert(s1 == 1 && "expected dim-1 broadcasting");
2143  res.insert(dstDim);
2144  }
2145  ++dstDim;
2146  }
2147  return res;
2148 }
2149 
2151  // Scalar broadcast is without any unit dim broadcast.
2152  auto srcVectorType = llvm::dyn_cast<VectorType>(getSourceType());
2153  if (!srcVectorType)
2154  return {};
2155  return ::computeBroadcastedUnitDims(srcVectorType.getShape(),
2156  getResultVectorType().getShape());
2157 }
2158 
2159 /// Broadcast `value` to a vector of `dstShape`, knowing that exactly the
2160 /// `broadcastedDims` dimensions in the dstShape are broadcasted.
2161 /// This requires (and asserts) that the broadcast is free of dim-1
2162 /// broadcasting.
2163 /// Since vector.broadcast only allows expanding leading dimensions, an extra
2164 /// vector.transpose may be inserted to make the broadcast possible.
2165 /// `value`, `dstShape` and `broadcastedDims` must be properly specified or
2166 /// the helper will assert. This means:
2167 /// 1. `dstShape` must not be empty.
2168 /// 2. `broadcastedDims` must be confined to [0 .. rank(value.getVectorType)]
2169 /// 2. `dstShape` trimmed of the dimensions specified in `broadcastedDims`
2170 // must match the `value` shape.
2171 Value BroadcastOp::createOrFoldBroadcastOp(
2172  OpBuilder &b, Value value, ArrayRef<int64_t> dstShape,
2173  const llvm::SetVector<int64_t> &broadcastedDims) {
2174  assert(!dstShape.empty() && "unexpected empty dst shape");
2175 
2176  // Well-formedness check.
2177  SmallVector<int64_t> checkShape;
2178  for (int i = 0, e = dstShape.size(); i < e; ++i) {
2179  if (broadcastedDims.contains(i))
2180  continue;
2181  checkShape.push_back(dstShape[i]);
2182  }
2183  assert(broadcastedDims.size() == dstShape.size() - checkShape.size() &&
2184  "ill-formed broadcastedDims contains values not confined to "
2185  "destVectorShape");
2186 
2187  Location loc = value.getLoc();
2188  Type elementType = getElementTypeOrSelf(value.getType());
2189  VectorType srcVectorType = llvm::dyn_cast<VectorType>(value.getType());
2190  VectorType dstVectorType = VectorType::get(dstShape, elementType);
2191 
2192  // Step 2. If scalar -> dstShape broadcast, just do it.
2193  if (!srcVectorType) {
2194  assert(checkShape.empty() &&
2195  "ill-formed createOrFoldBroadcastOp arguments");
2196  return b.createOrFold<vector::BroadcastOp>(loc, dstVectorType, value);
2197  }
2198 
2199  assert(srcVectorType.getShape().equals(checkShape) &&
2200  "ill-formed createOrFoldBroadcastOp arguments");
2201 
2202  // Step 3. Since vector.broadcast only allows creating leading dims,
2203  // vector -> dstShape broadcast may require a transpose.
2204  // Traverse the dims in order and construct:
2205  // 1. The leading entries of the broadcastShape that is guaranteed to be
2206  // achievable by a simple broadcast.
2207  // 2. The induced permutation for the subsequent vector.transpose that will
2208  // bring us from `broadcastShape` back to he desired `dstShape`.
2209  // If the induced permutation is not the identity, create a vector.transpose.
2210  SmallVector<int64_t> broadcastShape, permutation(dstShape.size(), -1);
2211  broadcastShape.reserve(dstShape.size());
2212  // Consider the example:
2213  // srcShape = 2x4
2214  // dstShape = 1x2x3x4x5
2215  // broadcastedDims = [0, 2, 4]
2216  //
2217  // We want to build:
2218  // broadcastShape = 1x3x5x2x4
2219  // permutation = [0, 2, 4, 1, 3]
2220  // ---V--- -----V-----
2221  // leading broadcast part src shape part
2222  //
2223  // Note that the trailing dims of broadcastShape are exactly the srcShape
2224  // by construction.
2225  // nextSrcShapeDim is used to keep track of where in the permutation the
2226  // "src shape part" occurs.
2227  int64_t nextSrcShapeDim = broadcastedDims.size();
2228  for (int64_t i = 0, e = dstShape.size(); i < e; ++i) {
2229  if (broadcastedDims.contains(i)) {
2230  // 3.a. For each dim in the dst shape, if it is a broadcasted dim,
2231  // bring it to the head of the broadcastShape.
2232  // It will need to be permuted back from `broadcastShape.size() - 1` into
2233  // position `i`.
2234  broadcastShape.push_back(dstShape[i]);
2235  permutation[i] = broadcastShape.size() - 1;
2236  } else {
2237  // 3.b. Otherwise, the dim is not broadcasted, it comes from the src
2238  // shape and needs to be permuted into position `i`.
2239  // Don't touch `broadcastShape` here, the whole srcShape will be
2240  // appended after.
2241  permutation[i] = nextSrcShapeDim++;
2242  }
2243  }
2244  // 3.c. Append the srcShape.
2245  llvm::append_range(broadcastShape, srcVectorType.getShape());
2246 
2247  // Ensure there are no dim-1 broadcasts.
2248  assert(::computeBroadcastedUnitDims(srcVectorType.getShape(), broadcastShape)
2249  .empty() &&
2250  "unexpected dim-1 broadcast");
2251 
2252  VectorType broadcastType = VectorType::get(broadcastShape, elementType);
2253  assert(vector::isBroadcastableTo(value.getType(), broadcastType) ==
2254  vector::BroadcastableToResult::Success &&
2255  "must be broadcastable");
2256  Value res = b.createOrFold<vector::BroadcastOp>(loc, broadcastType, value);
2257  // Step 4. If we find any dimension that indeed needs to be permuted,
2258  // immediately return a new vector.transpose.
2259  for (int64_t i = 0, e = permutation.size(); i < e; ++i)
2260  if (permutation[i] != i)
2261  return b.createOrFold<vector::TransposeOp>(loc, res, permutation);
2262  // Otherwise return res.
2263  return res;
2264 }
2265 
2267 mlir::vector::isBroadcastableTo(Type srcType, VectorType dstVectorType,
2268  std::pair<int, int> *mismatchingDims) {
2269  // Broadcast scalar to vector of the same element type.
2270  if (srcType.isIntOrIndexOrFloat() && dstVectorType &&
2271  getElementTypeOrSelf(srcType) == getElementTypeOrSelf(dstVectorType))
2272  return BroadcastableToResult::Success;
2273  // From now on, only vectors broadcast.
2274  VectorType srcVectorType = llvm::dyn_cast<VectorType>(srcType);
2275  if (!srcVectorType)
2276  return BroadcastableToResult::SourceTypeNotAVector;
2277 
2278  int64_t srcRank = srcVectorType.getRank();
2279  int64_t dstRank = dstVectorType.getRank();
2280  if (srcRank > dstRank)
2281  return BroadcastableToResult::SourceRankHigher;
2282  // Source has an exact match or singleton value for all trailing dimensions
2283  // (all leading dimensions are simply duplicated).
2284  int64_t lead = dstRank - srcRank;
2285  for (int64_t r = 0; r < srcRank; ++r) {
2286  int64_t srcDim = srcVectorType.getDimSize(r);
2287  int64_t dstDim = dstVectorType.getDimSize(lead + r);
2288  if (srcDim != 1 && srcDim != dstDim) {
2289  if (mismatchingDims) {
2290  mismatchingDims->first = srcDim;
2291  mismatchingDims->second = dstDim;
2292  }
2293  return BroadcastableToResult::DimensionMismatch;
2294  }
2295  }
2296 
2297  return BroadcastableToResult::Success;
2298 }
2299 
2301  std::pair<int, int> mismatchingDims;
2303  getSourceType(), getResultVectorType(), &mismatchingDims);
2304  if (res == BroadcastableToResult::Success)
2305  return success();
2306  if (res == BroadcastableToResult::SourceRankHigher)
2307  return emitOpError("source rank higher than destination rank");
2308  if (res == BroadcastableToResult::DimensionMismatch)
2309  return emitOpError("dimension mismatch (")
2310  << mismatchingDims.first << " vs. " << mismatchingDims.second << ")";
2311  if (res == BroadcastableToResult::SourceTypeNotAVector)
2312  return emitOpError("source type is not a vector");
2313  llvm_unreachable("unexpected vector.broadcast op error");
2314 }
2315 
2316 OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
2317  if (getSourceType() == getResultVectorType())
2318  return getSource();
2319  if (!adaptor.getSource())
2320  return {};
2321  auto vectorType = getResultVectorType();
2322  if (llvm::isa<IntegerAttr, FloatAttr>(adaptor.getSource()))
2323  return DenseElementsAttr::get(vectorType, adaptor.getSource());
2324  if (auto attr = llvm::dyn_cast<SplatElementsAttr>(adaptor.getSource()))
2325  return DenseElementsAttr::get(vectorType, attr.getSplatValue<Attribute>());
2326  return {};
2327 }
2328 
2329 namespace {
2330 
2331 // Fold broadcast1(broadcast2(x)) into broadcast1(x).
2332 struct BroadcastFolder : public OpRewritePattern<BroadcastOp> {
2334 
2335  LogicalResult matchAndRewrite(BroadcastOp broadcastOp,
2336  PatternRewriter &rewriter) const override {
2337  auto srcBroadcast = broadcastOp.getSource().getDefiningOp<BroadcastOp>();
2338  if (!srcBroadcast)
2339  return failure();
2340  rewriter.replaceOpWithNewOp<BroadcastOp>(broadcastOp,
2341  broadcastOp.getResultVectorType(),
2342  srcBroadcast.getSource());
2343  return success();
2344  }
2345 };
2346 } // namespace
2347 
2348 void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
2349  MLIRContext *context) {
2350  // BroadcastToShapeCast is not a default canonicalization, it is opt-in by
2351  // calling `populateCastAwayVectorLeadingOneDimPatterns`
2352  results.add<BroadcastFolder>(context);
2353 }
2354 
2355 //===----------------------------------------------------------------------===//
2356 // ShuffleOp
2357 //===----------------------------------------------------------------------===//
2358 
2359 void ShuffleOp::build(OpBuilder &builder, OperationState &result, Value v1,
2360  Value v2, ArrayRef<int64_t> mask) {
2361  build(builder, result, v1, v2, getVectorSubscriptAttr(builder, mask));
2362 }
2363 
2365  VectorType resultType = getResultVectorType();
2366  VectorType v1Type = getV1VectorType();
2367  VectorType v2Type = getV2VectorType();
2368  // Verify ranks.
2369  int64_t resRank = resultType.getRank();
2370  int64_t v1Rank = v1Type.getRank();
2371  int64_t v2Rank = v2Type.getRank();
2372  bool wellFormed0DCase = v1Rank == 0 && v2Rank == 0 && resRank == 1;
2373  bool wellFormedNDCase = v1Rank == resRank && v2Rank == resRank;
2374  if (!wellFormed0DCase && !wellFormedNDCase)
2375  return emitOpError("rank mismatch");
2376 
2377  // Verify all but leading dimension sizes.
2378  for (int64_t r = 1; r < v1Rank; ++r) {
2379  int64_t resDim = resultType.getDimSize(r);
2380  int64_t v1Dim = v1Type.getDimSize(r);
2381  int64_t v2Dim = v2Type.getDimSize(r);
2382  if (resDim != v1Dim || v1Dim != v2Dim)
2383  return emitOpError("dimension mismatch");
2384  }
2385  // Verify mask length.
2386  auto maskAttr = getMask().getValue();
2387  int64_t maskLength = maskAttr.size();
2388  if (maskLength <= 0)
2389  return emitOpError("invalid mask length");
2390  if (maskLength != resultType.getDimSize(0))
2391  return emitOpError("mask length mismatch");
2392  // Verify all indices.
2393  int64_t indexSize = (v1Type.getRank() == 0 ? 1 : v1Type.getDimSize(0)) +
2394  (v2Type.getRank() == 0 ? 1 : v2Type.getDimSize(0));
2395  for (const auto &en : llvm::enumerate(maskAttr)) {
2396  auto attr = llvm::dyn_cast<IntegerAttr>(en.value());
2397  if (!attr || attr.getInt() < 0 || attr.getInt() >= indexSize)
2398  return emitOpError("mask index #") << (en.index() + 1) << " out of range";
2399  }
2400  return success();
2401 }
2402 
2404 ShuffleOp::inferReturnTypes(MLIRContext *, std::optional<Location>,
2405  ShuffleOp::Adaptor adaptor,
2406  SmallVectorImpl<Type> &inferredReturnTypes) {
2407  auto v1Type = llvm::cast<VectorType>(adaptor.getV1().getType());
2408  auto v1Rank = v1Type.getRank();
2409  // Construct resulting type: leading dimension matches mask
2410  // length, all trailing dimensions match the operands.
2412  shape.reserve(v1Rank);
2413  shape.push_back(std::max<size_t>(1, adaptor.getMask().size()));
2414  // In the 0-D case there is no trailing shape to append.
2415  if (v1Rank > 0)
2416  llvm::append_range(shape, v1Type.getShape().drop_front());
2417  inferredReturnTypes.push_back(
2418  VectorType::get(shape, v1Type.getElementType()));
2419  return success();
2420 }
2421 
2422 static bool isStepIndexArray(ArrayAttr idxArr, uint64_t begin, size_t width) {
2423  uint64_t expected = begin;
2424  return idxArr.size() == width &&
2425  llvm::all_of(idxArr.getAsValueRange<IntegerAttr>(),
2426  [&expected](auto attr) {
2427  return attr.getZExtValue() == expected++;
2428  });
2429 }
2430 
2431 OpFoldResult vector::ShuffleOp::fold(FoldAdaptor adaptor) {
2432  VectorType v1Type = getV1VectorType();
2433  // For consistency: 0-D shuffle return type is 1-D, this cannot be a folding
2434  // but must be a canonicalization into a vector.broadcast.
2435  if (v1Type.getRank() == 0)
2436  return {};
2437 
2438  // fold shuffle V1, V2, [0, 1, 2, 3] : <4xi32>, <2xi32> -> V1
2439  if (!v1Type.isScalable() &&
2440  isStepIndexArray(getMask(), 0, v1Type.getDimSize(0)))
2441  return getV1();
2442  // fold shuffle V1, V2, [4, 5] : <4xi32>, <2xi32> -> V2
2443  if (!getV1VectorType().isScalable() && !getV2VectorType().isScalable() &&
2444  isStepIndexArray(getMask(), getV1VectorType().getDimSize(0),
2445  getV2VectorType().getDimSize(0)))
2446  return getV2();
2447 
2448  Attribute lhs = adaptor.getV1(), rhs = adaptor.getV2();
2449  if (!lhs || !rhs)
2450  return {};
2451 
2452  auto lhsType =
2453  llvm::cast<VectorType>(llvm::cast<DenseElementsAttr>(lhs).getType());
2454  // Only support 1-D for now to avoid complicated n-D DenseElementsAttr
2455  // manipulation.
2456  if (lhsType.getRank() != 1)
2457  return {};
2458  int64_t lhsSize = lhsType.getDimSize(0);
2459 
2460  SmallVector<Attribute> results;
2461  auto lhsElements = llvm::cast<DenseElementsAttr>(lhs).getValues<Attribute>();
2462  auto rhsElements = llvm::cast<DenseElementsAttr>(rhs).getValues<Attribute>();
2463  for (const auto &index : this->getMask().getAsValueRange<IntegerAttr>()) {
2464  int64_t i = index.getZExtValue();
2465  if (i >= lhsSize) {
2466  results.push_back(rhsElements[i - lhsSize]);
2467  } else {
2468  results.push_back(lhsElements[i]);
2469  }
2470  }
2471 
2472  return DenseElementsAttr::get(getResultVectorType(), results);
2473 }
2474 
2475 namespace {
2476 
2477 // Pattern to rewrite a 0-D shuffle with [0] or [1] mask returning a 1-D vector
2478 // to a broadcast.
2479 struct Canonicalize0DShuffleOp : public OpRewritePattern<ShuffleOp> {
2481 
2482  LogicalResult matchAndRewrite(ShuffleOp shuffleOp,
2483  PatternRewriter &rewriter) const override {
2484  VectorType v1VectorType = shuffleOp.getV1VectorType();
2485  ArrayAttr mask = shuffleOp.getMask();
2486  if (v1VectorType.getRank() > 0)
2487  return failure();
2488  if (mask.size() != 1)
2489  return failure();
2490  VectorType resType = VectorType::Builder(v1VectorType).setShape({1});
2491  if (llvm::cast<IntegerAttr>(mask[0]).getInt() == 0)
2492  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(shuffleOp, resType,
2493  shuffleOp.getV1());
2494  else
2495  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(shuffleOp, resType,
2496  shuffleOp.getV2());
2497  return success();
2498  }
2499 };
2500 
2501 /// Pattern to rewrite a ShuffleOp(SplatOp, SplatOp) to SplatOp.
2502 class ShuffleSplat final : public OpRewritePattern<ShuffleOp> {
2503 public:
2505 
2506  LogicalResult matchAndRewrite(ShuffleOp op,
2507  PatternRewriter &rewriter) const override {
2508  auto v1Splat = op.getV1().getDefiningOp<SplatOp>();
2509  auto v2Splat = op.getV2().getDefiningOp<SplatOp>();
2510 
2511  if (!v1Splat || !v2Splat)
2512  return failure();
2513 
2514  if (v1Splat.getInput() != v2Splat.getInput())
2515  return failure();
2516 
2517  rewriter.replaceOpWithNewOp<SplatOp>(op, op.getType(), v1Splat.getInput());
2518  return success();
2519  }
2520 };
2521 
2522 /// Pattern to rewrite a fixed-size interleave via vector.shuffle to
2523 /// vector.interleave.
2524 class ShuffleInterleave : public OpRewritePattern<ShuffleOp> {
2525 public:
2527 
2528  LogicalResult matchAndRewrite(ShuffleOp op,
2529  PatternRewriter &rewriter) const override {
2530  VectorType resultType = op.getResultVectorType();
2531  if (resultType.isScalable())
2532  return rewriter.notifyMatchFailure(
2533  op, "ShuffleOp can't represent a scalable interleave");
2534 
2535  if (resultType.getRank() != 1)
2536  return rewriter.notifyMatchFailure(
2537  op, "ShuffleOp can't represent an n-D interleave");
2538 
2539  VectorType sourceType = op.getV1VectorType();
2540  if (sourceType != op.getV2VectorType() ||
2541  sourceType.getNumElements() * 2 != resultType.getNumElements()) {
2542  return rewriter.notifyMatchFailure(
2543  op, "ShuffleOp types don't match an interleave");
2544  }
2545 
2546  ArrayAttr shuffleMask = op.getMask();
2547  int64_t resultVectorSize = resultType.getNumElements();
2548  for (int i = 0, e = resultVectorSize / 2; i < e; ++i) {
2549  int64_t maskValueA = cast<IntegerAttr>(shuffleMask[i * 2]).getInt();
2550  int64_t maskValueB = cast<IntegerAttr>(shuffleMask[(i * 2) + 1]).getInt();
2551  if (maskValueA != i || maskValueB != (resultVectorSize / 2) + i)
2552  return rewriter.notifyMatchFailure(op,
2553  "ShuffleOp mask not interleaving");
2554  }
2555 
2556  rewriter.replaceOpWithNewOp<InterleaveOp>(op, op.getV1(), op.getV2());
2557  return success();
2558  }
2559 };
2560 
2561 } // namespace
2562 
2563 void ShuffleOp::getCanonicalizationPatterns(RewritePatternSet &results,
2564  MLIRContext *context) {
2565  results.add<ShuffleSplat, ShuffleInterleave, Canonicalize0DShuffleOp>(
2566  context);
2567 }
2568 
2569 //===----------------------------------------------------------------------===//
2570 // InsertElementOp
2571 //===----------------------------------------------------------------------===//
2572 
2573 void InsertElementOp::build(OpBuilder &builder, OperationState &result,
2574  Value source, Value dest) {
2575  build(builder, result, source, dest, {});
2576 }
2577 
2579  auto dstVectorType = getDestVectorType();
2580  if (dstVectorType.getRank() == 0) {
2581  if (getPosition())
2582  return emitOpError("expected position to be empty with 0-D vector");
2583  return success();
2584  }
2585  if (dstVectorType.getRank() != 1)
2586  return emitOpError("unexpected >1 vector rank");
2587  if (!getPosition())
2588  return emitOpError("expected position for 1-D vector");
2589  return success();
2590 }
2591 
2592 OpFoldResult vector::InsertElementOp::fold(FoldAdaptor adaptor) {
2593  // Skip the 0-D vector here.
2594  if (!adaptor.getPosition())
2595  return {};
2596 
2597  auto src = dyn_cast_or_null<TypedAttr>(adaptor.getSource());
2598  auto dst = dyn_cast_or_null<DenseElementsAttr>(adaptor.getDest());
2599  auto pos = dyn_cast_or_null<IntegerAttr>(adaptor.getPosition());
2600  if (!src || !dst || !pos)
2601  return {};
2602 
2603  if (src.getType() != getDestVectorType().getElementType())
2604  return {};
2605 
2606  auto dstElements = dst.getValues<Attribute>();
2607 
2608  SmallVector<Attribute> results(dstElements);
2609 
2610  uint64_t posIdx = pos.getInt();
2611  if (posIdx >= results.size())
2612  return {};
2613  results[posIdx] = src;
2614 
2615  return DenseElementsAttr::get(getDestVectorType(), results);
2616 }
2617 
2618 //===----------------------------------------------------------------------===//
2619 // InsertOp
2620 //===----------------------------------------------------------------------===//
2621 
2622 void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
2623  Value source, Value dest, int64_t position) {
2624  build(builder, result, source, dest, ArrayRef<int64_t>{position});
2625 }
2626 
2627 void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
2628  Value source, Value dest, OpFoldResult position) {
2629  build(builder, result, source, dest, ArrayRef<OpFoldResult>{position});
2630 }
2631 
2632 void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
2633  Value source, Value dest,
2634  ArrayRef<int64_t> position) {
2635  SmallVector<OpFoldResult> posVals;
2636  posVals.reserve(position.size());
2637  llvm::transform(position, std::back_inserter(posVals),
2638  [&](int64_t pos) { return builder.getI64IntegerAttr(pos); });
2639  build(builder, result, source, dest, posVals);
2640 }
2641 
2642 void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
2643  Value source, Value dest,
2644  ArrayRef<OpFoldResult> position) {
2645  SmallVector<int64_t> staticPos;
2646  SmallVector<Value> dynamicPos;
2647  dispatchIndexOpFoldResults(position, dynamicPos, staticPos);
2648  build(builder, result, source, dest, dynamicPos,
2649  builder.getDenseI64ArrayAttr(staticPos));
2650 }
2651 
2653  SmallVector<OpFoldResult> position = getMixedPosition();
2654  auto destVectorType = getDestVectorType();
2655  if (position.size() > static_cast<unsigned>(destVectorType.getRank()))
2656  return emitOpError(
2657  "expected position attribute of rank no greater than dest vector rank");
2658  auto srcVectorType = llvm::dyn_cast<VectorType>(getSourceType());
2659  if (srcVectorType &&
2660  (static_cast<unsigned>(srcVectorType.getRank()) + position.size() !=
2661  static_cast<unsigned>(destVectorType.getRank())))
2662  return emitOpError("expected position attribute rank + source rank to "
2663  "match dest vector rank");
2664  if (!srcVectorType &&
2665  (position.size() != static_cast<unsigned>(destVectorType.getRank())))
2666  return emitOpError(
2667  "expected position attribute rank to match the dest vector rank");
2668  for (auto [idx, pos] : llvm::enumerate(position)) {
2669  if (auto attr = pos.dyn_cast<Attribute>()) {
2670  int64_t constIdx = cast<IntegerAttr>(attr).getInt();
2671  if (constIdx < 0 || constIdx >= destVectorType.getDimSize(idx)) {
2672  return emitOpError("expected position attribute #")
2673  << (idx + 1)
2674  << " to be a non-negative integer smaller than the "
2675  "corresponding "
2676  "dest vector dimension";
2677  }
2678  }
2679  }
2680  return success();
2681 }
2682 
2683 namespace {
2684 
2685 // If insertOp is only inserting unit dimensions it can be transformed to a
2686 // broadcast.
2687 class InsertToBroadcast final : public OpRewritePattern<InsertOp> {
2688 public:
2690 
2691  LogicalResult matchAndRewrite(InsertOp insertOp,
2692  PatternRewriter &rewriter) const override {
2693  auto srcVecType = llvm::dyn_cast<VectorType>(insertOp.getSourceType());
2694  if (!srcVecType || insertOp.getDestVectorType().getNumElements() !=
2695  srcVecType.getNumElements())
2696  return failure();
2697  rewriter.replaceOpWithNewOp<BroadcastOp>(
2698  insertOp, insertOp.getDestVectorType(), insertOp.getSource());
2699  return success();
2700  }
2701 };
2702 
2703 /// Pattern to rewrite a InsertOp(SplatOp, SplatOp) to SplatOp.
2704 class InsertSplatToSplat final : public OpRewritePattern<InsertOp> {
2705 public:
2707 
2708  LogicalResult matchAndRewrite(InsertOp op,
2709  PatternRewriter &rewriter) const override {
2710  auto srcSplat = op.getSource().getDefiningOp<SplatOp>();
2711  auto dstSplat = op.getDest().getDefiningOp<SplatOp>();
2712 
2713  if (!srcSplat || !dstSplat)
2714  return failure();
2715 
2716  if (srcSplat.getInput() != dstSplat.getInput())
2717  return failure();
2718 
2719  rewriter.replaceOpWithNewOp<SplatOp>(op, op.getType(), srcSplat.getInput());
2720  return success();
2721  }
2722 };
2723 
2724 // Pattern to rewrite a InsertOp(ConstantOp into ConstantOp) -> ConstantOp.
2725 class InsertOpConstantFolder final : public OpRewritePattern<InsertOp> {
2726 public:
2728 
2729  // Do not create constants with more than `vectorSizeFoldThreashold` elements,
2730  // unless the source vector constant has a single use.
2731  static constexpr int64_t vectorSizeFoldThreshold = 256;
2732 
2733  LogicalResult matchAndRewrite(InsertOp op,
2734  PatternRewriter &rewriter) const override {
2735  // TODO: Canonicalization for dynamic position not implemented yet.
2736  if (op.hasDynamicPosition())
2737  return failure();
2738 
2739  // Return if 'InsertOp' operand is not defined by a compatible vector
2740  // ConstantOp.
2741  TypedValue<VectorType> destVector = op.getDest();
2742  Attribute vectorDestCst;
2743  if (!matchPattern(destVector, m_Constant(&vectorDestCst)))
2744  return failure();
2745 
2746  VectorType destTy = destVector.getType();
2747  if (destTy.isScalable())
2748  return failure();
2749 
2750  // Make sure we do not create too many large constants.
2751  if (destTy.getNumElements() > vectorSizeFoldThreshold &&
2752  !destVector.hasOneUse())
2753  return failure();
2754 
2755  auto denseDest = llvm::cast<DenseElementsAttr>(vectorDestCst);
2756 
2757  Value sourceValue = op.getSource();
2758  Attribute sourceCst;
2759  if (!matchPattern(sourceValue, m_Constant(&sourceCst)))
2760  return failure();
2761 
2762  // Calculate the linearized position of the continuous chunk of elements to
2763  // insert.
2764  llvm::SmallVector<int64_t> completePositions(destTy.getRank(), 0);
2765  copy(op.getStaticPosition(), completePositions.begin());
2766  int64_t insertBeginPosition =
2767  linearize(completePositions, computeStrides(destTy.getShape()));
2768 
2769  SmallVector<Attribute> insertedValues;
2770  if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(sourceCst))
2771  llvm::append_range(insertedValues, denseSource.getValues<Attribute>());
2772  else
2773  insertedValues.push_back(sourceCst);
2774 
2775  auto allValues = llvm::to_vector(denseDest.getValues<Attribute>());
2776  copy(insertedValues, allValues.begin() + insertBeginPosition);
2777  auto newAttr = DenseElementsAttr::get(destTy, allValues);
2778 
2779  rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);
2780  return success();
2781  }
2782 };
2783 
2784 } // namespace
2785 
2786 void InsertOp::getCanonicalizationPatterns(RewritePatternSet &results,
2787  MLIRContext *context) {
2788  results.add<InsertToBroadcast, BroadcastFolder, InsertSplatToSplat,
2789  InsertOpConstantFolder>(context);
2790 }
2791 
2792 // Eliminates insert operations that produce values identical to their source
2793 // value. This happens when the source and destination vectors have identical
2794 // sizes.
2795 OpFoldResult vector::InsertOp::fold(FoldAdaptor adaptor) {
2796  if (getNumIndices() == 0)
2797  return getSource();
2798  return {};
2799 }
2800 
2801 //===----------------------------------------------------------------------===//
2802 // InsertStridedSliceOp
2803 //===----------------------------------------------------------------------===//
2804 
2805 void InsertStridedSliceOp::build(OpBuilder &builder, OperationState &result,
2806  Value source, Value dest,
2807  ArrayRef<int64_t> offsets,
2808  ArrayRef<int64_t> strides) {
2809  result.addOperands({source, dest});
2810  auto offsetsAttr = getVectorSubscriptAttr(builder, offsets);
2811  auto stridesAttr = getVectorSubscriptAttr(builder, strides);
2812  result.addTypes(dest.getType());
2813  result.addAttribute(InsertStridedSliceOp::getOffsetsAttrName(result.name),
2814  offsetsAttr);
2815  result.addAttribute(InsertStridedSliceOp::getStridesAttrName(result.name),
2816  stridesAttr);
2817 }
2818 
2819 // TODO: Should be moved to Tablegen ConfinedAttr attributes.
2820 template <typename OpType>
2822  ArrayAttr arrayAttr,
2823  ArrayRef<int64_t> shape,
2824  StringRef attrName) {
2825  if (arrayAttr.size() > shape.size())
2826  return op.emitOpError("expected ")
2827  << attrName << " attribute of rank no greater than vector rank";
2828  return success();
2829 }
2830 
2831 // Returns true if all integers in `arrayAttr` are in the half-open [min, max}
2832 // interval. If `halfOpen` is true then the admissible interval is [min, max).
2833 // Otherwise, the admissible interval is [min, max].
2834 template <typename OpType>
2835 static LogicalResult
2836 isIntegerArrayAttrConfinedToRange(OpType op, ArrayAttr arrayAttr, int64_t min,
2837  int64_t max, StringRef attrName,
2838  bool halfOpen = true) {
2839  for (auto attr : arrayAttr) {
2840  auto val = llvm::cast<IntegerAttr>(attr).getInt();
2841  auto upper = max;
2842  if (!halfOpen)
2843  upper += 1;
2844  if (val < min || val >= upper)
2845  return op.emitOpError("expected ") << attrName << " to be confined to ["
2846  << min << ", " << upper << ")";
2847  }
2848  return success();
2849 }
2850 
2851 // Returns true if all integers in `arrayAttr` are in the half-open [min, max}
2852 // interval. If `halfOpen` is true then the admissible interval is [min, max).
2853 // Otherwise, the admissible interval is [min, max].
2854 template <typename OpType>
2855 static LogicalResult
2856 isIntegerArrayAttrConfinedToShape(OpType op, ArrayAttr arrayAttr,
2857  ArrayRef<int64_t> shape, StringRef attrName,
2858  bool halfOpen = true, int64_t min = 0) {
2859  for (auto [index, attrDimPair] :
2860  llvm::enumerate(llvm::zip_first(arrayAttr, shape))) {
2861  int64_t val = llvm::cast<IntegerAttr>(std::get<0>(attrDimPair)).getInt();
2862  int64_t max = std::get<1>(attrDimPair);
2863  if (!halfOpen)
2864  max += 1;
2865  if (val < min || val >= max)
2866  return op.emitOpError("expected ")
2867  << attrName << " dimension " << index << " to be confined to ["
2868  << min << ", " << max << ")";
2869  }
2870  return success();
2871 }
2872 
2873 // Returns true if, for all indices i = 0..shape.size()-1, val is in the
2874 // [min, max} interval:
2875 // val = `arrayAttr1[i]` + `arrayAttr2[i]`,
2876 // If `halfOpen` is true then the admissible interval is [min, max). Otherwise,
2877 // the admissible interval is [min, max].
2878 template <typename OpType>
2880  OpType op, ArrayAttr arrayAttr1, ArrayAttr arrayAttr2,
2881  ArrayRef<int64_t> shape, StringRef attrName1, StringRef attrName2,
2882  bool halfOpen = true, int64_t min = 1) {
2883  assert(arrayAttr1.size() <= shape.size());
2884  assert(arrayAttr2.size() <= shape.size());
2885  for (auto [index, it] :
2886  llvm::enumerate(llvm::zip(arrayAttr1, arrayAttr2, shape))) {
2887  auto val1 = llvm::cast<IntegerAttr>(std::get<0>(it)).getInt();
2888  auto val2 = llvm::cast<IntegerAttr>(std::get<1>(it)).getInt();
2889  int64_t max = std::get<2>(it);
2890  if (!halfOpen)
2891  max += 1;
2892  if (val1 + val2 < 0 || val1 + val2 >= max)
2893  return op.emitOpError("expected sum(")
2894  << attrName1 << ", " << attrName2 << ") dimension " << index
2895  << " to be confined to [" << min << ", " << max << ")";
2896  }
2897  return success();
2898 }
2899 
2900 static ArrayAttr makeI64ArrayAttr(ArrayRef<int64_t> values,
2901  MLIRContext *context) {
2902  auto attrs = llvm::map_range(values, [context](int64_t v) -> Attribute {
2903  return IntegerAttr::get(IntegerType::get(context, 64), APInt(64, v));
2904  });
2905  return ArrayAttr::get(context, llvm::to_vector<8>(attrs));
2906 }
2907 
2909  auto sourceVectorType = getSourceVectorType();
2910  auto destVectorType = getDestVectorType();
2911  auto offsets = getOffsetsAttr();
2912  auto strides = getStridesAttr();
2913  if (offsets.size() != static_cast<unsigned>(destVectorType.getRank()))
2914  return emitOpError(
2915  "expected offsets of same size as destination vector rank");
2916  if (strides.size() != static_cast<unsigned>(sourceVectorType.getRank()))
2917  return emitOpError("expected strides of same size as source vector rank");
2918  if (sourceVectorType.getRank() > destVectorType.getRank())
2919  return emitOpError(
2920  "expected source rank to be no greater than destination rank");
2921 
2922  auto sourceShape = sourceVectorType.getShape();
2923  auto destShape = destVectorType.getShape();
2924  SmallVector<int64_t, 4> sourceShapeAsDestShape(
2925  destShape.size() - sourceShape.size(), 0);
2926  sourceShapeAsDestShape.append(sourceShape.begin(), sourceShape.end());
2927  auto offName = InsertStridedSliceOp::getOffsetsAttrName();
2928  auto stridesName = InsertStridedSliceOp::getStridesAttrName();
2929  if (failed(isIntegerArrayAttrConfinedToShape(*this, offsets, destShape,
2930  offName)) ||
2931  failed(isIntegerArrayAttrConfinedToRange(*this, strides, /*min=*/1,
2932  /*max=*/1, stridesName,
2933  /*halfOpen=*/false)) ||
2935  *this, offsets,
2936  makeI64ArrayAttr(sourceShapeAsDestShape, getContext()), destShape,
2937  offName, "source vector shape",
2938  /*halfOpen=*/false, /*min=*/1)))
2939  return failure();
2940 
2941  unsigned rankDiff = destShape.size() - sourceShape.size();
2942  for (unsigned idx = 0; idx < sourceShape.size(); ++idx) {
2943  if (sourceVectorType.getScalableDims()[idx] !=
2944  destVectorType.getScalableDims()[idx + rankDiff]) {
2945  return emitOpError("mismatching scalable flags (at source vector idx=")
2946  << idx << ")";
2947  }
2948  if (sourceVectorType.getScalableDims()[idx]) {
2949  auto sourceSize = sourceShape[idx];
2950  auto destSize = destShape[idx + rankDiff];
2951  if (sourceSize != destSize) {
2952  return emitOpError("expected size at idx=")
2953  << idx
2954  << (" to match the corresponding base size from the input "
2955  "vector (")
2956  << sourceSize << (" vs ") << destSize << (")");
2957  }
2958  }
2959  }
2960 
2961  return success();
2962 }
2963 
2964 namespace {
2965 /// Pattern to rewrite an InsertStridedSliceOp(SplatOp(X):src_type,
2966 /// SplatOp(X):dst_type) to SplatOp(X):dst_type.
2967 class FoldInsertStridedSliceSplat final
2968  : public OpRewritePattern<InsertStridedSliceOp> {
2969 public:
2971 
2972  LogicalResult matchAndRewrite(InsertStridedSliceOp insertStridedSliceOp,
2973  PatternRewriter &rewriter) const override {
2974  auto srcSplatOp =
2975  insertStridedSliceOp.getSource().getDefiningOp<vector::SplatOp>();
2976  auto destSplatOp =
2977  insertStridedSliceOp.getDest().getDefiningOp<vector::SplatOp>();
2978 
2979  if (!srcSplatOp || !destSplatOp)
2980  return failure();
2981 
2982  if (srcSplatOp.getInput() != destSplatOp.getInput())
2983  return failure();
2984 
2985  rewriter.replaceOp(insertStridedSliceOp, insertStridedSliceOp.getDest());
2986  return success();
2987  }
2988 };
2989 
2990 /// Pattern to rewrite an InsertStridedSliceOp(ExtractStridedSliceOp(dst), dst)
2991 /// to dst.
2992 class FoldInsertStridedSliceOfExtract final
2993  : public OpRewritePattern<InsertStridedSliceOp> {
2994 public:
2996 
2997  LogicalResult matchAndRewrite(InsertStridedSliceOp insertStridedSliceOp,
2998  PatternRewriter &rewriter) const override {
2999  auto extractStridedSliceOp =
3000  insertStridedSliceOp.getSource()
3001  .getDefiningOp<vector::ExtractStridedSliceOp>();
3002 
3003  if (!extractStridedSliceOp)
3004  return failure();
3005 
3006  if (extractStridedSliceOp.getOperand() != insertStridedSliceOp.getDest())
3007  return failure();
3008 
3009  // Check if have the same strides and offsets.
3010  if (extractStridedSliceOp.getStrides() !=
3011  insertStridedSliceOp.getStrides() ||
3012  extractStridedSliceOp.getOffsets() != insertStridedSliceOp.getOffsets())
3013  return failure();
3014 
3015  rewriter.replaceOp(insertStridedSliceOp, insertStridedSliceOp.getDest());
3016  return success();
3017  }
3018 };
3019 
3020 // Pattern to rewrite an InsertStridedSliceOp(ConstantOp into ConstantOp) ->
3021 // ConstantOp.
3022 class InsertStridedSliceConstantFolder final
3023  : public OpRewritePattern<InsertStridedSliceOp> {
3024 public:
3026 
3027  // Do not create constants with more than `vectorSizeFoldThreashold` elements,
3028  // unless the source vector constant has a single use.
3029  static constexpr int64_t vectorSizeFoldThreshold = 256;
3030 
3031  LogicalResult matchAndRewrite(InsertStridedSliceOp op,
3032  PatternRewriter &rewriter) const override {
3033  // Return if 'InsertOp' operand is not defined by a compatible vector
3034  // ConstantOp.
3035  TypedValue<VectorType> destVector = op.getDest();
3036  Attribute vectorDestCst;
3037  if (!matchPattern(destVector, m_Constant(&vectorDestCst)))
3038  return failure();
3039 
3040  VectorType destTy = destVector.getType();
3041  if (destTy.isScalable())
3042  return failure();
3043 
3044  // Make sure we do not create too many large constants.
3045  if (destTy.getNumElements() > vectorSizeFoldThreshold &&
3046  !destVector.hasOneUse())
3047  return failure();
3048 
3049  auto denseDest = llvm::cast<DenseElementsAttr>(vectorDestCst);
3050 
3051  TypedValue<VectorType> sourceValue = op.getSource();
3052  Attribute sourceCst;
3053  if (!matchPattern(sourceValue, m_Constant(&sourceCst)))
3054  return failure();
3055 
3056  // TODO: Handle non-unit strides when they become available.
3057  if (op.hasNonUnitStrides())
3058  return failure();
3059 
3060  VectorType sliceVecTy = sourceValue.getType();
3061  ArrayRef<int64_t> sliceShape = sliceVecTy.getShape();
3062  int64_t rankDifference = destTy.getRank() - sliceVecTy.getRank();
3063  SmallVector<int64_t, 4> offsets = getI64SubArray(op.getOffsets());
3064  SmallVector<int64_t, 4> destStrides = computeStrides(destTy.getShape());
3065 
3066  // Calcualte the destination element indices by enumerating all slice
3067  // positions within the destination and linearizing them. The enumeration
3068  // order is lexicographic which yields a sequence of monotonically
3069  // increasing linearized position indices.
3070  // Because the destination may have higher dimensionality then the slice,
3071  // we keep track of two overlapping sets of positions and offsets.
3072  auto denseSlice = llvm::cast<DenseElementsAttr>(sourceCst);
3073  auto sliceValuesIt = denseSlice.value_begin<Attribute>();
3074  auto newValues = llvm::to_vector(denseDest.getValues<Attribute>());
3075  SmallVector<int64_t> currDestPosition(offsets.begin(), offsets.end());
3076  MutableArrayRef<int64_t> currSlicePosition(
3077  currDestPosition.begin() + rankDifference, currDestPosition.end());
3078  ArrayRef<int64_t> sliceOffsets(offsets.begin() + rankDifference,
3079  offsets.end());
3080  do {
3081  int64_t linearizedPosition = linearize(currDestPosition, destStrides);
3082  assert(linearizedPosition < destTy.getNumElements() && "Invalid index");
3083  assert(sliceValuesIt != denseSlice.value_end<Attribute>() &&
3084  "Invalid slice element");
3085  newValues[linearizedPosition] = *sliceValuesIt;
3086  ++sliceValuesIt;
3087  } while (succeeded(
3088  incSlicePosition(currSlicePosition, sliceShape, sliceOffsets)));
3089 
3090  auto newAttr = DenseElementsAttr::get(destTy, newValues);
3091  rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);
3092  return success();
3093  }
3094 };
3095 
3096 } // namespace
3097 
3098 void vector::InsertStridedSliceOp::getCanonicalizationPatterns(
3099  RewritePatternSet &results, MLIRContext *context) {
3100  results.add<FoldInsertStridedSliceSplat, FoldInsertStridedSliceOfExtract,
3101  InsertStridedSliceConstantFolder>(context);
3102 }
3103 
3104 OpFoldResult InsertStridedSliceOp::fold(FoldAdaptor adaptor) {
3105  if (getSourceVectorType() == getDestVectorType())
3106  return getSource();
3107  return {};
3108 }
3109 
3110 //===----------------------------------------------------------------------===//
3111 // OuterProductOp
3112 //===----------------------------------------------------------------------===//
3113 
3114 /// Build an op without mask, use the type of `acc` as the return type.
3115 void OuterProductOp::build(OpBuilder &builder, OperationState &result,
3116  Value lhs, Value rhs, Value acc) {
3117  result.addOperands({lhs, rhs, acc});
3118  result.addTypes(acc.getType());
3119 }
3120 
3122  p << " " << getLhs() << ", " << getRhs();
3123  if (getAcc()) {
3124  p << ", " << getAcc();
3125  p.printOptionalAttrDict((*this)->getAttrs());
3126  }
3127  p << " : " << getLhs().getType() << ", " << getRhs().getType();
3128 }
3129 
3132  Type tLHS, tRHS;
3133  if (parser.parseOperandList(operandsInfo) ||
3134  parser.parseOptionalAttrDict(result.attributes) ||
3135  parser.parseColonType(tLHS) || parser.parseComma() ||
3136  parser.parseType(tRHS))
3137  return failure();
3138  if (operandsInfo.size() < 2)
3139  return parser.emitError(parser.getNameLoc(),
3140  "expected at least 2 operands");
3141  VectorType vLHS = llvm::dyn_cast<VectorType>(tLHS);
3142  VectorType vRHS = llvm::dyn_cast<VectorType>(tRHS);
3143  if (!vLHS)
3144  return parser.emitError(parser.getNameLoc(),
3145  "expected vector type for operand #1");
3146 
3147  VectorType resType;
3148  if (vRHS) {
3149  SmallVector<bool> scalableDimsRes{vLHS.getScalableDims()[0],
3150  vRHS.getScalableDims()[0]};
3151  resType = VectorType::get({vLHS.getDimSize(0), vRHS.getDimSize(0)},
3152  vLHS.getElementType(), scalableDimsRes);
3153  } else {
3154  // Scalar RHS operand
3155  SmallVector<bool> scalableDimsRes{vLHS.getScalableDims()[0]};
3156  resType = VectorType::get({vLHS.getDimSize(0)}, vLHS.getElementType(),
3157  scalableDimsRes);
3158  }
3159 
3160  if (!result.attributes.get(OuterProductOp::getKindAttrName(result.name))) {
3161  result.attributes.append(
3162  OuterProductOp::getKindAttrName(result.name),
3164  OuterProductOp::getDefaultKind()));
3165  }
3166 
3167  return failure(
3168  parser.resolveOperand(operandsInfo[0], tLHS, result.operands) ||
3169  parser.resolveOperand(operandsInfo[1], tRHS, result.operands) ||
3170  (operandsInfo.size() > 2 &&
3171  parser.resolveOperand(operandsInfo[2], resType, result.operands)) ||
3172  parser.addTypeToList(resType, result.types));
3173 }
3174 
3176  Type tRHS = getOperandTypeRHS();
3177  VectorType vLHS = getOperandVectorTypeLHS(),
3178  vRHS = llvm::dyn_cast<VectorType>(tRHS),
3179  vACC = getOperandVectorTypeACC(), vRES = getResultVectorType();
3180 
3181  if (vLHS.getRank() != 1)
3182  return emitOpError("expected 1-d vector for operand #1");
3183 
3184  if (vRHS) {
3185  // Proper OUTER operation.
3186  if (vRHS.getRank() != 1)
3187  return emitOpError("expected 1-d vector for operand #2");
3188  if (vRES.getRank() != 2)
3189  return emitOpError("expected 2-d vector result");
3190  if (vLHS.getDimSize(0) != vRES.getDimSize(0))
3191  return emitOpError("expected #1 operand dim to match result dim #1");
3192  if (vRHS.getDimSize(0) != vRES.getDimSize(1))
3193  return emitOpError("expected #2 operand dim to match result dim #2");
3194  if (vLHS.isScalable() && !vRHS.isScalable()) {
3195  // This restriction reflects what's currently supported in terms of
3196  // scalable vectors. However, we could relax this if there's a use case.
3197  return emitOpError(
3198  "expected either both or only #2 operand dim to be scalable");
3199  }
3200  } else {
3201  // An AXPY operation.
3202  if (vRES.getRank() != 1)
3203  return emitOpError("expected 1-d vector result");
3204  if (vLHS.getDimSize(0) != vRES.getDimSize(0))
3205  return emitOpError("expected #1 operand dim to match result dim #1");
3206  }
3207 
3208  if (vACC && vACC != vRES)
3209  return emitOpError("expected operand #3 of same type as result type");
3210 
3211  // Verify supported combining kind.
3212  if (!isSupportedCombiningKind(getKind(), vRES.getElementType()))
3213  return emitOpError("unsupported outerproduct type");
3214 
3215  return success();
3216 }
3217 
3218 // MaskableOpInterface methods.
3219 
3220 /// Returns the mask type expected by this operation. Mostly used for
3221 /// verification purposes. It requires the operation to be vectorized."
3222 Type OuterProductOp::getExpectedMaskType() {
3223  auto vecType = this->getResultVectorType();
3224  return VectorType::get(vecType.getShape(),
3225  IntegerType::get(vecType.getContext(), /*width=*/1),
3226  vecType.getScalableDims());
3227 }
3228 
3229 //===----------------------------------------------------------------------===//
3230 // ReshapeOp
3231 //===----------------------------------------------------------------------===//
3232 
3234  // Verify that rank(numInputs/outputs) + numFixedVec dim matches vec rank.
3235  auto inputVectorType = getInputVectorType();
3236  auto outputVectorType = getOutputVectorType();
3237  int64_t inputShapeRank = getNumInputShapeSizes();
3238  int64_t outputShapeRank = getNumOutputShapeSizes();
3239  SmallVector<int64_t, 4> fixedVectorSizes;
3240  getFixedVectorSizes(fixedVectorSizes);
3241  int64_t numFixedVectorSizes = fixedVectorSizes.size();
3242 
3243  if (inputVectorType.getRank() != inputShapeRank + numFixedVectorSizes)
3244  return emitError("invalid input shape for vector type ") << inputVectorType;
3245 
3246  if (outputVectorType.getRank() != outputShapeRank + numFixedVectorSizes)
3247  return emitError("invalid output shape for vector type ")
3248  << outputVectorType;
3249 
3250  // Verify that the 'fixedVectorSizes' match an input/output vector shape
3251  // suffix.
3252  unsigned inputVectorRank = inputVectorType.getRank();
3253  for (unsigned i = 0; i < numFixedVectorSizes; ++i) {
3254  unsigned index = inputVectorRank - numFixedVectorSizes - i;
3255  if (fixedVectorSizes[i] != inputVectorType.getShape()[index])
3256  return emitError("fixed vector size must match input vector for dim ")
3257  << i;
3258  }
3259 
3260  unsigned outputVectorRank = outputVectorType.getRank();
3261  for (unsigned i = 0; i < numFixedVectorSizes; ++i) {
3262  unsigned index = outputVectorRank - numFixedVectorSizes - i;
3263  if (fixedVectorSizes[i] != outputVectorType.getShape()[index])
3264  return emitError("fixed vector size must match output vector for dim ")
3265  << i;
3266  }
3267 
3268  // If all shape operands are produced by constant ops, verify that product
3269  // of dimensions for input/output shape match.
3270  auto isDefByConstant = [](Value operand) {
3271  return getConstantIntValue(operand).has_value();
3272  };
3273  if (llvm::all_of(getInputShape(), isDefByConstant) &&
3274  llvm::all_of(getOutputShape(), isDefByConstant)) {
3275  int64_t numInputElements = 1;
3276  for (auto operand : getInputShape())
3277  numInputElements *= getConstantIntValue(operand).value();
3278  int64_t numOutputElements = 1;
3279  for (auto operand : getOutputShape())
3280  numOutputElements *= getConstantIntValue(operand).value();
3281  if (numInputElements != numOutputElements)
3282  return emitError("product of input and output shape sizes must match");
3283  }
3284  return success();
3285 }
3286 
3287 void ReshapeOp::getFixedVectorSizes(SmallVectorImpl<int64_t> &results) {
3288  populateFromInt64AttrArray(getFixedVectorSizes(), results);
3289 }
3290 
3291 //===----------------------------------------------------------------------===//
3292 // ExtractStridedSliceOp
3293 //===----------------------------------------------------------------------===//
3294 
3295 // Inference works as follows:
3296 // 1. Add 'sizes' from prefix of dims in 'offsets'.
3297 // 2. Add sizes from 'vectorType' for remaining dims.
3298 // Scalable flags are inherited from 'vectorType'.
3299 static Type inferStridedSliceOpResultType(VectorType vectorType,
3300  ArrayAttr offsets, ArrayAttr sizes,
3301  ArrayAttr strides) {
3302  assert(offsets.size() == sizes.size() && offsets.size() == strides.size());
3304  shape.reserve(vectorType.getRank());
3305  unsigned idx = 0;
3306  for (unsigned e = offsets.size(); idx < e; ++idx)
3307  shape.push_back(llvm::cast<IntegerAttr>(sizes[idx]).getInt());
3308  for (unsigned e = vectorType.getShape().size(); idx < e; ++idx)
3309  shape.push_back(vectorType.getShape()[idx]);
3310 
3311  return VectorType::get(shape, vectorType.getElementType(),
3312  vectorType.getScalableDims());
3313 }
3314 
3315 void ExtractStridedSliceOp::build(OpBuilder &builder, OperationState &result,
3316  Value source, ArrayRef<int64_t> offsets,
3317  ArrayRef<int64_t> sizes,
3318  ArrayRef<int64_t> strides) {
3319  result.addOperands(source);
3320  auto offsetsAttr = getVectorSubscriptAttr(builder, offsets);
3321  auto sizesAttr = getVectorSubscriptAttr(builder, sizes);
3322  auto stridesAttr = getVectorSubscriptAttr(builder, strides);
3323  result.addTypes(
3324  inferStridedSliceOpResultType(llvm::cast<VectorType>(source.getType()),
3325  offsetsAttr, sizesAttr, stridesAttr));
3326  result.addAttribute(ExtractStridedSliceOp::getOffsetsAttrName(result.name),
3327  offsetsAttr);
3328  result.addAttribute(ExtractStridedSliceOp::getSizesAttrName(result.name),
3329  sizesAttr);
3330  result.addAttribute(ExtractStridedSliceOp::getStridesAttrName(result.name),
3331  stridesAttr);
3332 }
3333 
3335  auto type = getSourceVectorType();
3336  auto offsets = getOffsetsAttr();
3337  auto sizes = getSizesAttr();
3338  auto strides = getStridesAttr();
3339  if (offsets.size() != sizes.size() || offsets.size() != strides.size())
3340  return emitOpError(
3341  "expected offsets, sizes and strides attributes of same size");
3342 
3343  auto shape = type.getShape();
3344  auto offName = getOffsetsAttrName();
3345  auto sizesName = getSizesAttrName();
3346  auto stridesName = getStridesAttrName();
3347  if (failed(
3348  isIntegerArrayAttrSmallerThanShape(*this, offsets, shape, offName)) ||
3349  failed(
3350  isIntegerArrayAttrSmallerThanShape(*this, sizes, shape, sizesName)) ||
3351  failed(isIntegerArrayAttrSmallerThanShape(*this, strides, shape,
3352  stridesName)) ||
3353  failed(
3354  isIntegerArrayAttrConfinedToShape(*this, offsets, shape, offName)) ||
3355  failed(isIntegerArrayAttrConfinedToShape(*this, sizes, shape, sizesName,
3356  /*halfOpen=*/false,
3357  /*min=*/1)) ||
3358  failed(isIntegerArrayAttrConfinedToRange(*this, strides, /*min=*/1,
3359  /*max=*/1, stridesName,
3360  /*halfOpen=*/false)) ||
3361  failed(isSumOfIntegerArrayAttrConfinedToShape(*this, offsets, sizes,
3362  shape, offName, sizesName,
3363  /*halfOpen=*/false)))
3364  return failure();
3365 
3366  auto resultType = inferStridedSliceOpResultType(getSourceVectorType(),
3367  offsets, sizes, strides);
3368  if (getResult().getType() != resultType)
3369  return emitOpError("expected result type to be ") << resultType;
3370 
3371  for (unsigned idx = 0; idx < sizes.size(); ++idx) {
3372  if (type.getScalableDims()[idx]) {
3373  auto inputDim = type.getShape()[idx];
3374  auto inputSize = llvm::cast<IntegerAttr>(sizes[idx]).getInt();
3375  if (inputDim != inputSize)
3376  return emitOpError("expected size at idx=")
3377  << idx
3378  << (" to match the corresponding base size from the input "
3379  "vector (")
3380  << inputSize << (" vs ") << inputDim << (")");
3381  }
3382  }
3383 
3384  return success();
3385 }
3386 
3387 // When the source of ExtractStrided comes from a chain of InsertStrided ops try
3388 // to use the source of the InsertStrided ops if we can detect that the
3389 // extracted vector is a subset of one of the vector inserted.
3390 static LogicalResult
3391 foldExtractStridedOpFromInsertChain(ExtractStridedSliceOp op) {
3392  // Helper to extract integer out of ArrayAttr.
3393  auto getElement = [](ArrayAttr array, int idx) {
3394  return llvm::cast<IntegerAttr>(array[idx]).getInt();
3395  };
3396  ArrayAttr extractOffsets = op.getOffsets();
3397  ArrayAttr extractStrides = op.getStrides();
3398  ArrayAttr extractSizes = op.getSizes();
3399  auto insertOp = op.getVector().getDefiningOp<InsertStridedSliceOp>();
3400  while (insertOp) {
3401  if (op.getSourceVectorType().getRank() !=
3402  insertOp.getSourceVectorType().getRank())
3403  return failure();
3404  ArrayAttr insertOffsets = insertOp.getOffsets();
3405  ArrayAttr insertStrides = insertOp.getStrides();
3406  // If the rank of extract is greater than the rank of insert, we are likely
3407  // extracting a partial chunk of the vector inserted.
3408  if (extractOffsets.size() > insertOffsets.size())
3409  return failure();
3410  bool patialoverlap = false;
3411  bool disjoint = false;
3412  SmallVector<int64_t, 4> offsetDiffs;
3413  for (unsigned dim = 0, e = extractOffsets.size(); dim < e; ++dim) {
3414  if (getElement(extractStrides, dim) != getElement(insertStrides, dim))
3415  return failure();
3416  int64_t start = getElement(insertOffsets, dim);
3417  int64_t end = start + insertOp.getSourceVectorType().getDimSize(dim);
3418  int64_t offset = getElement(extractOffsets, dim);
3419  int64_t size = getElement(extractSizes, dim);
3420  // Check if the start of the extract offset is in the interval inserted.
3421  if (start <= offset && offset < end) {
3422  // If the extract interval overlaps but is not fully included we may
3423  // have a partial overlap that will prevent any folding.
3424  if (offset + size > end)
3425  patialoverlap = true;
3426  offsetDiffs.push_back(offset - start);
3427  continue;
3428  }
3429  disjoint = true;
3430  break;
3431  }
3432  // The extract element chunk is a subset of the insert element.
3433  if (!disjoint && !patialoverlap) {
3434  op.setOperand(insertOp.getSource());
3435  // OpBuilder is only used as a helper to build an I64ArrayAttr.
3436  OpBuilder b(op.getContext());
3437  op.setOffsetsAttr(b.getI64ArrayAttr(offsetDiffs));
3438  return success();
3439  }
3440  // If the chunk extracted is disjoint from the chunk inserted, keep looking
3441  // in the insert chain.
3442  if (disjoint)
3443  insertOp = insertOp.getDest().getDefiningOp<InsertStridedSliceOp>();
3444  else {
3445  // The extracted vector partially overlap the inserted vector, we cannot
3446  // fold.
3447  return failure();
3448  }
3449  }
3450  return failure();
3451 }
3452 
3453 OpFoldResult ExtractStridedSliceOp::fold(FoldAdaptor adaptor) {
3454  if (getSourceVectorType() == getResult().getType())
3455  return getVector();
3457  return getResult();
3458  return {};
3459 }
3460 
3461 void ExtractStridedSliceOp::getOffsets(SmallVectorImpl<int64_t> &results) {
3462  populateFromInt64AttrArray(getOffsets(), results);
3463 }
3464 
3465 namespace {
3466 
3467 // Pattern to rewrite an ExtractStridedSliceOp(ConstantMaskOp) to
3468 // ConstantMaskOp.
3469 class StridedSliceConstantMaskFolder final
3470  : public OpRewritePattern<ExtractStridedSliceOp> {
3471 public:
3473 
3474  LogicalResult matchAndRewrite(ExtractStridedSliceOp extractStridedSliceOp,
3475  PatternRewriter &rewriter) const override {
3476  // Return if 'extractStridedSliceOp' operand is not defined by a
3477  // ConstantMaskOp.
3478  auto *defOp = extractStridedSliceOp.getVector().getDefiningOp();
3479  auto constantMaskOp = dyn_cast_or_null<ConstantMaskOp>(defOp);
3480  if (!constantMaskOp)
3481  return failure();
3482  // Return if 'extractStridedSliceOp' has non-unit strides.
3483  if (extractStridedSliceOp.hasNonUnitStrides())
3484  return failure();
3485  // Gather constant mask dimension sizes.
3486  SmallVector<int64_t, 4> maskDimSizes;
3487  populateFromInt64AttrArray(constantMaskOp.getMaskDimSizes(), maskDimSizes);
3488  // Gather strided slice offsets and sizes.
3489  SmallVector<int64_t, 4> sliceOffsets;
3490  populateFromInt64AttrArray(extractStridedSliceOp.getOffsets(),
3491  sliceOffsets);
3492  SmallVector<int64_t, 4> sliceSizes;
3493  populateFromInt64AttrArray(extractStridedSliceOp.getSizes(), sliceSizes);
3494 
3495  // Compute slice of vector mask region.
3496  SmallVector<int64_t, 4> sliceMaskDimSizes;
3497  sliceMaskDimSizes.reserve(maskDimSizes.size());
3498  for (auto [maskDimSize, sliceOffset, sliceSize] :
3499  llvm::zip(maskDimSizes, sliceOffsets, sliceSizes)) {
3500  int64_t sliceMaskDimSize = std::max(
3501  static_cast<int64_t>(0),
3502  std::min(sliceOffset + sliceSize, maskDimSize) - sliceOffset);
3503  sliceMaskDimSizes.push_back(sliceMaskDimSize);
3504  }
3505  // Add unchanged dimensions.
3506  if (sliceMaskDimSizes.size() < maskDimSizes.size())
3507  for (size_t i = sliceMaskDimSizes.size(); i < maskDimSizes.size(); ++i)
3508  sliceMaskDimSizes.push_back(maskDimSizes[i]);
3509  // If any of 'sliceMaskDimSizes' are zero, then set all to zero (masked
3510  // region is a conjunction of mask dim intervals).
3511  if (llvm::is_contained(sliceMaskDimSizes, 0))
3512  sliceMaskDimSizes.assign(maskDimSizes.size(), 0);
3513 
3514  // Replace 'extractStridedSliceOp' with ConstantMaskOp with sliced mask
3515  // region.
3516  rewriter.replaceOpWithNewOp<ConstantMaskOp>(
3517  extractStridedSliceOp, extractStridedSliceOp.getResult().getType(),
3518  vector::getVectorSubscriptAttr(rewriter, sliceMaskDimSizes));
3519  return success();
3520  }
3521 };
3522 
3523 // Pattern to rewrite a ExtractStridedSliceOp(splat ConstantOp) -> ConstantOp.
3524 class StridedSliceSplatConstantFolder final
3525  : public OpRewritePattern<ExtractStridedSliceOp> {
3526 public:
3528 
3529  LogicalResult matchAndRewrite(ExtractStridedSliceOp extractStridedSliceOp,
3530  PatternRewriter &rewriter) const override {
3531  // Return if 'ExtractStridedSliceOp' operand is not defined by a splat
3532  // ConstantOp.
3533  Value sourceVector = extractStridedSliceOp.getVector();
3534  Attribute vectorCst;
3535  if (!matchPattern(sourceVector, m_Constant(&vectorCst)))
3536  return failure();
3537 
3538  auto splat = llvm::dyn_cast<SplatElementsAttr>(vectorCst);
3539  if (!splat)
3540  return failure();
3541 
3542  auto newAttr = SplatElementsAttr::get(extractStridedSliceOp.getType(),
3543  splat.getSplatValue<Attribute>());
3544  rewriter.replaceOpWithNewOp<arith::ConstantOp>(extractStridedSliceOp,
3545  newAttr);
3546  return success();
3547  }
3548 };
3549 
3550 // Pattern to rewrite a ExtractStridedSliceOp(non-splat ConstantOp) ->
3551 // ConstantOp.
3552 class StridedSliceNonSplatConstantFolder final
3553  : public OpRewritePattern<ExtractStridedSliceOp> {
3554 public:
3556 
3557  LogicalResult matchAndRewrite(ExtractStridedSliceOp extractStridedSliceOp,
3558  PatternRewriter &rewriter) const override {
3559  // Return if 'ExtractStridedSliceOp' operand is not defined by a non-splat
3560  // ConstantOp.
3561  Value sourceVector = extractStridedSliceOp.getVector();
3562  Attribute vectorCst;
3563  if (!matchPattern(sourceVector, m_Constant(&vectorCst)))
3564  return failure();
3565 
3566  // The splat case is handled by `StridedSliceSplatConstantFolder`.
3567  auto dense = llvm::dyn_cast<DenseElementsAttr>(vectorCst);
3568  if (!dense || dense.isSplat())
3569  return failure();
3570 
3571  // TODO: Handle non-unit strides when they become available.
3572  if (extractStridedSliceOp.hasNonUnitStrides())
3573  return failure();
3574 
3575  auto sourceVecTy = llvm::cast<VectorType>(sourceVector.getType());
3576  ArrayRef<int64_t> sourceShape = sourceVecTy.getShape();
3577  SmallVector<int64_t, 4> sourceStrides = computeStrides(sourceShape);
3578 
3579  VectorType sliceVecTy = extractStridedSliceOp.getType();
3580  ArrayRef<int64_t> sliceShape = sliceVecTy.getShape();
3581  int64_t sliceRank = sliceVecTy.getRank();
3582 
3583  // Expand offsets and sizes to match the vector rank.
3584  SmallVector<int64_t, 4> offsets(sliceRank, 0);
3585  copy(getI64SubArray(extractStridedSliceOp.getOffsets()), offsets.begin());
3586 
3587  SmallVector<int64_t, 4> sizes(sourceShape.begin(), sourceShape.end());
3588  copy(getI64SubArray(extractStridedSliceOp.getSizes()), sizes.begin());
3589 
3590  // Calculate the slice elements by enumerating all slice positions and
3591  // linearizing them. The enumeration order is lexicographic which yields a
3592  // sequence of monotonically increasing linearized position indices.
3593  auto denseValuesBegin = dense.value_begin<Attribute>();
3594  SmallVector<Attribute> sliceValues;
3595  sliceValues.reserve(sliceVecTy.getNumElements());
3596  SmallVector<int64_t> currSlicePosition(offsets.begin(), offsets.end());
3597  do {
3598  int64_t linearizedPosition = linearize(currSlicePosition, sourceStrides);
3599  assert(linearizedPosition < sourceVecTy.getNumElements() &&
3600  "Invalid index");
3601  sliceValues.push_back(*(denseValuesBegin + linearizedPosition));
3602  } while (
3603  succeeded(incSlicePosition(currSlicePosition, sliceShape, offsets)));
3604 
3605  assert(static_cast<int64_t>(sliceValues.size()) ==
3606  sliceVecTy.getNumElements() &&
3607  "Invalid number of slice elements");
3608  auto newAttr = DenseElementsAttr::get(sliceVecTy, sliceValues);
3609  rewriter.replaceOpWithNewOp<arith::ConstantOp>(extractStridedSliceOp,
3610  newAttr);
3611  return success();
3612  }
3613 };
3614 
3615 // Pattern to rewrite an ExtractStridedSliceOp(BroadcastOp) to
3616 // BroadcastOp(ExtractStrideSliceOp).
3617 class StridedSliceBroadcast final
3618  : public OpRewritePattern<ExtractStridedSliceOp> {
3619 public:
3621 
3622  LogicalResult matchAndRewrite(ExtractStridedSliceOp op,
3623  PatternRewriter &rewriter) const override {
3624  auto broadcast = op.getVector().getDefiningOp<BroadcastOp>();
3625  if (!broadcast)
3626  return failure();
3627  auto srcVecType =
3628  llvm::dyn_cast<VectorType>(broadcast.getSource().getType());
3629  unsigned srcRank = srcVecType ? srcVecType.getRank() : 0;
3630  auto dstVecType = llvm::cast<VectorType>(op.getType());
3631  unsigned dstRank = dstVecType.getRank();
3632  unsigned rankDiff = dstRank - srcRank;
3633  // Check if the most inner dimensions of the source of the broadcast are the
3634  // same as the destination of the extract. If this is the case we can just
3635  // use a broadcast as the original dimensions are untouched.
3636  bool lowerDimMatch = true;
3637  for (unsigned i = 0; i < srcRank; i++) {
3638  if (srcVecType.getDimSize(i) != dstVecType.getDimSize(i + rankDiff)) {
3639  lowerDimMatch = false;
3640  break;
3641  }
3642  }
3643  Value source = broadcast.getSource();
3644  // If the inner dimensions don't match, it means we need to extract from the
3645  // source of the orignal broadcast and then broadcast the extracted value.
3646  // We also need to handle degenerated cases where the source is effectively
3647  // just a single scalar.
3648  bool isScalarSrc = (srcRank == 0 || srcVecType.getNumElements() == 1);
3649  if (!lowerDimMatch && !isScalarSrc) {
3650  source = rewriter.create<ExtractStridedSliceOp>(
3651  op->getLoc(), source,
3652  getI64SubArray(op.getOffsets(), /* dropFront=*/rankDiff),
3653  getI64SubArray(op.getSizes(), /* dropFront=*/rankDiff),
3654  getI64SubArray(op.getStrides(), /* dropFront=*/rankDiff));
3655  }
3656  rewriter.replaceOpWithNewOp<BroadcastOp>(op, op.getType(), source);
3657  return success();
3658  }
3659 };
3660 
3661 /// Pattern to rewrite an ExtractStridedSliceOp(SplatOp) to SplatOp.
3662 class StridedSliceSplat final : public OpRewritePattern<ExtractStridedSliceOp> {
3663 public:
3665 
3666  LogicalResult matchAndRewrite(ExtractStridedSliceOp op,
3667  PatternRewriter &rewriter) const override {
3668  auto splat = op.getVector().getDefiningOp<SplatOp>();
3669  if (!splat)
3670  return failure();
3671  rewriter.replaceOpWithNewOp<SplatOp>(op, op.getType(), splat.getInput());
3672  return success();
3673  }
3674 };
3675 
3676 } // namespace
3677 
3678 void ExtractStridedSliceOp::getCanonicalizationPatterns(
3679  RewritePatternSet &results, MLIRContext *context) {
3680  // Pattern to rewrite a ExtractStridedSliceOp(ConstantMaskOp) ->
3681  // ConstantMaskOp and ExtractStridedSliceOp(ConstantOp) -> ConstantOp.
3682  results.add<StridedSliceConstantMaskFolder, StridedSliceSplatConstantFolder,
3683  StridedSliceNonSplatConstantFolder, StridedSliceBroadcast,
3684  StridedSliceSplat>(context);
3685 }
3686 
3687 //===----------------------------------------------------------------------===//
3688 // TransferReadOp
3689 //===----------------------------------------------------------------------===//
3690 
3691 /// 1. Builder that sets padding to zero and an empty mask (variant with attrs).
3692 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
3693  VectorType vectorType, Value source,
3694  ValueRange indices, AffineMapAttr permutationMapAttr,
3695  /*optional*/ ArrayAttr inBoundsAttr) {
3696  Type elemType = llvm::cast<ShapedType>(source.getType()).getElementType();
3697  Value padding = builder.create<arith::ConstantOp>(
3698  result.location, elemType, builder.getZeroAttr(elemType));
3699  build(builder, result, vectorType, source, indices, permutationMapAttr,
3700  padding, /*mask=*/Value(), inBoundsAttr);
3701 }
3702 
3703 /// 2. Builder that sets padding to zero an empty mask (variant without attrs).
3704 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
3705  VectorType vectorType, Value source,
3706  ValueRange indices, AffineMap permutationMap,
3707  std::optional<ArrayRef<bool>> inBounds) {
3708  auto permutationMapAttr = AffineMapAttr::get(permutationMap);
3709  auto inBoundsAttr = (inBounds && !inBounds.value().empty())
3710  ? builder.getBoolArrayAttr(inBounds.value())
3711  : ArrayAttr();
3712  build(builder, result, vectorType, source, indices, permutationMapAttr,
3713  inBoundsAttr);
3714 }
3715 
3716 /// 3. Builder that sets permutation map to 'getMinorIdentityMap'.
3717 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
3718  VectorType vectorType, Value source,
3719  ValueRange indices, Value padding,
3720  std::optional<ArrayRef<bool>> inBounds) {
3721  AffineMap permutationMap = getTransferMinorIdentityMap(
3722  llvm::cast<ShapedType>(source.getType()), vectorType);
3723  auto permutationMapAttr = AffineMapAttr::get(permutationMap);
3724  auto inBoundsAttr = (inBounds && !inBounds.value().empty())
3725  ? builder.getBoolArrayAttr(inBounds.value())
3726  : ArrayAttr();
3727  build(builder, result, vectorType, source, indices, permutationMapAttr,
3728  padding,
3729  /*mask=*/Value(), inBoundsAttr);
3730 }
3731 
3732 /// 4. Builder that sets padding to zero and permutation map to
3733 /// 'getMinorIdentityMap'.
3734 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
3735  VectorType vectorType, Value source,
3736  ValueRange indices,
3737  std::optional<ArrayRef<bool>> inBounds) {
3738  Type elemType = llvm::cast<ShapedType>(source.getType()).getElementType();
3739  Value padding = builder.create<arith::ConstantOp>(
3740  result.location, elemType, builder.getZeroAttr(elemType));
3741  build(builder, result, vectorType, source, indices, padding, inBounds);
3742 }
3743 
3744 template <typename EmitFun>
3746  EmitFun emitOpError) {
3747  SmallVector<bool, 8> seen(permutationMap.getNumInputs(), false);
3748  for (auto expr : permutationMap.getResults()) {
3749  auto dim = dyn_cast<AffineDimExpr>(expr);
3750  auto zero = dyn_cast<AffineConstantExpr>(expr);
3751  if (zero) {
3752  if (zero.getValue() != 0) {
3753  return emitOpError(
3754  "requires a projected permutation_map (at most one dim or the zero "
3755  "constant can appear in each result)");
3756  }
3757  continue;
3758  }
3759  if (!dim) {
3760  return emitOpError("requires a projected permutation_map (at most one "
3761  "dim or the zero constant can appear in each result)");
3762  }
3763  if (seen[dim.getPosition()]) {
3764  return emitOpError(
3765  "requires a permutation_map that is a permutation (found one dim "
3766  "used more than once)");
3767  }
3768  seen[dim.getPosition()] = true;
3769  }
3770  return success();
3771 }
3772 
3773 static LogicalResult
3774 verifyTransferOp(VectorTransferOpInterface op, ShapedType shapedType,
3775  VectorType vectorType, VectorType maskType,
3776  VectorType inferredMaskType, AffineMap permutationMap,
3777  ArrayAttr inBounds) {
3778  if (op->hasAttr("masked")) {
3779  return op->emitOpError("masked attribute has been removed. "
3780  "Use in_bounds instead.");
3781  }
3782 
3783  if (!llvm::isa<MemRefType, RankedTensorType>(shapedType))
3784  return op->emitOpError(
3785  "requires source to be a memref or ranked tensor type");
3786 
3787  auto elementType = shapedType.getElementType();
3788  DataLayout dataLayout = DataLayout::closest(op);
3789  if (auto vectorElementType = llvm::dyn_cast<VectorType>(elementType)) {
3790  // Memref or tensor has vector element type.
3791  unsigned sourceVecSize =
3792  dataLayout.getTypeSizeInBits(vectorElementType.getElementType()) *
3793  vectorElementType.getShape().back();
3794  unsigned resultVecSize =
3795  dataLayout.getTypeSizeInBits(vectorType.getElementType()) *
3796  vectorType.getShape().back();
3797  if (resultVecSize % sourceVecSize != 0)
3798  return op->emitOpError(
3799  "requires the bitwidth of the minor 1-D vector to be an integral "
3800  "multiple of the bitwidth of the minor 1-D vector of the source");
3801 
3802  unsigned sourceVecEltRank = vectorElementType.getRank();
3803  unsigned resultVecRank = vectorType.getRank();
3804  if (sourceVecEltRank > resultVecRank)
3805  return op->emitOpError(
3806  "requires source vector element and vector result ranks to match.");
3807  unsigned rankOffset = resultVecRank - sourceVecEltRank;
3808  // Check that permutation map results match 'rankOffset' of vector type.
3809  if (permutationMap.getNumResults() != rankOffset)
3810  return op->emitOpError("requires a permutation_map with result dims of "
3811  "the same rank as the vector type");
3812 
3813  if (maskType)
3814  return op->emitOpError("does not support masks with vector element type");
3815  } else {
3816  // Memref or tensor has scalar element type.
3817  unsigned minorSize =
3818  vectorType.getRank() == 0 ? 1 : vectorType.getShape().back();
3819  unsigned resultVecSize =
3820  dataLayout.getTypeSizeInBits(vectorType.getElementType()) * minorSize;
3821  if (resultVecSize % dataLayout.getTypeSizeInBits(elementType) != 0)
3822  return op->emitOpError(
3823  "requires the bitwidth of the minor 1-D vector to be an integral "
3824  "multiple of the bitwidth of the source element type");
3825 
3826  // Check that permutation map results match rank of vector type.
3827  if (permutationMap.getNumResults() != vectorType.getRank())
3828  return op->emitOpError("requires a permutation_map with result dims of "
3829  "the same rank as the vector type");
3830  }
3831 
3832  if (permutationMap.getNumSymbols() != 0)
3833  return op->emitOpError("requires permutation_map without symbols");
3834 
3835  if (permutationMap.getNumInputs() != shapedType.getRank())
3836  return op->emitOpError("requires a permutation_map with input dims of the "
3837  "same rank as the source type");
3838 
3839  if (maskType && maskType != inferredMaskType)
3840  return op->emitOpError("inferred mask type (")
3841  << inferredMaskType << ") and mask operand type (" << maskType
3842  << ") don't match";
3843 
3844  if (inBounds) {
3845  if (permutationMap.getNumResults() != static_cast<int64_t>(inBounds.size()))
3846  return op->emitOpError("expects the optional in_bounds attr of same rank "
3847  "as permutation_map results: ")
3848  << AffineMapAttr::get(permutationMap)
3849  << " vs inBounds of size: " << inBounds.size();
3850  for (unsigned int i = 0; i < permutationMap.getNumResults(); ++i)
3851  if (isa<AffineConstantExpr>(permutationMap.getResult(i)) &&
3852  !llvm::cast<BoolAttr>(inBounds.getValue()[i]).getValue())
3853  return op->emitOpError("requires broadcast dimensions to be in-bounds");
3854  }
3855 
3856  return success();
3857 }
3858 
3859 static void printTransferAttrs(OpAsmPrinter &p, VectorTransferOpInterface op) {
3860  SmallVector<StringRef, 3> elidedAttrs;
3861  elidedAttrs.push_back(TransferReadOp::getOperandSegmentSizeAttr());
3862  if (op.getPermutationMap().isMinorIdentity())
3863  elidedAttrs.push_back(op.getPermutationMapAttrName());
3864  // Elide in_bounds attribute if all dims are out-of-bounds.
3865  if (llvm::none_of(op.getInBoundsValues(), [](bool b) { return b; }))
3866  elidedAttrs.push_back(op.getInBoundsAttrName());
3867  p.printOptionalAttrDict(op->getAttrs(), elidedAttrs);
3868 }
3869 
3871  p << " " << getSource() << "[" << getIndices() << "], " << getPadding();
3872  if (getMask())
3873  p << ", " << getMask();
3874  printTransferAttrs(p, *this);
3875  p << " : " << getShapedType() << ", " << getVectorType();
3876 }
3877 
3878 VectorType mlir::vector::inferTransferOpMaskType(VectorType vecType,
3879  AffineMap permMap) {
3880  auto i1Type = IntegerType::get(permMap.getContext(), 1);
3881  AffineMap invPermMap = inversePermutation(compressUnusedDims(permMap));
3882  assert(invPermMap && "Inversed permutation map couldn't be computed");
3883  SmallVector<int64_t, 8> maskShape = invPermMap.compose(vecType.getShape());
3884 
3885  SmallVector<bool> scalableDims =
3886  applyPermutationMap(invPermMap, vecType.getScalableDims());
3887 
3888  return VectorType::get(maskShape, i1Type, scalableDims);
3889 }
3890 
3892  auto &builder = parser.getBuilder();
3893  SMLoc typesLoc;
3894  OpAsmParser::UnresolvedOperand sourceInfo;
3896  OpAsmParser::UnresolvedOperand paddingInfo;
3897  SmallVector<Type, 2> types;
3899  // Parsing with support for paddingValue.
3900  if (parser.parseOperand(sourceInfo) ||
3901  parser.parseOperandList(indexInfo, OpAsmParser::Delimiter::Square) ||
3902  parser.parseComma() || parser.parseOperand(paddingInfo))
3903  return failure();
3904  ParseResult hasMask = parser.parseOptionalComma();
3905  if (hasMask.succeeded()) {
3906  if (parser.parseOperand(maskInfo))
3907  return failure();
3908  }
3909  if (parser.parseOptionalAttrDict(result.attributes) ||
3910  parser.getCurrentLocation(&typesLoc) || parser.parseColonTypeList(types))
3911  return failure();
3912  if (types.size() != 2)
3913  return parser.emitError(typesLoc, "requires two types");
3914  auto indexType = builder.getIndexType();
3915  auto shapedType = llvm::dyn_cast<ShapedType>(types[0]);
3916  if (!shapedType || !llvm::isa<MemRefType, RankedTensorType>(shapedType))
3917  return parser.emitError(typesLoc, "requires memref or ranked tensor type");
3918  VectorType vectorType = llvm::dyn_cast<VectorType>(types[1]);
3919  if (!vectorType)
3920  return parser.emitError(typesLoc, "requires vector type");
3921  auto permMapAttrName = TransferReadOp::getPermutationMapAttrName(result.name);
3922  Attribute permMapAttr = result.attributes.get(permMapAttrName);
3923  AffineMap permMap;
3924  if (!permMapAttr) {
3925  permMap = getTransferMinorIdentityMap(shapedType, vectorType);
3926  result.attributes.set(permMapAttrName, AffineMapAttr::get(permMap));
3927  } else {
3928  permMap = llvm::cast<AffineMapAttr>(permMapAttr).getValue();
3929  }
3930  if (parser.resolveOperand(sourceInfo, shapedType, result.operands) ||
3931  parser.resolveOperands(indexInfo, indexType, result.operands) ||
3932  parser.resolveOperand(paddingInfo, shapedType.getElementType(),
3933  result.operands))
3934  return failure();
3935  if (hasMask.succeeded()) {
3936  if (llvm::dyn_cast<VectorType>(shapedType.getElementType()))
3937  return parser.emitError(
3938  maskInfo.location, "does not support masks with vector element type");
3939  if (vectorType.getRank() != permMap.getNumResults()) {
3940  return parser.emitError(typesLoc,
3941  "expected the same rank for the vector and the "
3942  "results of the permutation map");
3943  }
3944  // Instead of adding the mask type as an op type, compute it based on the
3945  // vector type and the permutation map (to keep the type signature small).
3946  auto maskType = inferTransferOpMaskType(vectorType, permMap);
3947  if (parser.resolveOperand(maskInfo, maskType, result.operands))
3948  return failure();
3949  }
3950  result.addAttribute(TransferReadOp::getOperandSegmentSizeAttr(),
3951  builder.getDenseI32ArrayAttr(
3952  {1, static_cast<int32_t>(indexInfo.size()), 1,
3953  static_cast<int32_t>(hasMask.succeeded())}));
3954  return parser.addTypeToList(vectorType, result.types);
3955 }
3956 
3958  // Consistency of elemental types in source and vector.
3959  ShapedType shapedType = getShapedType();
3960  VectorType vectorType = getVectorType();
3961  VectorType maskType = getMaskType();
3962  auto paddingType = getPadding().getType();
3963  auto permutationMap = getPermutationMap();
3964  VectorType inferredMaskType =
3965  maskType ? inferTransferOpMaskType(vectorType, permutationMap)
3966  : VectorType();
3967  auto sourceElementType = shapedType.getElementType();
3968 
3969  if (static_cast<int64_t>(getIndices().size()) != shapedType.getRank())
3970  return emitOpError("requires ") << shapedType.getRank() << " indices";
3971 
3972  if (failed(verifyTransferOp(cast<VectorTransferOpInterface>(getOperation()),
3973  shapedType, vectorType, maskType,
3974  inferredMaskType, permutationMap,
3975  getInBounds() ? *getInBounds() : ArrayAttr())))
3976  return failure();
3977 
3978  if (auto sourceVectorElementType =
3979  llvm::dyn_cast<VectorType>(sourceElementType)) {
3980  // Source has vector element type.
3981  // Check that 'sourceVectorElementType' and 'paddingType' types match.
3982  if (sourceVectorElementType != paddingType)
3983  return emitOpError(
3984  "requires source element type and padding type to match.");
3985 
3986  } else {
3987  // Check that 'paddingType' is valid to store in a vector type.
3988  if (!VectorType::isValidElementType(paddingType))
3989  return emitOpError("requires valid padding vector elemental type");
3990 
3991  // Check that padding type and vector element types match.
3992  if (paddingType != sourceElementType)
3993  return emitOpError(
3994  "requires formal padding and source of the same elemental type");
3995  }
3996 
3997  return verifyPermutationMap(permutationMap,
3998  [&](Twine t) { return emitOpError(t); });
3999 }
4000 
4001 // MaskableOpInterface methods.
4002 
4003 /// Returns the mask type expected by this operation. Mostly used for
4004 /// verification purposes. It requires the operation to be vectorized."
4005 Type TransferReadOp::getExpectedMaskType() {
4006  return inferTransferOpMaskType(getVectorType(), getPermutationMap());
4007 }
4008 
4009 template <typename TransferOp>
4010 static bool isInBounds(TransferOp op, int64_t resultIdx, int64_t indicesIdx) {
4011  // TODO: support more aggressive createOrFold on:
4012  // op.getIndices()[indicesIdx] + vectorType < dim(op.getSource(), indicesIdx)
4013  if (op.getShapedType().isDynamicDim(indicesIdx))
4014  return false;
4015  Value index = op.getIndices()[indicesIdx];
4016  std::optional<int64_t> cstOp = getConstantIntValue(index);
4017  if (!cstOp.has_value())
4018  return false;
4019 
4020  int64_t sourceSize = op.getShapedType().getDimSize(indicesIdx);
4021  int64_t vectorSize = op.getVectorType().getDimSize(resultIdx);
4022 
4023  return cstOp.value() + vectorSize <= sourceSize;
4024 }
4025 
4026 template <typename TransferOp>
4028  // TODO: support 0-d corner case.
4029  // TODO: Be less conservative.
4030  if (op.getTransferRank() == 0)
4031  return failure();
4032  AffineMap permutationMap = op.getPermutationMap();
4033  bool changed = false;
4034  SmallVector<bool, 4> newInBounds;
4035  newInBounds.reserve(op.getTransferRank());
4036  for (unsigned i = 0; i < op.getTransferRank(); ++i) {
4037  // Already marked as in-bounds, nothing to see here.
4038  if (op.isDimInBounds(i)) {
4039  newInBounds.push_back(true);
4040  continue;
4041  }
4042  // Currently out-of-bounds, check whether we can statically determine it is
4043  // inBounds.
4044  auto dimExpr = dyn_cast<AffineDimExpr>(permutationMap.getResult(i));
4045  assert(dimExpr && "Broadcast dims must be in-bounds");
4046  auto inBounds =
4047  isInBounds(op, /*resultIdx=*/i, /*indicesIdx=*/dimExpr.getPosition());
4048  newInBounds.push_back(inBounds);
4049  // We commit the pattern if it is "more inbounds".
4050  changed |= inBounds;
4051  }
4052  if (!changed)
4053  return failure();
4054  // OpBuilder is only used as a helper to build an I64ArrayAttr.
4055  OpBuilder b(op.getContext());
4056  op.setInBoundsAttr(b.getBoolArrayAttr(newInBounds));
4057  return success();
4058 }
4059 
4060 template <typename TransferOp>
4061 static LogicalResult foldTransferFullMask(TransferOp op) {
4062  auto mask = op.getMask();
4063  if (!mask)
4064  return failure();
4065 
4066  auto constantMask = mask.template getDefiningOp<vector::ConstantMaskOp>();
4067  if (!constantMask)
4068  return failure();
4069 
4070  if (!constantMask.isAllOnesMask())
4071  return failure();
4072 
4073  op.getMaskMutable().clear();
4074  return success();
4075 }
4076 
4077 /// ```
4078 /// %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
4079 /// : vector<1x4xf32>, tensor<4x4xf32>
4080 /// %0 = vector.transfer_read %w0[%c1, %c0], %cf0 {in_bounds = [true, true]}
4081 /// : tensor<4x4xf32>, vector<1x4xf32>
4082 /// ```
4083 /// -> Folds into
4084 /// ```
4085 /// %v0
4086 /// ```
4087 static Value foldRAW(TransferReadOp readOp) {
4088  if (!llvm::isa<RankedTensorType>(readOp.getShapedType()))
4089  return {};
4090  auto defWrite = readOp.getSource().getDefiningOp<vector::TransferWriteOp>();
4091  while (defWrite) {
4092  if (checkSameValueRAW(defWrite, readOp))
4093  return defWrite.getVector();
4095  cast<VectorTransferOpInterface>(defWrite.getOperation()),
4096  cast<VectorTransferOpInterface>(readOp.getOperation())))
4097  break;
4098  defWrite = defWrite.getSource().getDefiningOp<vector::TransferWriteOp>();
4099  }
4100  return {};
4101 }
4102 
4103 OpFoldResult TransferReadOp::fold(FoldAdaptor) {
4104  if (Value vec = foldRAW(*this))
4105  return vec;
4106  /// transfer_read(memrefcast) -> transfer_read
4108  return getResult();
4109  if (succeeded(foldTransferFullMask(*this)))
4110  return getResult();
4111  if (succeeded(memref::foldMemRefCast(*this)))
4112  return getResult();
4113  if (succeeded(tensor::foldTensorCast(*this)))
4114  return getResult();
4115  return OpFoldResult();
4116 }
4117 
4118 std::optional<SmallVector<int64_t, 4>> TransferReadOp::getShapeForUnroll() {
4119  return llvm::to_vector<4>(getVectorType().getShape());
4120 }
4121 
4122 void TransferReadOp::getEffects(
4124  &effects) {
4125  if (llvm::isa<MemRefType>(getShapedType()))
4126  effects.emplace_back(MemoryEffects::Read::get(), getSource(),
4128 }
4129 
4130 namespace {
4131 /// Store to load forwarding for transfer operations with permuation maps.
4132 /// Even if the permutation maps are different we can still propagate the store
4133 /// into the load if the size of the dimensions read and written match. Then we
4134 /// can replace the transfer_read + transfer_write by vector.broadcast and
4135 /// vector.transpose.
4136 /// Example:
4137 /// ```
4138 /// %w0 = vector.transfer_write %v0, %arg0[%c0, %c0, %c0]
4139 /// {in_bounds = [true, true],
4140 /// permutation_map = affine_map<(d0, d1, d2) -> (d2, d1)>} :
4141 /// vector<4x1xf32>, tensor<4x4x4xf32>
4142 /// %r = vector.transfer_read %w0[%c0, %c0, %c0], %cf0
4143 /// {in_bounds = [true, true, true, true],
4144 /// permutation_map = affine_map<(d0, d1, d2) -> (d1, 0, d2, 0)>} :
4145 /// tensor<4x4x4xf32>, vector<1x100x4x5xf32>
4146 /// ```
4147 /// To:
4148 /// ```
4149 /// %0 = vector.broadcast %arg1 : vector<4x1xf32> to vector<100x5x4x1xf32>
4150 /// %r = vector.transpose %0, [3, 0, 2, 1] :
4151 /// vector<100x5x4x1xf32> to vector<1x100x4x5xf32>
4152 /// ```
4153 struct TransferReadAfterWriteToBroadcast
4154  : public OpRewritePattern<TransferReadOp> {
4156 
4157  LogicalResult matchAndRewrite(TransferReadOp readOp,
4158  PatternRewriter &rewriter) const override {
4159  if (readOp.hasOutOfBoundsDim() ||
4160  !llvm::isa<RankedTensorType>(readOp.getShapedType()))
4161  return failure();
4162  auto defWrite = readOp.getSource().getDefiningOp<vector::TransferWriteOp>();
4163  if (!defWrite)
4164  return failure();
4165  // TODO: If the written transfer chunk is a superset of the read transfer
4166  // chunk we could do an extract_strided_slice.
4167  if (readOp.getTransferChunkAccessed() !=
4168  defWrite.getTransferChunkAccessed())
4169  return failure();
4170  // TODO: Support cases where a dim is explicitly written but implicitly
4171  // read (i.e., a unit dim that is rank reduced).
4172  if (getUnusedDimsBitVector({readOp.getPermutationMap()}) !=
4173  getUnusedDimsBitVector({defWrite.getPermutationMap()}))
4174  return failure();
4175  if (readOp.getIndices() != defWrite.getIndices() ||
4176  readOp.getMask() != defWrite.getMask())
4177  return failure();
4178  Value vec = defWrite.getVector();
4179  // TODO: loop through the chain of transfer_write if we can prove that they
4180  // don't overlap with the transfer_read. This requires improving
4181  // `isDisjointTransferIndices` helper.
4182  AffineMap readMap = compressUnusedDims(readOp.getPermutationMap());
4183  AffineMap writeMap = compressUnusedDims(defWrite.getPermutationMap());
4184  AffineMap map = readMap.compose(writeMap);
4185  if (map.getNumResults() == 0)
4186  return failure();
4187  // Calculate the permutation to apply to go from the vector stored to the
4188  // vector read.
4189  SmallVector<unsigned> permutation;
4190  if (!map.isPermutationOfMinorIdentityWithBroadcasting(permutation))
4191  return failure();
4192 
4193  Location loc = readOp.getLoc();
4194  // Calculate the broadcast shape by applying the reverse permutation to the
4195  // final shape we want.
4196  ArrayRef<int64_t> destShape = readOp.getVectorType().getShape();
4197  SmallVector<int64_t> broadcastShape(destShape.size());
4198  SmallVector<bool> broadcastScalableFlags(destShape.size());
4199  for (const auto &pos : llvm::enumerate(permutation)) {
4200  broadcastShape[pos.value()] = destShape[pos.index()];
4201  broadcastScalableFlags[pos.value()] =
4202  readOp.getVectorType().getScalableDims()[pos.index()];
4203  }
4204  VectorType broadcastedType = VectorType::get(
4205  broadcastShape, defWrite.getVectorType().getElementType(),
4206  broadcastScalableFlags);
4207  vec = rewriter.create<vector::BroadcastOp>(loc, broadcastedType, vec);
4208  SmallVector<int64_t> transposePerm(permutation.begin(), permutation.end());
4209  rewriter.replaceOpWithNewOp<vector::TransposeOp>(readOp, vec,
4210  transposePerm);
4211  return success();
4212  }
4213 };
4214 } // namespace
4215 
4216 void TransferReadOp::getCanonicalizationPatterns(RewritePatternSet &results,
4217  MLIRContext *context) {
4218  results.add<TransferReadAfterWriteToBroadcast>(context);
4219 }
4220 
4221 //===----------------------------------------------------------------------===//
4222 // TransferWriteOp
4223 //===----------------------------------------------------------------------===//
4224 
4225 /// 1. Builder with type inference.
4226 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
4227  Value vector, Value dest, ValueRange indices,
4228  AffineMapAttr permutationMapAttr,
4229  /*optional*/ Value mask,
4230  /*optional*/ ArrayAttr inBoundsAttr) {
4231  Type resultType = llvm::dyn_cast<RankedTensorType>(dest.getType());
4232  build(builder, result, resultType, vector, dest, indices, permutationMapAttr,
4233  mask, inBoundsAttr);
4234 }
4235 
4236 /// 2. Builder with type inference that sets an empty mask (variant with attrs).
4237 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
4238  Value vector, Value dest, ValueRange indices,
4239  AffineMapAttr permutationMapAttr,
4240  /*optional*/ ArrayAttr inBoundsAttr) {
4241  build(builder, result, vector, dest, indices, permutationMapAttr,
4242  /*mask=*/Value(), inBoundsAttr);
4243 }
4244 
4245 /// 3. Builder with type inference that sets an empty mask (variant without
4246 /// attrs)
4247 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
4248  Value vector, Value dest, ValueRange indices,
4249  AffineMap permutationMap,
4250  std::optional<ArrayRef<bool>> inBounds) {
4251  auto permutationMapAttr = AffineMapAttr::get(permutationMap);
4252  auto inBoundsAttr = (inBounds && !inBounds.value().empty())
4253  ? builder.getBoolArrayAttr(inBounds.value())
4254  : ArrayAttr();
4255  build(builder, result, vector, dest, indices, permutationMapAttr,
4256  /*mask=*/Value(), inBoundsAttr);
4257 }
4258 
4259 /// 4. Builder with type inference that sets an empty mask and sets permutation
4260 /// map to 'getMinorIdentityMap'.
4261 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
4262  Value vector, Value dest, ValueRange indices,
4263  std::optional<ArrayRef<bool>> inBounds) {
4264  auto vectorType = llvm::cast<VectorType>(vector.getType());
4265  AffineMap permutationMap = getTransferMinorIdentityMap(
4266  llvm::cast<ShapedType>(dest.getType()), vectorType);
4267  build(builder, result, vector, dest, indices, permutationMap, inBounds);
4268 }
4269 
4271  OperationState &result) {
4272  auto &builder = parser.getBuilder();
4273  SMLoc typesLoc;
4274  OpAsmParser::UnresolvedOperand vectorInfo, sourceInfo;
4276  SmallVector<Type, 2> types;
4278  if (parser.parseOperand(vectorInfo) || parser.parseComma() ||
4279  parser.parseOperand(sourceInfo) ||
4280  parser.parseOperandList(indexInfo, OpAsmParser::Delimiter::Square))
4281  return failure();
4282  ParseResult hasMask = parser.parseOptionalComma();
4283  if (hasMask.succeeded() && parser.parseOperand(maskInfo))
4284  return failure();
4285  if (parser.parseOptionalAttrDict(result.attributes) ||
4286  parser.getCurrentLocation(&typesLoc) || parser.parseColonTypeList(types))
4287  return failure();
4288  if (types.size() != 2)
4289  return parser.emitError(typesLoc, "requires two types");
4290  auto indexType = builder.getIndexType();
4291  VectorType vectorType = llvm::dyn_cast<VectorType>(types[0]);
4292  if (!vectorType)
4293  return parser.emitError(typesLoc, "requires vector type");
4294  ShapedType shapedType = llvm::dyn_cast<ShapedType>(types[1]);
4295  if (!shapedType || !llvm::isa<MemRefType, RankedTensorType>(shapedType))
4296  return parser.emitError(typesLoc, "requires memref or ranked tensor type");
4297  auto permMapAttrName =
4298  TransferWriteOp::getPermutationMapAttrName(result.name);
4299  auto permMapAttr = result.attributes.get(permMapAttrName);
4300  AffineMap permMap;
4301  if (!permMapAttr) {
4302  permMap = getTransferMinorIdentityMap(shapedType, vectorType);
4303  result.attributes.set(permMapAttrName, AffineMapAttr::get(permMap));
4304  } else {
4305  permMap = llvm::cast<AffineMapAttr>(permMapAttr).getValue();
4306  }
4307  if (parser.resolveOperand(vectorInfo, vectorType, result.operands) ||
4308  parser.resolveOperand(sourceInfo, shapedType, result.operands) ||
4309  parser.resolveOperands(indexInfo, indexType, result.operands))
4310  return failure();
4311  if (hasMask.succeeded()) {
4312  if (llvm::dyn_cast<VectorType>(shapedType.getElementType()))
4313  return parser.emitError(
4314  maskInfo.location, "does not support masks with vector element type");
4315  if (vectorType.getRank() != permMap.getNumResults()) {
4316  return parser.emitError(typesLoc,
4317  "expected the same rank for the vector and the "
4318  "results of the permutation map");
4319  }
4320  auto maskType = inferTransferOpMaskType(vectorType, permMap);
4321  if (parser.resolveOperand(maskInfo, maskType, result.operands))
4322  return failure();
4323  }
4324  result.addAttribute(TransferWriteOp::getOperandSegmentSizeAttr(),
4325  builder.getDenseI32ArrayAttr(
4326  {1, 1, static_cast<int32_t>(indexInfo.size()),
4327  static_cast<int32_t>(hasMask.succeeded())}));
4328  return failure(llvm::isa<RankedTensorType>(shapedType) &&
4329  parser.addTypeToList(shapedType, result.types));
4330 }
4331 
4333  p << " " << getVector() << ", " << getSource() << "[" << getIndices() << "]";
4334  if (getMask())
4335  p << ", " << getMask();
4336  printTransferAttrs(p, *this);
4337  p << " : " << getVectorType() << ", " << getShapedType();
4338 }
4339 
4341  // Consistency of elemental types in shape and vector.
4342  ShapedType shapedType = getShapedType();
4343  VectorType vectorType = getVectorType();
4344  VectorType maskType = getMaskType();
4345  auto permutationMap = getPermutationMap();
4346  VectorType inferredMaskType =
4347  maskType ? inferTransferOpMaskType(vectorType, permutationMap)
4348  : VectorType();
4349 
4350  if (llvm::size(getIndices()) != shapedType.getRank())
4351  return emitOpError("requires ") << shapedType.getRank() << " indices";
4352 
4353  // We do not allow broadcast dimensions on TransferWriteOps for the moment,
4354  // as the semantics is unclear. This can be revisited later if necessary.
4355  if (hasBroadcastDim())
4356  return emitOpError("should not have broadcast dimensions");
4357 
4358  if (failed(verifyTransferOp(cast<VectorTransferOpInterface>(getOperation()),
4359  shapedType, vectorType, maskType,
4360  inferredMaskType, permutationMap,
4361  getInBounds() ? *getInBounds() : ArrayAttr())))
4362  return failure();
4363 
4364  return verifyPermutationMap(permutationMap,
4365  [&](Twine t) { return emitOpError(t); });
4366 }
4367 
4368 // MaskableOpInterface methods.
4369 
4370 /// Returns the mask type expected by this operation. Mostly used for
4371 /// verification purposes.
4372 Type TransferWriteOp::getExpectedMaskType() {
4373  return inferTransferOpMaskType(getVectorType(), getPermutationMap());
4374 }
4375 
4376 /// Fold:
4377 /// ```
4378 /// %t1 = ...
4379 /// %v = vector.transfer_read %t0[%c0...], {in_bounds = [true...]} :
4380 /// tensor<static_sizesxf32>, vector<static_sizesxf32>
4381 /// %t2 = vector.transfer_write %v, %t1[%c0...] {in_bounds = [true...]} :
4382 /// vector<static_sizesxf32>, tensor<static_sizesxf32>
4383 /// ```
4384 ///
4385 /// into:
4386 ///
4387 /// ```
4388 /// %t0
4389 /// ```
4390 ///
4391 /// The producer of t1 may or may not be DCE'd depending on whether it is a
4392 /// block argument or has side effects.
4393 static LogicalResult foldReadInitWrite(TransferWriteOp write,
4395  SmallVectorImpl<OpFoldResult> &results) {
4396  // TODO: support 0-d corner case.
4397  if (write.getTransferRank() == 0)
4398  return failure();
4399  auto rankedTensorType =
4400  llvm::dyn_cast<RankedTensorType>(write.getSource().getType());
4401  // If not operating on tensors, bail.
4402  if (!rankedTensorType)
4403  return failure();
4404  // If no read, bail.
4405  auto read = write.getVector().getDefiningOp<vector::TransferReadOp>();
4406  if (!read)
4407  return failure();
4408  // TODO: support 0-d corner case.
4409  if (read.getTransferRank() == 0)
4410  return failure();
4411  // For now, only accept minor identity. Future: composition is minor identity.
4412  if (!read.getPermutationMap().isMinorIdentity() ||
4413  !write.getPermutationMap().isMinorIdentity())
4414  return failure();
4415  // Bail on mismatching ranks.
4416  if (read.getTransferRank() != write.getTransferRank())
4417  return failure();
4418  // Bail on potential out-of-bounds accesses.
4419  if (read.hasOutOfBoundsDim() || write.hasOutOfBoundsDim())
4420  return failure();
4421  // Tensor types must be the same.
4422  if (read.getSource().getType() != rankedTensorType)
4423  return failure();
4424  // Vector types must be the same.
4425  if (read.getVectorType() != write.getVectorType())
4426  return failure();
4427  // Vector and Tensor shapes must match.
4428  if (read.getVectorType().getShape() != rankedTensorType.getShape())
4429  return failure();
4430  // If any index is nonzero.
4431  auto isNotConstantZero = [](Value v) {
4432  auto cstOp = getConstantIntValue(v);
4433  return !cstOp.has_value() || cstOp.value() != 0;
4434  };
4435  if (llvm::any_of(read.getIndices(), isNotConstantZero) ||
4436  llvm::any_of(write.getIndices(), isNotConstantZero))
4437  return failure();
4438  // Success.
4439  results.push_back(read.getSource());
4440  return success();
4441 }
4442 
4443 static bool checkSameValueWAR(vector::TransferReadOp read,
4444  vector::TransferWriteOp write) {
4445  return read.getSource() == write.getSource() &&
4446  read.getIndices() == write.getIndices() &&
4447  read.getPermutationMap() == write.getPermutationMap() &&
4448  read.getVectorType() == write.getVectorType() && !read.getMask() &&
4449  !write.getMask();
4450 }
4451 /// Fold transfer_write write after read:
4452 /// ```
4453 /// %t0 = ...
4454 /// %v = vector.transfer_read %t0[%c0...] :
4455 /// tensor<static_sizesxf32>, vector<static_sizesxf32>
4456 /// %t1 = vector.transfer_write %v, %t0[%c0...] :
4457 /// vector<static_sizesxf32>, tensor<static_sizesxf32>
4458 /// ```
4459 ///
4460 /// into:
4461 ///
4462 /// ```
4463 /// %t0
4464 /// ```
4465 static LogicalResult foldWAR(TransferWriteOp write,
4466  SmallVectorImpl<OpFoldResult> &results) {
4467  if (!llvm::isa<RankedTensorType>(write.getSource().getType()))
4468  return failure();
4469  auto read = write.getVector().getDefiningOp<vector::TransferReadOp>();
4470  if (!read)
4471  return failure();
4472 
4473  if (!checkSameValueWAR(read, write))
4474  return failure();
4475  results.push_back(read.getSource());
4476  return success();
4477 }
4478 
4479 LogicalResult TransferWriteOp::fold(FoldAdaptor adaptor,
4480  SmallVectorImpl<OpFoldResult> &results) {
4481  if (succeeded(foldReadInitWrite(*this, adaptor.getOperands(), results)))
4482  return success();
4483  if (succeeded(foldWAR(*this, results)))
4484  return success();
4486  return success();
4487  if (succeeded(foldTransferFullMask(*this)))
4488  return success();
4489  return memref::foldMemRefCast(*this);
4490 }
4491 
4492 std::optional<SmallVector<int64_t, 4>> TransferWriteOp::getShapeForUnroll() {
4493  return llvm::to_vector<4>(getVectorType().getShape());
4494 }
4495 
4496 void TransferWriteOp::getEffects(
4498  &effects) {
4499  if (llvm::isa<MemRefType>(getShapedType()))
4500  effects.emplace_back(MemoryEffects::Write::get(), getSource(),
4502 }
4503 
4504 namespace {
4505 /// Remove dead transfer write from the SSA chain so that it an be eliminated by
4506 /// DCE
4507 /// ```
4508 /// %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
4509 /// : vector<1x4xf32>, tensor<4x4xf32>
4510 /// %w1 = vector.transfer_write %v0, %w0[%c2, %c0] {in_bounds = [true, true]}
4511 /// : vector<1x4xf32>, tensor<4x4xf32>
4512 /// %w2 = vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]}
4513 /// : vector<1x4xf32>, tensor<4x4xf32>
4514 /// ```
4515 ///
4516 /// into:
4517 ///
4518 /// ```
4519 /// %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
4520 /// : vector<1x4xf32>, tensor<4x4xf32>
4521 /// %w1 = vector.transfer_write %v0, %arg0[%c2, %c0] {in_bounds = [true, true]}
4522 /// : vector<1x4xf32>, tensor<4x4xf32>
4523 /// %w2 = vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]}
4524 /// : vector<1x4xf32>, tensor<4x4xf32>
4525 /// ```
4526 ///
4527 /// `%w0 = vector.transfer_write` op will be removed by DCE if it doesn't have
4528 /// any other uses.
4529 class FoldWaw final : public OpRewritePattern<TransferWriteOp> {
4530 public:
4532  LogicalResult matchAndRewrite(TransferWriteOp writeOp,
4533  PatternRewriter &rewriter) const override {
4534  if (!llvm::isa<RankedTensorType>(writeOp.getShapedType()))
4535  return failure();
4536  vector::TransferWriteOp writeToModify = writeOp;
4537 
4538  auto defWrite =
4539  writeOp.getSource().getDefiningOp<vector::TransferWriteOp>();
4540  while (defWrite) {
4541  if (checkSameValueWAW(writeOp, defWrite)) {
4542  rewriter.modifyOpInPlace(writeToModify, [&]() {
4543  writeToModify.getSourceMutable().assign(defWrite.getSource());
4544  });
4545  return success();
4546  }
4548  cast<VectorTransferOpInterface>(defWrite.getOperation()),
4549  cast<VectorTransferOpInterface>(writeOp.getOperation())))
4550  break;
4551  // If the previous write op doesn't have any other use we an safely look
4552  // at the previous store to see if it can be removed.
4553  if (!defWrite->hasOneUse())
4554  break;
4555  writeToModify = defWrite;
4556  defWrite = defWrite.getSource().getDefiningOp<vector::TransferWriteOp>();
4557  }
4558  return failure();
4559  }
4560 };
4561 
4562 /// Rewrite tensor::ExtractSliceOp(vector::TransferWriteOp) to
4563 /// vector::TransferWriteOp(tensor::ExtractSliceOp) if the full slice is
4564 /// overwritten and inserted into another tensor. After this rewrite, the
4565 /// operations bufferize in-place since all of them work on the same slice.
4566 ///
4567 /// For example:
4568 /// ```mlir
4569 /// %0 = vector.transfer_write %vec, %init_tensor[%c0, %c0]
4570 /// : vector<8x16xf32>, tensor<8x16xf32>
4571 /// %1 = tensor.extract_slice %0[0, 0] [%sz0, %sz1] [1, 1]
4572 /// : tensor<8x16xf32> to tensor<?x?xf32>
4573 /// %r = tensor.insert_slice %1 into %iter_arg[%iv0, %iv1] [%sz0, %sz1] [1, 1]
4574 /// : tensor<?x?xf32> into tensor<27x37xf32>
4575 /// ```
4576 /// folds to
4577 /// ```mlir
4578 /// %0 = tensor.extract_slice %iter_arg[%iv0, %iv1] [%sz0, %sz1] [1, 1]
4579 /// : tensor<27x37xf32> to tensor<?x?xf32>
4580 /// %1 = vector.transfer_write %vec, %0[%c0, %c0]
4581 /// : vector<8x16xf32>, tensor<?x?xf32>
4582 /// %r = tensor.insert_slice %1 into %iter_arg[%iv0, %iv1] [%sz0, %sz1] [1, 1]
4583 /// : tensor<?x?xf32> into tensor<27x37xf32>
4584 /// ```
4585 struct SwapExtractSliceOfTransferWrite
4586  : public OpRewritePattern<tensor::InsertSliceOp> {
4587 public:
4589 
4590  LogicalResult matchAndRewrite(tensor::InsertSliceOp insertOp,
4591  PatternRewriter &rewriter) const override {
4592  if (!insertOp.hasUnitStride())
4593  return failure();
4594  auto extractOp =
4595  insertOp.getSource().getDefiningOp<tensor::ExtractSliceOp>();
4596  if (!extractOp || !extractOp.hasUnitStride() || !extractOp->hasOneUse())
4597  return failure();
4598  auto transferOp = extractOp.getSource().getDefiningOp<TransferWriteOp>();
4599  if (!transferOp || !transferOp->hasOneUse())
4600  return failure();
4601 
4602  // Fail if vector::TransferWriteOp or tensor::ExtractSliceOp is
4603  // rank-reducing.
4604  if (insertOp.getSourceType().getRank() != transferOp.getTransferRank()) {
4605  return rewriter.notifyMatchFailure(insertOp,
4606  "use-def chain is rank-reducing");
4607  }
4608 
4609  // Fail if tensor::ExtractSliceOp has non-zero offset.
4610  if (!extractOp.hasZeroOffset()) {
4611  return rewriter.notifyMatchFailure(insertOp,
4612  "ExtractSliceOp has non-zero offset");
4613  }
4614 
4615  // Fail if tensor::TransferWriteOp has non-zero offset.
4616  if (!llvm::all_of(transferOp.getIndices(), [](Value value) {
4617  return getConstantIntValue(value) == static_cast<int64_t>(0);
4618  })) {
4619  return rewriter.notifyMatchFailure(insertOp,
4620  "TranferWriteOp has non-zero offset");
4621  }
4622 
4623  // Fail if tensor::ExtractSliceOp and tensor::InsertSliceOp sizes differ.
4624  if (insertOp.getMixedSizes().size() != extractOp.getMixedSizes().size()) {
4625  return rewriter.notifyMatchFailure(
4626  insertOp, "InsertSliceOp and ExtractSliceOp ranks differ");
4627  }
4628 
4629  for (auto [insertSize, extractSize] :
4630  llvm::zip_equal(insertOp.getMixedSizes(), extractOp.getMixedSizes())) {
4631  if (!isEqualConstantIntOrValue(insertSize, extractSize)) {
4632  return rewriter.notifyMatchFailure(
4633  insertOp, "InsertSliceOp and ExtractSliceOp sizes differ");
4634  }
4635  }
4636 
4637  // Fail if the vector::TransferWriteOp may not overwrite the full tensor.
4638  assert(transferOp.getVectorType().hasStaticShape() &&
4639  "expected vector to have a static shape");
4640  ArrayRef<int64_t> vectorShape = transferOp.getVectorType().getShape();
4642  transferOp.getPermutationMap(), transferOp.getShapedType().getShape());
4643  if (transferOp.getMask() || !vectorShape.equals(resultShape)) {
4644  return rewriter.notifyMatchFailure(
4645  insertOp, "TransferWriteOp may not write the full tensor.");
4646  }
4647 
4648  // Swap the tensor::ExtractSliceOp in front of the vector::TransferWriteOp.
4649  // Set all in_bounds to false and let the folder infer them.
4650  SmallVector<bool> newInBounds(vectorShape.size(), false);
4651  auto newExtractOp = rewriter.create<tensor::ExtractSliceOp>(
4652  extractOp.getLoc(), insertOp.getSourceType(), insertOp.getDest(),
4653  insertOp.getMixedOffsets(), insertOp.getMixedSizes(),
4654  insertOp.getMixedStrides());
4655  auto newTransferWriteOp = rewriter.create<TransferWriteOp>(
4656  transferOp.getLoc(), transferOp.getVector(), newExtractOp.getResult(),
4657  transferOp.getIndices(), transferOp.getPermutationMapAttr(),
4658  rewriter.getBoolArrayAttr(newInBounds));
4659  rewriter.modifyOpInPlace(insertOp, [&]() {
4660  insertOp.getSourceMutable().assign(newTransferWriteOp.getResult());
4661  });
4662  return success();
4663  }
4664 };
4665 
4666 } // namespace
4667 
4668 void TransferWriteOp::getCanonicalizationPatterns(RewritePatternSet &results,
4669  MLIRContext *context) {
4670  results.add<FoldWaw, SwapExtractSliceOfTransferWrite>(context);
4671 }
4672 
4673 //===----------------------------------------------------------------------===//
4674 // LoadOp
4675 //===----------------------------------------------------------------------===//
4676 
4677 static LogicalResult verifyLoadStoreMemRefLayout(Operation *op,
4678  MemRefType memRefTy) {
4679  if (!isLastMemrefDimUnitStride(memRefTy))
4680  return op->emitOpError("most minor memref dim must have unit stride");
4681  return success();
4682 }
4683 
4685  VectorType resVecTy = getVectorType();
4686  MemRefType memRefTy = getMemRefType();
4687 
4688  if (failed(verifyLoadStoreMemRefLayout(*this, memRefTy)))
4689  return failure();
4690 
4691  // Checks for vector memrefs.
4692  Type memElemTy = memRefTy.getElementType();
4693  if (auto memVecTy = llvm::dyn_cast<VectorType>(memElemTy)) {
4694  if (memVecTy != resVecTy)
4695  return emitOpError("base memref and result vector types should match");
4696  memElemTy = memVecTy.getElementType();
4697  }
4698 
4699  if (resVecTy.getElementType() != memElemTy)
4700  return emitOpError("base and result element types should match");
4701  if (llvm::size(getIndices()) != memRefTy.getRank())
4702  return emitOpError("requires ") << memRefTy.getRank() << " indices";
4703  return success();
4704 }
4705 
4706 OpFoldResult LoadOp::fold(FoldAdaptor) {
4707  if (succeeded(memref::foldMemRefCast(*this)))
4708  return getResult();
4709  return OpFoldResult();
4710 }
4711 
4712 //===----------------------------------------------------------------------===//
4713 // StoreOp
4714 //===----------------------------------------------------------------------===//
4715 
4717  VectorType valueVecTy = getVectorType();
4718  MemRefType memRefTy = getMemRefType();
4719 
4720  if (failed(verifyLoadStoreMemRefLayout(*this, memRefTy)))
4721  return failure();
4722 
4723  // Checks for vector memrefs.
4724  Type memElemTy = memRefTy.getElementType();
4725  if (auto memVecTy = llvm::dyn_cast<VectorType>(memElemTy)) {
4726  if (memVecTy != valueVecTy)
4727  return emitOpError(
4728  "base memref and valueToStore vector types should match");
4729  memElemTy = memVecTy.getElementType();
4730  }
4731 
4732  if (valueVecTy.getElementType() != memElemTy)
4733  return emitOpError("base and valueToStore element type should match");
4734  if (llvm::size(getIndices()) != memRefTy.getRank())
4735  return emitOpError("requires ") << memRefTy.getRank() << " indices";
4736  return success();
4737 }
4738 
4739 LogicalResult StoreOp::fold(FoldAdaptor adaptor,
4740  SmallVectorImpl<OpFoldResult> &results) {
4741  return memref::foldMemRefCast(*this);
4742 }
4743 
4744 //===----------------------------------------------------------------------===//
4745 // MaskedLoadOp
4746 //===----------------------------------------------------------------------===//
4747 
4749  VectorType maskVType = getMaskVectorType();
4750  VectorType passVType = getPassThruVectorType();
4751  VectorType resVType = getVectorType();
4752  MemRefType memType = getMemRefType();
4753 
4754  if (resVType.getElementType() != memType.getElementType())
4755  return emitOpError("base and result element type should match");
4756  if (llvm::size(getIndices()) != memType.getRank())
4757  return emitOpError("requires ") << memType.getRank() << " indices";
4758  if (resVType.getDimSize(0) != maskVType.getDimSize(0))
4759  return emitOpError("expected result dim to match mask dim");
4760  if (resVType != passVType)
4761  return emitOpError("expected pass_thru of same type as result type");
4762  return success();
4763 }
4764 
4765 namespace {
4766 class MaskedLoadFolder final : public OpRewritePattern<MaskedLoadOp> {
4767 public:
4769  LogicalResult matchAndRewrite(MaskedLoadOp load,
4770  PatternRewriter &rewriter) const override {
4771  switch (getMaskFormat(load.getMask())) {
4772  case MaskFormat::AllTrue:
4773  rewriter.replaceOpWithNewOp<vector::LoadOp>(
4774  load, load.getType(), load.getBase(), load.getIndices());
4775  return success();
4776  case MaskFormat::AllFalse:
4777  rewriter.replaceOp(load, load.getPassThru());
4778  return success();
4779  case MaskFormat::Unknown:
4780  return failure();
4781  }
4782  llvm_unreachable("Unexpected 1DMaskFormat on MaskedLoad");
4783  }
4784 };
4785 } // namespace
4786 
4787 void MaskedLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
4788  MLIRContext *context) {
4789  results.add<MaskedLoadFolder>(context);
4790 }
4791 
4792 OpFoldResult MaskedLoadOp::fold(FoldAdaptor) {
4793  if (succeeded(memref::foldMemRefCast(*this)))
4794  return getResult();
4795  return OpFoldResult();
4796 }
4797 
4798 //===----------------------------------------------------------------------===//
4799 // MaskedStoreOp
4800 //===----------------------------------------------------------------------===//
4801 
4803  VectorType maskVType = getMaskVectorType();
4804  VectorType valueVType = getVectorType();
4805  MemRefType memType = getMemRefType();
4806 
4807  if (valueVType.getElementType() != memType.getElementType())
4808  return emitOpError("base and valueToStore element type should match");
4809  if (llvm::size(getIndices()) != memType.getRank())
4810  return emitOpError("requires ") << memType.getRank() << " indices";
4811  if (valueVType.getDimSize(0) != maskVType.getDimSize(0))
4812  return emitOpError("expected valueToStore dim to match mask dim");
4813  return success();
4814 }
4815 
4816 namespace {
4817 class MaskedStoreFolder final : public OpRewritePattern<MaskedStoreOp> {
4818 public:
4820  LogicalResult matchAndRewrite(MaskedStoreOp store,
4821  PatternRewriter &rewriter) const override {
4822  switch (getMaskFormat(store.getMask())) {
4823  case MaskFormat::AllTrue:
4824  rewriter.replaceOpWithNewOp<vector::StoreOp>(
4825  store, store.getValueToStore(), store.getBase(), store.getIndices());
4826  return success();
4827  case MaskFormat::AllFalse:
4828  rewriter.eraseOp(store);
4829  return success();
4830  case MaskFormat::Unknown:
4831  return failure();
4832  }
4833  llvm_unreachable("Unexpected 1DMaskFormat on MaskedStore");
4834  }
4835 };
4836 } // namespace
4837 
4838 void MaskedStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
4839  MLIRContext *context) {
4840  results.add<MaskedStoreFolder>(context);
4841 }
4842 
4843 LogicalResult MaskedStoreOp::fold(FoldAdaptor adaptor,
4844  SmallVectorImpl<OpFoldResult> &results) {
4845  return memref::foldMemRefCast(*this);
4846 }
4847 
4848 //===----------------------------------------------------------------------===//
4849 // GatherOp
4850 //===----------------------------------------------------------------------===//
4851 
4853  VectorType indVType = getIndexVectorType();
4854  VectorType maskVType = getMaskVectorType();
4855  VectorType resVType = getVectorType();
4856  ShapedType baseType = getBaseType();
4857 
4858  if (!llvm::isa<MemRefType, RankedTensorType>(baseType))
4859  return emitOpError("requires base to be a memref or ranked tensor type");
4860 
4861  if (resVType.getElementType() != baseType.getElementType())
4862  return emitOpError("base and result element type should match");
4863  if (llvm::size(getIndices()) != baseType.getRank())
4864  return emitOpError("requires ") << baseType.getRank() << " indices";
4865  if (resVType.getShape() != indVType.getShape())
4866  return emitOpError("expected result dim to match indices dim");
4867  if (resVType.getShape() != maskVType.getShape())
4868  return emitOpError("expected result dim to match mask dim");
4869  if (resVType != getPassThruVectorType())
4870  return emitOpError("expected pass_thru of same type as result type");
4871  return success();
4872 }
4873 
4874 // MaskableOpInterface methods.
4875 
4876 /// Returns the mask type expected by this operation. Mostly used for
4877 /// verification purposes. It requires the operation to be vectorized."
4878 Type GatherOp::getExpectedMaskType() {
4879  auto vecType = this->getIndexVectorType();
4880  return VectorType::get(vecType.getShape(),
4881  IntegerType::get(vecType.getContext(), /*width=*/1),
4882  vecType.getScalableDims());
4883 }
4884 
4885 std::optional<SmallVector<int64_t, 4>> GatherOp::getShapeForUnroll() {
4886  return llvm::to_vector<4>(getVectorType().getShape());
4887 }
4888 
4889 namespace {
4890 class GatherFolder final : public OpRewritePattern<GatherOp> {
4891 public:
4893  LogicalResult matchAndRewrite(GatherOp gather,
4894  PatternRewriter &rewriter) const override {
4895  switch (getMaskFormat(gather.getMask())) {
4896  case MaskFormat::AllTrue:
4897  return failure(); // no unmasked equivalent
4898  case MaskFormat::AllFalse:
4899  rewriter.replaceOp(gather, gather.getPassThru());
4900  return success();
4901  case MaskFormat::Unknown:
4902  return failure();
4903  }
4904  llvm_unreachable("Unexpected 1DMaskFormat on GatherFolder");
4905  }
4906 };
4907 } // namespace
4908 
4909 void GatherOp::getCanonicalizationPatterns(RewritePatternSet &results,
4910  MLIRContext *context) {
4911  results.add<GatherFolder>(context);
4912 }
4913 
4914 //===----------------------------------------------------------------------===//
4915 // ScatterOp
4916 //===----------------------------------------------------------------------===//
4917 
4919  VectorType indVType = getIndexVectorType();
4920  VectorType maskVType = getMaskVectorType();
4921  VectorType valueVType = getVectorType();
4922  MemRefType memType = getMemRefType();
4923 
4924  if (valueVType.getElementType() != memType.getElementType())
4925  return emitOpError("base and valueToStore element type should match");
4926  if (llvm::size(getIndices()) != memType.getRank())
4927  return emitOpError("requires ") << memType.getRank() << " indices";
4928  if (valueVType.getDimSize(0) != indVType.getDimSize(0))
4929  return emitOpError("expected valueToStore dim to match indices dim");
4930  if (valueVType.getDimSize(0) != maskVType.getDimSize(0))
4931  return emitOpError("expected valueToStore dim to match mask dim");
4932  return success();
4933 }
4934 
4935 namespace {
4936 class ScatterFolder final : public OpRewritePattern<ScatterOp> {
4937 public:
4939  LogicalResult matchAndRewrite(ScatterOp scatter,
4940  PatternRewriter &rewriter) const override {
4941  switch (getMaskFormat(scatter.getMask())) {
4942  case MaskFormat::AllTrue:
4943  return failure(); // no unmasked equivalent
4944  case MaskFormat::AllFalse:
4945  rewriter.eraseOp(scatter);
4946  return success();
4947  case MaskFormat::Unknown:
4948  return failure();
4949  }
4950  llvm_unreachable("Unexpected 1DMaskFormat on ScatterFolder");
4951  }
4952 };
4953 } // namespace
4954 
4955 void ScatterOp::getCanonicalizationPatterns(RewritePatternSet &results,
4956  MLIRContext *context) {
4957  results.add<ScatterFolder>(context);
4958 }
4959 
4960 //===----------------------------------------------------------------------===//
4961 // ExpandLoadOp
4962 //===----------------------------------------------------------------------===//
4963 
4965  VectorType maskVType = getMaskVectorType();
4966  VectorType passVType = getPassThruVectorType();
4967  VectorType resVType = getVectorType();
4968  MemRefType memType = getMemRefType();
4969 
4970  if (resVType.getElementType() != memType.getElementType())
4971  return emitOpError("base and result element type should match");
4972  if (llvm::size(getIndices()) != memType.getRank())
4973  return emitOpError("requires ") << memType.getRank() << " indices";
4974  if (resVType.getDimSize(0) != maskVType.getDimSize(0))
4975  return emitOpError("expected result dim to match mask dim");
4976  if (resVType != passVType)
4977  return emitOpError("expected pass_thru of same type as result type");
4978  return success();
4979 }
4980 
4981 namespace {
4982 class ExpandLoadFolder final : public OpRewritePattern<ExpandLoadOp> {
4983 public:
4985  LogicalResult matchAndRewrite(ExpandLoadOp expand,
4986  PatternRewriter &rewriter) const override {
4987  switch (getMaskFormat(expand.getMask())) {
4988  case MaskFormat::AllTrue:
4989  rewriter.replaceOpWithNewOp<vector::LoadOp>(
4990  expand, expand.getType(), expand.getBase(), expand.getIndices());
4991  return success();
4992  case MaskFormat::AllFalse:
4993  rewriter.replaceOp(expand, expand.getPassThru());
4994  return success();
4995  case MaskFormat::Unknown:
4996  return failure();
4997  }
4998  llvm_unreachable("Unexpected 1DMaskFormat on ExpandLoadFolder");
4999  }
5000 };
5001 } // namespace
5002 
5003 void ExpandLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
5004  MLIRContext *context) {
5005  results.add<ExpandLoadFolder>(context);
5006 }
5007 
5008 //===----------------------------------------------------------------------===//
5009 // CompressStoreOp
5010 //===----------------------------------------------------------------------===//
5011 
5013  VectorType maskVType = getMaskVectorType();
5014  VectorType valueVType = getVectorType();
5015  MemRefType memType = getMemRefType();
5016 
5017  if (valueVType.getElementType() != memType.getElementType())
5018  return emitOpError("base and valueToStore element type should match");
5019  if (llvm::size(getIndices()) != memType.getRank())
5020  return emitOpError("requires ") << memType.getRank() << " indices";
5021  if (valueVType.getDimSize(0) != maskVType.getDimSize(0))
5022  return emitOpError("expected valueToStore dim to match mask dim");
5023  return success();
5024 }
5025 
5026 namespace {
5027 class CompressStoreFolder final : public OpRewritePattern<CompressStoreOp> {
5028 public:
5030  LogicalResult matchAndRewrite(CompressStoreOp compress,
5031  PatternRewriter &rewriter) const override {
5032  switch (getMaskFormat(compress.getMask())) {
5033  case MaskFormat::AllTrue:
5034  rewriter.replaceOpWithNewOp<vector::StoreOp>(
5035  compress, compress.getValueToStore(), compress.getBase(),
5036  compress.getIndices());
5037  return success();
5038  case MaskFormat::AllFalse:
5039  rewriter.eraseOp(compress);
5040  return success();
5041  case MaskFormat::Unknown:
5042  return failure();
5043  }
5044  llvm_unreachable("Unexpected 1DMaskFormat on CompressStoreFolder");
5045  }
5046 };
5047 } // namespace
5048 
5049 void CompressStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
5050  MLIRContext *context) {
5051  results.add<CompressStoreFolder>(context);
5052 }
5053 
5054 //===----------------------------------------------------------------------===//
5055 // ShapeCastOp
5056 //===----------------------------------------------------------------------===//
5057 
5058 /// Returns true if each element of 'a' is equal to the product of a contiguous
5059 /// sequence of the elements of 'b'. Returns false otherwise.
5060 static bool isValidShapeCast(ArrayRef<int64_t> a, ArrayRef<int64_t> b) {
5061  unsigned rankA = a.size();
5062  unsigned rankB = b.size();
5063  assert(rankA < rankB);
5064 
5065  auto isOne = [](int64_t v) { return v == 1; };
5066 
5067  // Special-case for n-D to 0-d shape cast. 'b' must be all ones to be shape
5068  // casted to a 0-d vector.
5069  if (rankA == 0 && llvm::all_of(b, isOne))
5070  return true;
5071 
5072  unsigned i = 0;
5073  unsigned j = 0;
5074  while (i < rankA && j < rankB) {
5075  int64_t dimA = a[i];
5076  int64_t dimB = 1;
5077  while (dimB < dimA && j < rankB)
5078  dimB *= b[j++];
5079  if (dimA != dimB)
5080  break;
5081  ++i;
5082 
5083  // Handle the case when trailing dimensions are of size 1.
5084  // Include them into the contiguous sequence.
5085  if (i < rankA && llvm::all_of(a.slice(i), isOne))
5086  i = rankA;
5087  if (j < rankB && llvm::all_of(b.slice(j), isOne))
5088  j = rankB;
5089  }
5090 
5091  return i == rankA && j == rankB;
5092 }
5093 
5094 static LogicalResult verifyVectorShapeCast(Operation *op,
5095  VectorType sourceVectorType,
5096  VectorType resultVectorType) {
5097  // Check that element type is the same.
5098  if (sourceVectorType.getElementType() != resultVectorType.getElementType())
5099  return op->emitOpError("source/result vectors must have same element type");
5100  auto sourceShape = sourceVectorType.getShape();
5101  auto resultShape = resultVectorType.getShape();
5102 
5103  // Check that product of source dim sizes matches product of result dim sizes.
5104  int64_t sourceDimProduct = std::accumulate(
5105  sourceShape.begin(), sourceShape.end(), 1LL, std::multiplies<int64_t>{});
5106  int64_t resultDimProduct = std::accumulate(
5107  resultShape.begin(), resultShape.end(), 1LL, std::multiplies<int64_t>{});
5108  if (sourceDimProduct != resultDimProduct)
5109  return op->emitOpError("source/result number of elements must match");
5110 
5111  // Check that expanding/contracting rank cases.
5112  unsigned sourceRank = sourceVectorType.getRank();
5113  unsigned resultRank = resultVectorType.getRank();
5114  if (sourceRank < resultRank) {
5115  if (!isValidShapeCast(sourceShape, resultShape))
5116  return op->emitOpError("invalid shape cast");
5117  } else if (sourceRank > resultRank) {
5118  if (!isValidShapeCast(resultShape, sourceShape))
5119  return op->emitOpError("invalid shape cast");
5120  }
5121  return success();
5122 }
5123 
5125  auto sourceVectorType =
5126  llvm::dyn_cast_or_null<VectorType>(getSource().getType());
5127  auto resultVectorType =
5128  llvm::dyn_cast_or_null<VectorType>(getResult().getType());
5129 
5130  // Check if source/result are of vector type.
5131  if (sourceVectorType && resultVectorType)
5132  return verifyVectorShapeCast(*this, sourceVectorType, resultVectorType);
5133 
5134  return success();
5135 }
5136 
5137 OpFoldResult ShapeCastOp::fold(FoldAdaptor adaptor) {
5138  // No-op shape cast.
5139  if (getSource().getType() == getResult().getType())
5140  return getSource();
5141 
5142  // Canceling shape casts.
5143  if (auto otherOp = getSource().getDefiningOp<ShapeCastOp>()) {
5144  if (getResult().getType() == otherOp.getSource().getType())
5145  return otherOp.getSource();
5146 
5147  // Only allows valid transitive folding.
5148  VectorType srcType = llvm::cast<VectorType>(otherOp.getSource().getType());
5149  VectorType resultType = llvm::cast<VectorType>(getResult().getType());
5150  if (srcType.getRank() < resultType.getRank()) {
5151  if (!isValidShapeCast(srcType.getShape(), resultType.getShape()))
5152  return {};
5153  } else if (srcType.getRank() > resultType.getRank()) {
5154  if (!isValidShapeCast(resultType.getShape(), srcType.getShape()))
5155  return {};
5156  } else {
5157  return {};
5158  }
5159 
5160  setOperand(otherOp.getSource());
5161  return getResult();
5162  }
5163 
5164  // Cancelling broadcast and shape cast ops.
5165  if (auto bcastOp = getSource().getDefiningOp<BroadcastOp>()) {
5166  if (bcastOp.getSourceType() == getType())
5167  return bcastOp.getSource();
5168  }
5169 
5170  return {};
5171 }
5172 
5173 namespace {
5174 // Pattern to rewrite a ShapeCast(splat ConstantOp) -> ConstantOp.
5175 class ShapeCastConstantFolder final : public OpRewritePattern<ShapeCastOp> {
5176 public:
5178 
5179  LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
5180  PatternRewriter &rewriter) const override {
5181  auto constantOp =
5182  shapeCastOp.getSource().getDefiningOp<arith::ConstantOp>();
5183  if (!constantOp)
5184  return failure();
5185  // Only handle splat for now.
5186  auto dense = llvm::dyn_cast<SplatElementsAttr>(constantOp.getValue());
5187  if (!dense)
5188  return failure();
5189  auto newAttr =
5190  DenseElementsAttr::get(llvm::cast<VectorType>(shapeCastOp.getType()),
5191  dense.getSplatValue<Attribute>());
5192  rewriter.replaceOpWithNewOp<arith::ConstantOp>(shapeCastOp, newAttr);
5193  return success();
5194  }
5195 };
5196 
5197 /// Helper function that computes a new vector type based on the input vector
5198 /// type by removing the trailing one dims:
5199 ///
5200 /// vector<4x1x1xi1> --> vector<4x1>
5201 ///
5202 static VectorType trimTrailingOneDims(VectorType oldType) {
5203  ArrayRef<int64_t> oldShape = oldType.getShape();
5204  ArrayRef<int64_t> newShape = oldShape;
5205 
5206  ArrayRef<bool> oldScalableDims = oldType.getScalableDims();
5207  ArrayRef<bool> newScalableDims = oldScalableDims;
5208 
5209  while (!newShape.empty() && newShape.back() == 1 && !newScalableDims.back()) {
5210  newShape = newShape.drop_back(1);
5211  newScalableDims = newScalableDims.drop_back(1);
5212  }
5213 
5214  // Make sure we have at least 1 dimension.
5215  // TODO: Add support for 0-D vectors.
5216  if (newShape.empty()) {
5217  newShape = oldShape.take_back();
5218  newScalableDims = oldScalableDims.take_back();
5219  }
5220 
5221  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
5222 }
5223 
5224 /// Folds qualifying shape_cast(create_mask) into a new create_mask
5225 ///
5226 /// Looks at `vector.shape_cast` Ops that simply "drop" the trailing unit
5227 /// dimension. If the input vector comes from `vector.create_mask` for which
5228 /// the corresponding mask input value is 1 (e.g. `%c1` below), then it is safe
5229 /// to fold shape_cast into create_mask.
5230 ///
5231 /// BEFORE:
5232 /// %1 = vector.create_mask %c1, %dim, %c1, %c1 : vector<1x[4]x1x1xi1>
5233 /// %2 = vector.shape_cast %1 : vector<1x[4]x1x1xi1> to vector<1x[4]xi1>
5234 /// AFTER:
5235 /// %0 = vector.create_mask %c1, %dim : vector<1x[4]xi1>
5236 class ShapeCastCreateMaskFolderTrailingOneDim final
5237  : public OpRewritePattern<ShapeCastOp> {
5238 public:
5240 
5241  LogicalResult matchAndRewrite(ShapeCastOp shapeOp,
5242  PatternRewriter &rewriter) const override {
5243  Value shapeOpSrc = shapeOp->getOperand(0);
5244  auto createMaskOp = shapeOpSrc.getDefiningOp<vector::CreateMaskOp>();
5245  auto constantMaskOp = shapeOpSrc.getDefiningOp<vector::ConstantMaskOp>();
5246  if (!createMaskOp && !constantMaskOp)
5247  return failure();
5248 
5249  VectorType shapeOpResTy = shapeOp.getResultVectorType();
5250  VectorType shapeOpSrcTy = shapeOp.getSourceVectorType();
5251 
5252  VectorType newVecType = trimTrailingOneDims(shapeOpSrcTy);
5253  if (newVecType != shapeOpResTy)
5254  return failure();
5255 
5256  auto numDimsToDrop =
5257  shapeOpSrcTy.getShape().size() - shapeOpResTy.getShape().size();
5258 
5259  // No unit dims to drop
5260  if (!numDimsToDrop)
5261  return failure();
5262 
5263  if (createMaskOp) {
5264  auto maskOperands = createMaskOp.getOperands();
5265  auto numMaskOperands = maskOperands.size();
5266 
5267  // Check every mask dim size to see whether it can be dropped
5268  for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
5269  --i) {
5270  auto constant = maskOperands[i].getDefiningOp<arith::ConstantIndexOp>();
5271  if (!constant || (constant.value() != 1))
5272  return failure();
5273  }
5274  SmallVector<Value> newMaskOperands =
5275  maskOperands.drop_back(numDimsToDrop);
5276 
5277  rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(shapeOp, shapeOpResTy,
5278  newMaskOperands);
5279  return success();
5280  }
5281 
5282  if (constantMaskOp) {
5283  auto maskDimSizes = constantMaskOp.getMaskDimSizes().getValue();
5284  auto numMaskOperands = maskDimSizes.size();
5285 
5286  // Check every mask dim size to see whether it can be dropped
5287  for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
5288  --i) {
5289  if (cast<IntegerAttr>(maskDimSizes[i]).getValue() != 1)
5290  return failure();
5291  }
5292 
5293  auto newMaskOperands = maskDimSizes.drop_back(numDimsToDrop);
5294  ArrayAttr newMaskOperandsAttr = rewriter.getArrayAttr(newMaskOperands);
5295 
5296  rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(shapeOp, shapeOpResTy,
5297  newMaskOperandsAttr);
5298  return success();
5299  }
5300 
5301  return failure();
5302  }
5303 };
5304 
5305 /// Pattern to rewrite a ShapeCast(Broadcast) -> Broadcast.
5306 /// This only applies when the shape of the broadcast source
5307 /// 1. is a suffix of the shape of the result (i.e. when broadcast without
5308 /// reshape is expressive enough to capture the result in a single op), or
5309 /// 2. has the same element count as the shape cast result.
5310 class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
5311 public:
5313 
5314  LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
5315  PatternRewriter &rewriter) const override {
5316  auto broadcastOp =
5317  shapeCastOp.getSource().getDefiningOp<vector::BroadcastOp>();
5318  if (!broadcastOp)
5319  return failure();
5320 
5321  ArrayRef<int64_t> broadcastSourceShape;
5322  if (auto srcType = dyn_cast<VectorType>(broadcastOp.getSourceType()))
5323  broadcastSourceShape = srcType.getShape();
5324  ArrayRef<int64_t> shapeCastTargetShape =
5325  shapeCastOp.getResultVectorType().getShape();
5326 
5327  // If `broadcastSourceShape` is a suffix of the result, we can just replace
5328  // with a broadcast to the final shape.
5329  if (broadcastSourceShape ==
5330  shapeCastTargetShape.take_back(broadcastSourceShape.size())) {
5331  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
5332  shapeCastOp, shapeCastOp.getResultVectorType(),
5333  broadcastOp.getSource());
5334  return success();
5335  }
5336 
5337  // Otherwise, if the final result has the same element count, we can replace
5338  // with a shape cast.
5339  if (auto srcType = dyn_cast<VectorType>(broadcastOp.getSourceType())) {
5340  if (srcType.getNumElements() ==
5341  shapeCastOp.getResultVectorType().getNumElements()) {
5342  rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
5343  shapeCastOp, shapeCastOp.getResultVectorType(),
5344  broadcastOp.getSource());
5345  return success();
5346  }
5347  }
5348 
5349  return failure();
5350  }
5351 };
5352 
5353 } // namespace
5354 
5355 void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
5356  MLIRContext *context) {
5357  results.add<ShapeCastConstantFolder, ShapeCastCreateMaskFolderTrailingOneDim,
5358  ShapeCastBroadcastFolder>(context);
5359 }
5360 
5361 //===----------------------------------------------------------------------===//
5362 // VectorBitCastOp
5363 //===----------------------------------------------------------------------===//
5364 
5366  auto sourceVectorType = getSourceVectorType();
5367  auto resultVectorType = getResultVectorType();
5368 
5369  for (int64_t i = 0, e = sourceVectorType.getRank() - 1; i < e; i++) {
5370  if (sourceVectorType.getDimSize(i) != resultVectorType.getDimSize(i))
5371  return emitOpError("dimension size mismatch at: ") << i;
5372  }
5373 
5374  DataLayout dataLayout = DataLayout::closest(*this);
5375  auto sourceElementBits =
5376  dataLayout.getTypeSizeInBits(sourceVectorType.getElementType());
5377  auto resultElementBits =
5378  dataLayout.getTypeSizeInBits(resultVectorType.getElementType());
5379 
5380  if (sourceVectorType.getRank() == 0) {
5381  if (sourceElementBits != resultElementBits)
5382  return emitOpError("source/result bitwidth of the 0-D vector element "
5383  "types must be equal");
5384  } else if (sourceElementBits * sourceVectorType.getShape().back() !=
5385  resultElementBits * resultVectorType.getShape().back()) {
5386  return emitOpError(
5387  "source/result bitwidth of the minor 1-D vectors must be equal");
5388  }
5389 
5390  return success();
5391 }
5392 
5393 OpFoldResult BitCastOp::fold(FoldAdaptor adaptor) {
5394  // Nop cast.
5395  if (getSource().getType() == getResult().getType())
5396  return getSource();
5397 
5398  // Canceling bitcasts.
5399  if (auto otherOp = getSource().getDefiningOp<BitCastOp>()) {
5400  if (getResult().getType() == otherOp.getSource().getType())
5401  return otherOp.getSource();
5402 
5403  setOperand(otherOp.getSource());
5404  return getResult();
5405  }
5406 
5407  Attribute sourceConstant = adaptor.getSource();
5408  if (!sourceConstant)
5409  return {};
5410 
5411  Type srcElemType = getSourceVectorType().getElementType();
5412  Type dstElemType = getResultVectorType().getElementType();
5413 
5414  if (auto floatPack = llvm::dyn_cast<DenseFPElementsAttr>(sourceConstant)) {
5415  if (floatPack.isSplat()) {
5416  auto splat = floatPack.getSplatValue<FloatAttr>();
5417 
5418  // Casting fp16 into fp32.
5419  if (srcElemType.isF16() && dstElemType.isF32()) {
5420  uint32_t bits = static_cast<uint32_t>(
5421  splat.getValue().bitcastToAPInt().getZExtValue());
5422  // Duplicate the 16-bit pattern.
5423  bits = (bits << 16) | (bits & 0xffff);
5424  APInt intBits(32, bits);
5425  APFloat floatBits(llvm::APFloat::IEEEsingle(), intBits);
5426  return DenseElementsAttr::get(getResultVectorType(), floatBits);
5427  }
5428  }
5429  }
5430 
5431  if (auto intPack = llvm::dyn_cast<DenseIntElementsAttr>(sourceConstant)) {
5432  if (intPack.isSplat()) {
5433  auto splat = intPack.getSplatValue<IntegerAttr>();
5434 
5435  if (llvm::isa<IntegerType>(dstElemType)) {
5436  uint64_t srcBitWidth = srcElemType.getIntOrFloatBitWidth();
5437  uint64_t dstBitWidth = dstElemType.getIntOrFloatBitWidth();
5438 
5439  // Casting to a larger integer bit width.
5440  if (dstBitWidth > srcBitWidth && dstBitWidth % srcBitWidth == 0) {
5441  APInt intBits = splat.getValue().zext(dstBitWidth);
5442 
5443  // Duplicate the lower width element.
5444  for (uint64_t i = 0; i < dstBitWidth / srcBitWidth - 1; i++)
5445  intBits = (intBits << srcBitWidth) | intBits;
5446  return DenseElementsAttr::get(getResultVectorType(), intBits);
5447  }
5448  }
5449  }
5450  }
5451 
5452  return {};
5453 }
5454 
5455 //===----------------------------------------------------------------------===//
5456 // TypeCastOp
5457 //===----------------------------------------------------------------------===//
5458 
5459 static SmallVector<int64_t, 8> extractShape(MemRefType memRefType) {
5460  auto vectorType = llvm::dyn_cast<VectorType>(memRefType.getElementType());
5461  SmallVector<int64_t, 8> res(memRefType.getShape().begin(),
5462  memRefType.getShape().end());
5463  if (vectorType)
5464  res.append(vectorType.getShape().begin(), vectorType.getShape().end());
5465  return res;
5466 }
5467 
5468 /// Build the canonical memRefType with a single vector.
5469 /// E.g. memref<4 x 5 x vector<6 x f32>> -> memref<vector<4 x 5 x 6 x f32>>.
5470 void TypeCastOp::build(OpBuilder &builder, OperationState &result,
5471  Value source) {
5472  result.addOperands(source);
5473  MemRefType memRefType = llvm::cast<MemRefType>(source.getType());
5474  VectorType vectorType =
5475  VectorType::get(extractShape(memRefType),
5477  result.addTypes(MemRefType::get({}, vectorType, MemRefLayoutAttrInterface(),
5478  memRefType.getMemorySpace()));
5479 }
5480 
5482  MemRefType canonicalType = canonicalizeStridedLayout(getMemRefType());
5483  if (!canonicalType.getLayout().isIdentity())
5484  return emitOpError("expects operand to be a memref with identity layout");
5485  if (!getResultMemRefType().getLayout().isIdentity())
5486  return emitOpError("expects result to be a memref with identity layout");
5487  if (getResultMemRefType().getMemorySpace() !=
5488  getMemRefType().getMemorySpace())
5489  return emitOpError("expects result in same memory space");
5490 
5491  auto sourceType = getMemRefType();
5492  auto resultType = getResultMemRefType();
5493  if (getElementTypeOrSelf(getElementTypeOrSelf(sourceType)) !=
5495  return emitOpError(
5496  "expects result and operand with same underlying scalar type: ")
5497  << resultType;
5498  if (extractShape(sourceType) != extractShape(resultType))
5499  return emitOpError(
5500  "expects concatenated result and operand shapes to be equal: ")
5501  << resultType;
5502  return success();
5503 }
5504 
5505 //===----------------------------------------------------------------------===//
5506 // TransposeOp
5507 //===----------------------------------------------------------------------===//
5508 
5509 void vector::TransposeOp::build(OpBuilder &builder, OperationState &result,
5510  Value vector, ArrayRef<int64_t> permutation) {
5511  VectorType vt = llvm::cast<VectorType>(vector.getType());
5512  SmallVector<int64_t, 4> transposedShape(vt.getRank());
5513  SmallVector<bool, 4> transposedScalableDims(vt.getRank());
5514  for (unsigned i = 0; i < permutation.size(); ++i) {
5515  transposedShape[i] = vt.getShape()[permutation[i]];
5516  transposedScalableDims[i] = vt.getScalableDims()[permutation[i]];
5517  }
5518 
5519  result.addOperands(vector);
5520  result.addTypes(VectorType::get(transposedShape, vt.getElementType(),
5521  transposedScalableDims));
5522  result.addAttribute(TransposeOp::getPermutationAttrName(result.name),
5523  builder.getDenseI64ArrayAttr(permutation));
5524 }
5525 
5526 OpFoldResult vector::TransposeOp::fold(FoldAdaptor adaptor) {
5527  // Eliminate splat constant transpose ops.
5528  if (auto attr =
5529  llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getVector()))
5530  if (attr.isSplat())
5531  return attr.reshape(getResultVectorType());
5532 
5533  // Eliminate identity transpose ops. This happens when the dimensions of the
5534  // input vector remain in their original order after the transpose operation.
5535  ArrayRef<int64_t> perm = getPermutation();
5536 
5537  // Check if the permutation of the dimensions contains sequential values:
5538  // {0, 1, 2, ...}.
5539  for (int64_t i = 0, e = perm.size(); i < e; i++) {
5540  if (perm[i] != i)
5541  return {};
5542  }
5543 
5544  return getVector();
5545 }
5546 
5548  VectorType vectorType = getSourceVectorType();
5549  VectorType resultType = getResultVectorType();
5550  int64_t rank = resultType.getRank();
5551  if (vectorType.getRank() != rank)
5552  return emitOpError("vector result rank mismatch: ") << rank;
5553  // Verify transposition array.
5554  ArrayRef<int64_t> perm = getPermutation();
5555  int64_t size = perm.size();
5556  if (rank != size)
5557  return emitOpError("transposition length mismatch: ") << size;
5558  SmallVector<bool, 8> seen(rank, false);
5559  for (const auto &ta : llvm::enumerate(perm)) {
5560  if (ta.value() < 0 || ta.value() >= rank)
5561  return emitOpError("transposition index out of range: ") << ta.value();
5562  if (seen[ta.value()])
5563  return emitOpError("duplicate position index: ") << ta.value();
5564  seen[ta.value()] = true;
5565  if (resultType.getDimSize(ta.index()) != vectorType.getDimSize(ta.value()))
5566  return emitOpError("dimension size mismatch at: ") << ta.value();
5567  }
5568  return success();
5569 }
5570 
5571 std::optional<SmallVector<int64_t, 4>> TransposeOp::getShapeForUnroll() {
5572  return llvm::to_vector<4>(getResultVectorType().getShape());
5573 }
5574 
5575 namespace {
5576 
5577 // Rewrites two back-to-back TransposeOp operations into a single TransposeOp.
5578 class TransposeFolder final : public OpRewritePattern<vector::TransposeOp> {
5579 public:
5581 
5582  LogicalResult matchAndRewrite(vector::TransposeOp transposeOp,
5583  PatternRewriter &rewriter) const override {
5584  // Composes two permutations: result[i] = permutation1[permutation2[i]].
5585  auto composePermutations = [](ArrayRef<int64_t> permutation1,
5586  ArrayRef<int64_t> permutation2) {
5587  SmallVector<int64_t, 4> result;
5588  for (auto index : permutation2)
5589  result.push_back(permutation1[index]);
5590  return result;
5591  };
5592 
5593  // Return if the input of 'transposeOp' is not defined by another transpose.
5594  vector::TransposeOp parentTransposeOp =
5595  transposeOp.getVector().getDefiningOp<vector::TransposeOp>();
5596  if (!parentTransposeOp)
5597  return failure();
5598 
5599  SmallVector<int64_t, 4> permutation = composePermutations(
5600  parentTransposeOp.getPermutation(), transposeOp.getPermutation());
5601  // Replace 'transposeOp' with a new transpose operation.
5602  rewriter.replaceOpWithNewOp<vector::TransposeOp>(
5603  transposeOp, transposeOp.getResult().getType(),
5604  parentTransposeOp.getVector(), permutation);
5605  return success();
5606  }
5607 };
5608 
5609 // Folds transpose(broadcast(<scalar>)) into brodcast(<scalar>).
5610 struct FoldTransposedScalarBroadcast final
5611  : public OpRewritePattern<vector::TransposeOp> {
5613 
5614  LogicalResult matchAndRewrite(vector::TransposeOp transposeOp,
5615  PatternRewriter &rewriter) const override {
5616  auto bcastOp = transposeOp.getVector().getDefiningOp<vector::BroadcastOp>();
5617  if (!bcastOp)
5618  return failure();
5619 
5620  auto srcVectorType = llvm::dyn_cast<VectorType>(bcastOp.getSourceType());
5621  if (!srcVectorType || srcVectorType.getNumElements() == 1) {
5622  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
5623  transposeOp, transposeOp.getResultVectorType(), bcastOp.getSource());
5624  return success();
5625  }
5626 
5627  return failure();
5628  }
5629 };
5630 
5631 // Folds transpose(splat x : src_type) : res_type into splat x : res_type.
5632 class FoldTransposeSplat final : public OpRewritePattern<TransposeOp> {
5633 public:
5635 
5636  LogicalResult matchAndRewrite(TransposeOp transposeOp,
5637  PatternRewriter &rewriter) const override {
5638  auto splatOp = transposeOp.getVector().getDefiningOp<vector::SplatOp>();
5639  if (!splatOp)
5640  return failure();
5641 
5642  rewriter.replaceOpWithNewOp<vector::SplatOp>(
5643  transposeOp, transposeOp.getResultVectorType(), splatOp.getInput());
5644  return success();
5645  }
5646 };
5647 
5648 /// Folds transpose(create_mask) into a new transposed create_mask.
5649 class FoldTransposeCreateMask final : public OpRewritePattern<TransposeOp> {
5650 public:
5652 
5653  LogicalResult matchAndRewrite(TransposeOp transpOp,
5654  PatternRewriter &rewriter) const override {
5655  Value transposeSrc = transpOp.getVector();
5656  auto createMaskOp = transposeSrc.getDefiningOp<vector::CreateMaskOp>();
5657  auto constantMaskOp = transposeSrc.getDefiningOp<vector::ConstantMaskOp>();
5658  if (!createMaskOp && !constantMaskOp)
5659  return failure();
5660 
5661  // Get the transpose permutation and apply it to the vector.create_mask or
5662  // vector.constant_mask operands.
5663  ArrayRef<int64_t> permutation = transpOp.getPermutation();
5664 
5665  if (createMaskOp) {
5666  auto maskOperands = createMaskOp.getOperands();
5667  SmallVector<Value> newOperands(maskOperands.begin(), maskOperands.end());
5668  applyPermutationToVector(newOperands, permutation);
5669 
5670  rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(
5671  transpOp, transpOp.getResultVectorType(), newOperands);
5672  return success();
5673  }
5674 
5675  // ConstantMaskOp case.
5676  auto maskDimSizes = constantMaskOp.getMaskDimSizes();
5677  SmallVector<Attribute> newMaskDimSizes(maskDimSizes.getValue());
5678  applyPermutationToVector(newMaskDimSizes, permutation);
5679 
5680  rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(
5681  transpOp, transpOp.getResultVectorType(),
5682  ArrayAttr::get(transpOp.getContext(), newMaskDimSizes));
5683  return success();
5684  }
5685 };
5686 
5687 } // namespace
5688 
5689 void vector::TransposeOp::getCanonicalizationPatterns(
5690  RewritePatternSet &results, MLIRContext *context) {
5691  results.add<FoldTransposeCreateMask, FoldTransposedScalarBroadcast,
5692  TransposeFolder, FoldTransposeSplat>(context);
5693 }
5694 
5695 //===----------------------------------------------------------------------===//
5696 // ConstantMaskOp
5697 //===----------------------------------------------------------------------===//
5698 
5700  auto resultType = llvm::cast<VectorType>(getResult().getType());
5701  // Check the corner case of 0-D vectors first.
5702  if (resultType.getRank() == 0) {
5703  if (getMaskDimSizes().size() != 1)
5704  return emitError("array attr must have length 1 for 0-D vectors");
5705  auto dim = llvm::cast<IntegerAttr>(getMaskDimSizes()[0]).getInt();
5706  if (dim != 0 && dim != 1)
5707  return emitError("mask dim size must be either 0 or 1 for 0-D vectors");
5708  return success();
5709  }
5710 
5711  // Verify that array attr size matches the rank of the vector result.
5712  if (static_cast<int64_t>(getMaskDimSizes().size()) != resultType.getRank())
5713  return emitOpError(
5714  "must specify array attr of size equal vector result rank");
5715  // Verify that each array attr element is in bounds of corresponding vector
5716  // result dimension size.
5717  auto resultShape = resultType.getShape();
5718  auto resultScalableDims = resultType.getScalableDims();
5719  SmallVector<int64_t, 4> maskDimSizes;
5720  for (const auto [index, intAttr] : llvm::enumerate(getMaskDimSizes())) {
5721  int64_t maskDimSize = llvm::cast<IntegerAttr>(intAttr).getInt();
5722  if (maskDimSize < 0 || maskDimSize > resultShape[index])
5723  return emitOpError(
5724  "array attr of size out of bounds of vector result dimension size");
5725  if (resultScalableDims[index] && maskDimSize != 0 &&
5726  maskDimSize != resultShape[index])
5727  return emitOpError(
5728  "only supports 'none set' or 'all set' scalable dimensions");
5729  maskDimSizes.push_back(maskDimSize);
5730  }
5731  // Verify that if one mask dim size is zero, they all should be zero (because
5732  // the mask region is a conjunction of each mask dimension interval).
5733  bool anyZeros = llvm::is_contained(maskDimSizes, 0);
5734  bool allZeros = llvm::all_of(maskDimSizes, [](int64_t s) { return s == 0; });
5735  if (anyZeros && !allZeros)
5736  return emitOpError("expected all mask dim sizes to be zeros, "
5737  "as a result of conjunction with zero mask dim");
5738  return success();
5739 }
5740 
5741 bool ConstantMaskOp::isAllOnesMask() {
5742  auto resultType = getVectorType();
5743  // Check the corner case of 0-D vectors first.
5744  if (resultType.getRank() == 0) {
5745  assert(getMaskDimSizes().size() == 1 && "invalid sizes for zero rank mask");
5746  return llvm::cast<IntegerAttr>(getMaskDimSizes()[0]).getInt() == 1;
5747  }
5748  for (const auto [resultSize, intAttr] :
5749  llvm::zip_equal(resultType.getShape(), getMaskDimSizes())) {
5750  int64_t maskDimSize = llvm::cast<IntegerAttr>(intAttr).getInt();
5751  if (maskDimSize < resultSize)
5752  return false;
5753  }
5754  return true;
5755 }
5756 
5757 //===----------------------------------------------------------------------===//
5758 // CreateMaskOp
5759 //===----------------------------------------------------------------------===//
5760 
5761 void CreateMaskOp::build(OpBuilder &builder, OperationState &result,
5762  VectorType type,
5763  ArrayRef<OpFoldResult> mixedOperands) {
5764  SmallVector<Value> operands =
5765  getValueOrCreateConstantIndexOp(builder, result.location, mixedOperands);
5766  build(builder, result, type, operands);
5767 }
5768 
5770  auto vectorType = llvm::cast<VectorType>(getResult().getType());
5771  // Verify that an operand was specified for each result vector each dimension.
5772  if (vectorType.getRank() == 0) {
5773  if (getNumOperands() != 1)
5774  return emitOpError(
5775  "must specify exactly one operand for 0-D create_mask");
5776  } else if (getNumOperands() !=
5777  llvm::cast<VectorType>(getResult().getType()).getRank()) {
5778  return emitOpError(
5779  "must specify an operand for each result vector dimension");
5780  }
5781  return success();
5782 }
5783 
5784 namespace {
5785 
5786 /// Pattern to rewrite a CreateMaskOp with a ConstantMaskOp.
5787 ///
5788 /// Ex 1:
5789 /// %c2 = arith.constant 2 : index
5790 /// %c3 = arith.constant 3 : index
5791 /// %0 = vector.create_mask %c3, %c2 : vector<4x3xi1>
5792 /// Becomes:
5793 /// vector.constant_mask [3, 2] : vector<4x3xi1>
5794 ///
5795 /// Ex 2:
5796 /// %c_neg_1 = arith.constant -1 : index
5797 /// %0 = vector.create_mask %c_neg_1 : vector<[8]xi1>
5798 /// becomes:
5799 /// vector.constant_mask [0] : vector<[8]xi1>
5800 ///
5801 /// Ex 3:
5802 /// %c8 = arith.constant 8 : index
5803 /// %c16 = arith.constant 16 : index
5804 /// %0 = vector.vscale
5805 /// %1 = arith.muli %0, %c16 : index
5806 /// %10 = vector.create_mask %c8, %1 : vector<8x[16]xi1>
5807 /// becomes:
5808 /// %0 = vector.constant_mask [8, 16] : vector<8x[16]xi1>
5809 class CreateMaskFolder final : public OpRewritePattern<CreateMaskOp> {
5810 public:
5812 
5813  LogicalResult matchAndRewrite(CreateMaskOp createMaskOp,
5814  PatternRewriter &rewriter) const override {
5815  VectorType retTy = createMaskOp.getResult().getType();
5816  bool isScalable = retTy.isScalable();
5817 
5818  // Check every mask operand
5819  for (auto [opIdx, operand] : llvm::enumerate(createMaskOp.getOperands())) {
5820  if (auto cst = getConstantIntValue(operand)) {
5821  // Most basic case - this operand is a constant value. Note that for
5822  // scalable dimensions, CreateMaskOp can be folded only if the
5823  // corresponding operand is negative or zero.
5824  if (retTy.getScalableDims()[opIdx] && *cst > 0)
5825  return failure();
5826 
5827  continue;
5828  }
5829 
5830  // Non-constant operands are not allowed for non-scalable vectors.
5831  if (!isScalable)
5832  return failure();
5833 
5834  // For scalable vectors, "arith.muli %vscale, %dimSize" means an "all
5835  // true" mask, so can also be treated as constant.
5836  auto mul = operand.getDefiningOp<arith::MulIOp>();
5837  if (!mul)
5838  return failure();
5839  auto mulLHS = mul.getRhs();
5840  auto mulRHS = mul.getLhs();
5841  bool isOneOpVscale =
5842  (isa<vector::VectorScaleOp>(mulLHS.getDefiningOp()) ||
5843  isa<vector::VectorScaleOp>(mulRHS.getDefiningOp()));
5844 
5845  auto isConstantValMatchingDim =
5846  [=, dim = retTy.getShape()[opIdx]](Value operand) {
5847  auto constantVal = getConstantIntValue(operand);
5848  return (constantVal.has_value() && constantVal.value() == dim);
5849  };
5850 
5851  bool isOneOpConstantMatchingDim =
5852  isConstantValMatchingDim(mulLHS) || isConstantValMatchingDim(mulRHS);
5853 
5854  if (!isOneOpVscale || !isOneOpConstantMatchingDim)
5855  return failure();
5856  }
5857 
5858  // Gather constant mask dimension sizes.
5859  SmallVector<int64_t, 4> maskDimSizes;
5860  maskDimSizes.reserve(createMaskOp->getNumOperands());
5861  for (auto [operand, maxDimSize] : llvm::zip_equal(
5862  createMaskOp.getOperands(), createMaskOp.getType().getShape())) {
5863  std::optional dimSize = getConstantIntValue(operand);
5864  if (!dimSize) {
5865  // Although not a constant, it is safe to assume that `operand` is
5866  // "vscale * maxDimSize".
5867  maskDimSizes.push_back(maxDimSize);
5868  continue;
5869  }
5870  int64_t dimSizeVal = std::min(dimSize.value(), maxDimSize);
5871  // If one of dim sizes is zero, set all dims to zero.
5872  if (dimSize <= 0) {
5873  maskDimSizes.assign(createMaskOp.getType().getRank(), 0);
5874  break;
5875  }
5876  maskDimSizes.push_back(dimSizeVal);
5877  }
5878 
5879  // Replace 'createMaskOp' with ConstantMaskOp.
5880  rewriter.replaceOpWithNewOp<ConstantMaskOp>(
5881  createMaskOp, retTy,
5882  vector::getVectorSubscriptAttr(rewriter, maskDimSizes));
5883  return success();
5884  }
5885 };
5886 
5887 } // namespace
5888 
5889 void CreateMaskOp::getCanonicalizationPatterns(RewritePatternSet &results,
5890  MLIRContext *context) {
5891  results.add<CreateMaskFolder>(context);
5892 }
5893 
5894 //===----------------------------------------------------------------------===//
5895 // MaskOp
5896 //===----------------------------------------------------------------------===//
5897 
5898 void MaskOp::build(
5899  OpBuilder &builder, OperationState &result, Value mask,
5900  Operation *maskableOp,
5901  function_ref<void(OpBuilder &, Operation *)> maskRegionBuilder) {
5902  assert(maskRegionBuilder &&
5903  "builder callback for 'maskRegion' must be present");
5904 
5905  result.addOperands(mask);
5906  OpBuilder::InsertionGuard guard(builder);
5907  Region *maskRegion = result.addRegion();
5908  builder.createBlock(maskRegion);
5909  maskRegionBuilder(builder, maskableOp);
5910 }
5911 
5912 void MaskOp::build(
5913  OpBuilder &builder, OperationState &result, TypeRange resultTypes,
5914  Value mask, Operation *maskableOp,
5915  function_ref<void(OpBuilder &, Operation *)> maskRegionBuilder) {
5916  build(builder, result, resultTypes, mask, /*passthru=*/Value(), maskableOp,
5917  maskRegionBuilder);
5918 }
5919 
5920 void MaskOp::build(
5921  OpBuilder &builder, OperationState &result, TypeRange resultTypes,
5922  Value mask, Value passthru, Operation *maskableOp,
5923  function_ref<void(OpBuilder &, Operation *)> maskRegionBuilder) {
5924  build(builder, result, mask, maskableOp, maskRegionBuilder);
5925  if (passthru)
5926  result.addOperands(passthru);
5927  result.addTypes(resultTypes);
5928 }
5929 
5931  // Create the op region.
5932  result.regions.reserve(1);
5933  Region &maskRegion = *result.addRegion();
5934 
5935  auto &builder = parser.getBuilder();
5936 
5937  // Parse all the operands.