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