MLIR  14.0.0git
Vectorization.cpp
Go to the documentation of this file.
1 //===- Vectorization.cpp - Implementation of linalg Vectorization ---------===//
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 the linalg dialect Vectorization transformations.
10 //
11 //===----------------------------------------------------------------------===//
12 
24 #include "mlir/IR/AffineExpr.h"
25 #include "mlir/IR/Matchers.h"
26 #include "mlir/IR/PatternMatch.h"
27 #include "mlir/Pass/Pass.h"
28 #include "mlir/Support/LLVM.h"
30 #include "llvm/ADT/ScopeExit.h"
31 #include "llvm/ADT/Sequence.h"
32 #include "llvm/ADT/SmallVector.h"
33 #include "llvm/ADT/TypeSwitch.h"
34 #include "llvm/Support/Debug.h"
35 #include "llvm/Support/raw_ostream.h"
36 #include <type_traits>
37 
38 using namespace mlir;
39 using namespace mlir::linalg;
40 
41 #define DEBUG_TYPE "linalg-vectorization"
42 
43 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
44 #define LDBG(X) LLVM_DEBUG(DBGS() << X)
45 
46 /// Try to vectorize `convOp` as a convolution.
48  LinalgOp convOp);
49 
50 /// Return the unique instance of OpType in `block` if it is indeed unique.
51 /// Return null if none or more than 1 instances exist.
52 template <typename OpType>
53 static OpType getSingleOpOfType(Block &block) {
54  OpType res;
55  block.walk([&](OpType op) {
56  if (res) {
57  res = nullptr;
58  return WalkResult::interrupt();
59  }
60  res = op;
61  return WalkResult::advance();
62  });
63  return res;
64 }
65 
66 /// Given an indexing `map` coming from a LinalgOp indexing, restricted to a
67 /// projectedPermutation, compress the unused dimensions to serve as a
68 /// permutation_map for a vector transfer operation.
69 /// For example, given a linalg op such as:
70 ///
71 /// ```
72 /// %0 = linalg.generic {
73 /// indexing_maps = affine_map<(d0, d1, d2, d3, d4) -> (d4, d0, d2)>,
74 /// indexing_maps = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3)>
75 /// }
76 /// ins(%0 : tensor<2x3x4xf32>)
77 /// outs(%1 : tensor<5x6xf32>)
78 /// ```
79 ///
80 /// the iteration domain size of the linalg op is 3x5x4x6x2. The first affine
81 /// map is reindexed to `affine_map<(d0, d1, d2) -> (d2, d0, d1)>`, the second
82 /// affine map is reindexed to `affine_map<(d0, d1) -> (d0, d1)>`.
84  assert(map.isProjectedPermutation(/*allowZeroInResults=*/true) &&
85  "expected projected permutation");
86  auto res = compressUnusedDims(map);
87  assert(res.getNumDims() == res.getNumResults() &&
88  "expected reindexed map with same number of dims and results");
89  return res;
90 }
91 
92 /// Helper data structure to represent the result of vectorization.
93 /// In certain specific cases, like terminators, we do not want to propagate/
95  /// Op failed to vectorize.
96  Failure = 0,
97  /// Op vectorized and custom function took care of replacement logic
99  /// Op vectorized into a new Op whose results will replace original Op's
100  /// results.
102  // TODO: support values if Op vectorized to Many-Ops whose results we need to
103  // aggregate for replacement.
104 };
106  /// Return status from vectorizing the current op.
108  /// New vectorized operation to replace the current op.
109  /// Replacement behavior is specified by `status`.
111 };
112 
115  using ::mlir::vector::CombiningKind;
116 
117  if (!combinerOp)
118  return llvm::None;
120  combinerOp)
121  .Case<arith::AddIOp, arith::AddFOp>(
122  [&](auto op) { return CombiningKind::ADD; })
123  .Case<arith::AndIOp>([&](auto op) { return CombiningKind::AND; })
124  .Case<arith::MaxSIOp>([&](auto op) { return CombiningKind::MAXSI; })
125  .Case<arith::MaxFOp>([&](auto op) { return CombiningKind::MAXF; })
126  .Case<arith::MinSIOp>([&](auto op) { return CombiningKind::MINSI; })
127  .Case<arith::MinFOp>([&](auto op) { return CombiningKind::MINF; })
128  .Case<arith::MulIOp, arith::MulFOp>(
129  [&](auto op) { return CombiningKind::MUL; })
130  .Case<arith::OrIOp>([&](auto op) { return CombiningKind::OR; })
131  .Case<arith::XOrIOp>([&](auto op) { return CombiningKind::XOR; })
132  .Default([&](auto op) { return llvm::None; });
133 }
134 
135 /// Check whether `outputOperand` is a reduction with a single combiner
136 /// operation. Return the combiner operation of the reduction. Return
137 /// nullptr otherwise. Multiple reduction operations would impose an
138 /// ordering between reduction dimensions and is currently unsupported in
139 /// Linalg. This limitation is motivated by the fact that e.g. min(max(X)) !=
140 /// max(min(X))
141 // TODO: use in LinalgOp verification, there is a circular dependency atm.
142 static Operation *matchLinalgReduction(OpOperand *outputOperand) {
143  auto linalgOp = cast<LinalgOp>(outputOperand->getOwner());
144  unsigned outputPos =
145  outputOperand->getOperandNumber() - linalgOp.getNumInputs();
146  // Only single combiner operations are supported for now.
147  SmallVector<Operation *, 4> combinerOps;
148  if (!matchReduction(linalgOp.getRegionOutputArgs(), outputPos, combinerOps) ||
149  combinerOps.size() != 1)
150  return nullptr;
151 
152  // Return the combiner operation.
153  return combinerOps[0];
154 }
155 
156 /// Broadcast `value` to a vector of `shape` if possible. Return value
157 /// otherwise.
159  ArrayRef<int64_t> shape) {
160  // If no shape to broadcast to, just return `value`.
161  if (shape.empty())
162  return value;
163  VectorType targetVectorType =
164  VectorType::get(shape, getElementTypeOrSelf(value));
165  if (vector::isBroadcastableTo(value.getType(), targetVectorType) !=
167  return value;
168  Location loc = b.getInsertionPoint()->getLoc();
169  return b.createOrFold<vector::BroadcastOp>(loc, targetVectorType, value);
170 }
171 
172 /// Create MultiDimReductionOp to compute the reduction for `reductionOp`. This
173 /// assumes that `reductionOp` has two operands and one of them is the reduction
174 /// initial value.
176  Value valueToReduce,
177  const SmallVector<bool> &reductionMask) {
178  auto maybeKind = getCombinerOpKind(reduceOp);
179  assert(maybeKind && "Failed precondition: could not get reduction kind");
180  return b.create<vector::MultiDimReductionOp>(
181  reduceOp->getLoc(), valueToReduce, reductionMask, *maybeKind);
182 }
183 
184 static SmallVector<bool> getReductionMask(LinalgOp linalgOp) {
185  unsigned idx = 0;
186  SmallVector<bool> reductionMask(linalgOp.iterator_types().size(), false);
187  for (auto attr : linalgOp.iterator_types()) {
188  if (isReductionIterator(attr))
189  reductionMask[idx] = true;
190  ++idx;
191  }
192  return reductionMask;
193 }
194 
195 /// Build a vector.transfer_write of `value` into `outputOperand` at indices set
196 /// to all `0`; where `outputOperand` is an output operand of the LinalgOp
197 /// currently being vectorized. If `dest` has null rank, build an memref.store.
198 /// Return the produced value or null if no value is produced.
200  OpOperand *outputOperand) {
201  Operation *write;
202  Location loc = value.getLoc();
203  auto linalgOp = cast<LinalgOp>(outputOperand->getOwner());
204  ArrayRef<int64_t> shape = linalgOp.getShape(outputOperand);
205  auto vectorType = VectorType::get(
206  shape, getElementTypeOrSelf(outputOperand->get().getType()));
207  if (vectorType.getRank() > 0) {
208  // 0-d case is still special: do not invert the reindexing map.
209  AffineMap map =
210  reindexIndexingMap(linalgOp.getTiedIndexingMap(outputOperand));
211  SmallVector<int64_t> transposeShape =
213  assert(!transposeShape.empty() && "unexpected empty transpose shape");
214  vectorType = VectorType::get(transposeShape, vectorType.getElementType());
215  SmallVector<Value> indices(linalgOp.getRank(outputOperand),
216  b.create<arith::ConstantIndexOp>(loc, 0));
217  value = broadcastIfNeeded(b, value, vectorType.getShape());
218  write = b.create<vector::TransferWriteOp>(loc, value, outputOperand->get(),
219  indices, map);
220  } else {
221  if (!value.getType().isa<VectorType>())
222  value = b.create<vector::BroadcastOp>(loc, vectorType, value);
223  assert(value.getType() == vectorType && "incorrect type");
224  write = b.create<vector::TransferWriteOp>(loc, value, outputOperand->get(),
225  ValueRange{});
226  }
227  LDBG("vectorized op: " << *write);
228  if (!write->getResults().empty())
229  return write->getResult(0);
230  return Value();
231 }
232 
233 // Custom vectorization function type. Produce a vector form of Operation*
234 // assuming all its vectorized operands are already in the BlockAndValueMapping.
235 // Return nullptr if the Operation cannot be vectorized.
236 using CustomVectorizationHook = std::function<VectorizationResult(
238 
239 /// Helper function to vectorize the terminator of a `linalgOp`. New result
240 /// vector values are appended to `newResults`. Return
241 /// VectorizationStatus::NoReplace to signal the vectorization algorithm that it
242 /// should not try to map produced operations and instead return the results
243 /// using the `newResults` vector making them available to the
244 /// vectorization algorithm for RAUW. This function is meant to be used as a
245 /// CustomVectorizationHook.
246 static VectorizationResult
248  const BlockAndValueMapping &bvm, LinalgOp linalgOp,
249  SmallVectorImpl<Value> &newResults) {
250  auto yieldOp = dyn_cast<linalg::YieldOp>(op);
251  if (!yieldOp)
253  for (const auto &outputs : llvm::enumerate(yieldOp.values())) {
254  // TODO: Scan for an opportunity for reuse.
255  // TODO: use a map.
256  Value vectorValue = bvm.lookup(outputs.value());
257  Value newResult = buildVectorWrite(
258  b, vectorValue, linalgOp.getOutputOperand(outputs.index()));
259  if (newResult)
260  newResults.push_back(newResult);
261  }
263 }
264 
265 /// Helper function to vectorize the index operations of a `linalgOp`. Return
266 /// VectorizationStatus::NewOp to signal the vectorization algorithm that it
267 /// should map the produced operations. This function is meant to be used as a
268 /// CustomVectorizationHook.
270  LinalgOp linalgOp) {
271  IndexOp indexOp = dyn_cast<linalg::IndexOp>(op);
272  if (!indexOp)
274  auto loc = indexOp.getLoc();
275  // Compute the static loop sizes of the index op.
276  auto targetShape = linalgOp.computeStaticLoopSizes();
277  // Compute a one-dimensional index vector for the index op dimension.
278  SmallVector<int64_t> constantSeq =
279  llvm::to_vector<16>(llvm::seq<int64_t>(0, targetShape[indexOp.dim()]));
280  auto constantOp =
281  b.create<arith::ConstantOp>(loc, b.getIndexVectorAttr(constantSeq));
282  // Return the one-dimensional index vector if it lives in the trailing
283  // dimension of the iteration space since the vectorization algorithm in this
284  // case can handle the broadcast.
285  if (indexOp.dim() == targetShape.size() - 1)
287  // Otherwise permute the targetShape to move the index dimension last,
288  // broadcast the one-dimensional index vector to the permuted shape, and
289  // finally transpose the broadcasted index vector to undo the permutation.
290  std::swap(targetShape[indexOp.dim()], targetShape.back());
291  auto broadCastOp = b.create<vector::BroadcastOp>(
292  loc, VectorType::get(targetShape, b.getIndexType()), constantOp);
293  SmallVector<int64_t> transposition =
294  llvm::to_vector<16>(llvm::seq<int64_t>(0, linalgOp.getNumLoops()));
295  std::swap(transposition.back(), transposition[indexOp.dim()]);
296  auto transposeOp =
297  b.create<vector::TransposeOp>(loc, broadCastOp, transposition);
298  return VectorizationResult{VectorizationStatus::NewOp, transposeOp};
299 }
300 
301 /// Create a new vectorized verstion of `op` with the given operands and types.
303  ValueRange newOperands,
304  ArrayRef<Type> types) {
305  OperationState state(op->getLoc(), op->getName());
306  state.addAttributes(op->getAttrs());
307  state.addOperands(newOperands);
308  state.addTypes(types);
309  return b.createOperation(state);
310 }
311 
312 /// Emit reduction operations if the shapes of the value to reduce is different
313 /// that the result shape.
314 static Operation *reduceIfNeeded(OpBuilder &b, LinalgOp linalgOp, Operation *op,
315  Value reduceValue, Value initialValue,
316  const BlockAndValueMapping &bvm) {
317  Value reduceVec = bvm.lookup(reduceValue);
318  Value outputVec = bvm.lookup(initialValue);
319  auto reduceType = reduceVec.getType().dyn_cast<VectorType>();
320  auto outputType = outputVec.getType().dyn_cast<VectorType>();
321  // Reduce only if needed as the value may already have been reduce for
322  // contraction vectorization.
323  if (!reduceType ||
324  (outputType && reduceType.getShape() == outputType.getShape()))
325  return nullptr;
326  SmallVector<bool> reductionMask = getReductionMask(linalgOp);
327  Value reduce = buildMultiDimReduce(b, op, reduceVec, reductionMask);
328  return createVectorizedOp(b, op, {reduce, outputVec}, reduce.getType());
329 }
330 
331 /// Generic vectorization for a single operation `op`, given already vectorized
332 /// operands carried by `bvm`. Vectorization occurs as follows:
333 /// 1. Try to apply any of the `customVectorizationHooks` and return its
334 /// result on success.
335 /// 2. Clone any constant in the current scope without vectorization: each
336 /// consumer of the constant will later determine the shape to which the
337 /// constant needs to be broadcast to.
338 /// 3. Fail on any remaining non `ElementwiseMappable` op. It is the purpose
339 /// of the `customVectorizationHooks` to cover such cases.
340 /// 4. Clone `op` in vector form to a vector of shape prescribed by the first
341 /// operand of maximal rank. Other operands have smaller rank and are
342 /// broadcast accordingly. It is assumed this broadcast is always legal,
343 /// otherwise, it means one of the `customVectorizationHooks` is incorrect.
344 ///
345 /// This function assumes all operands of `op` have been vectorized and are in
346 /// the `bvm` mapping. As a consequence, this function is meant to be called on
347 /// a topologically-sorted list of ops.
348 /// This function does not update `bvm` but returns a VectorizationStatus that
349 /// instructs the caller what `bvm` update needs to occur.
350 static VectorizationResult
351 vectorizeOneOp(OpBuilder &b, LinalgOp linalgOp, Operation *op,
352  const BlockAndValueMapping &bvm,
353  ArrayRef<CustomVectorizationHook> customVectorizationHooks) {
354  LDBG("vectorize op " << *op);
355 
356  // 1. Try to apply any CustomVectorizationHook.
357  if (!customVectorizationHooks.empty()) {
358  for (auto &customFunc : customVectorizationHooks) {
359  VectorizationResult result = customFunc(op, bvm);
360  if (result.status == VectorizationStatus::Failure)
361  continue;
362  return result;
363  }
364  }
365 
366  // 2. Constant ops don't get vectorized but rather broadcasted at their users.
367  // Clone so that the constant is not confined to the linalgOp block .
368  if (isa<arith::ConstantOp, ConstantOp>(op))
370 
371  // 3. Only ElementwiseMappable are allowed in the generic vectorization.
374 
375  // 4 . Check if the operation is a reduction.
376  SmallVector<std::pair<Value, Value>> reductionOperands;
377  for (Value operand : op->getOperands()) {
378  auto arg = operand.dyn_cast<BlockArgument>();
379  if (!arg || arg.getArgNumber() < linalgOp.getNumInputs())
380  continue;
381  SmallVector<Operation *> reductionOps;
382  Value reduceValue = matchReduction(
383  linalgOp.getRegionOutputArgs(),
384  arg.getArgNumber() - linalgOp.getNumInputs(), reductionOps);
385  if (!reduceValue)
386  continue;
387  reductionOperands.push_back(std::make_pair(reduceValue, operand));
388  }
389  if (!reductionOperands.empty()) {
390  assert(reductionOperands.size() == 1);
391  Operation *reduceOp =
392  reduceIfNeeded(b, linalgOp, op, reductionOperands[0].first,
393  reductionOperands[0].second, bvm);
394  if (reduceOp)
396  }
397 
398  // 5. Generic vectorization path for ElementwiseMappable ops.
399  // a. first get the first max ranked shape.
400  SmallVector<int64_t, 4> firstMaxRankedShape;
401  for (Value operand : op->getOperands()) {
402  auto vt = bvm.lookup(operand).getType().dyn_cast<VectorType>();
403  if (vt && firstMaxRankedShape.size() < vt.getShape().size())
404  firstMaxRankedShape.assign(vt.getShape().begin(), vt.getShape().end());
405  }
406  // b. broadcast each op if needed.
407  auto vectorizedOperands = llvm::map_range(op->getOperands(), [&](Value v) {
408  return firstMaxRankedShape.empty()
409  ? bvm.lookup(v)
410  : broadcastIfNeeded(b, bvm.lookup(v), firstMaxRankedShape);
411  });
412  // c. for elementwise, the result is the vector with the firstMaxRankedShape
413  auto returnTypes = llvm::map_range(op->getResultTypes(), [&](Type t) {
414  return firstMaxRankedShape.empty()
415  ? t
416  : VectorType::get(firstMaxRankedShape, t);
417  });
418 
419  // Build and return the new op.
420  return VectorizationResult{
422  createVectorizedOp(b, op, llvm::to_vector<4>(vectorizedOperands),
423  llvm::to_vector<4>(returnTypes))};
424 }
425 
426 /// Detect whether `r` has only ConstantOp, ElementwiseMappable and YieldOp.
428  if (!llvm::hasSingleElement(r))
429  return false;
430  for (Operation &op : r.front()) {
431  if (!(isa<arith::ConstantOp, ConstantOp, linalg::YieldOp, linalg::IndexOp>(
432  op) ||
434  llvm::any_of(op.getResultTypes(),
435  [](Type type) { return !type.isIntOrIndexOrFloat(); }))
436  return false;
437  }
438  return true;
439 }
440 
441 // Return true if the op is an element-wise linalg op.
442 static bool isElementwise(Operation *op) {
443  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
444  if (!linalgOp)
445  return false;
446  if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops())
447  return false;
448  // TODO: relax the restrictions on indexing map.
449  for (OpOperand *opOperand : linalgOp.getOutputOperands()) {
450  if (!linalgOp.getTiedIndexingMap(opOperand).isIdentity())
451  return false;
452  }
453  return hasOnlyScalarElementwiseOp(linalgOp->getRegion(0));
454 }
455 
456 /// Generic vectorization function that rewrites the body of a `linalgOp` into
457 /// vector form. Generic vectorization proceeds as follows:
458 /// 1. Verify the `linalgOp` has one non-empty region.
459 /// 2. Values defined above the region are mapped to themselves and will be
460 /// broadcasted on a per-need basis by their consumers.
461 /// 3. Each region argument is vectorized into a vector.transfer_read (or 0-d
462 /// load).
463 /// TODO: Reuse opportunities for RAR dependencies.
464 /// 4a. Register CustomVectorizationHook for YieldOp to capture the results.
465 /// 4b. Register CustomVectorizationHook for IndexOp to access the iteration
466 /// indices.
467 /// 5. Iteratively call vectorizeOneOp on the region operations.
468 ///
469 /// When `broadcastToMaximalCommonShape` is set to true, eager broadcasting is
470 /// performed to the maximal common vector size implied by the `linalgOp`
471 /// iteration space. This eager broadcasting is introduced in the
472 /// permutation_map of the vector.transfer_read operations. The eager
473 /// broadcasting makes it trivial to detrmine where broadcast, transposes and
474 /// reductions should occur, without any bookkeeping. The tradeoff is that, in
475 /// the absence of good canonicalizations, the amount of work increases.
476 /// This is not deemed a problem as we expect canonicalizations and foldings to
477 /// aggressively clean up the useless work.
478 static LogicalResult
479 vectorizeAsLinalgGeneric(OpBuilder &b, LinalgOp linalgOp,
480  SmallVectorImpl<Value> &newResults) {
481  Block *block = linalgOp.getBlock();
482 
483  // 2. Values defined above the region can only be broadcast for now. Make them
484  // map to themselves.
486  SetVector<Value> valuesSet;
487  mlir::getUsedValuesDefinedAbove(linalgOp->getRegion(0), valuesSet);
488  bvm.map(valuesSet.getArrayRef(), valuesSet.getArrayRef());
489 
490  if (linalgOp.getNumOutputs() == 0)
491  return failure();
492 
493  // TODO: the common vector shape is equal to the static loop sizes only when
494  // all indexing maps are projected permutations. For convs and stencils the
495  // logic will need to evolve.
496  SmallVector<int64_t> commonVectorShape = linalgOp.computeStaticLoopSizes();
497 
498  // 3. Turn all BBArgs into vector.transfer_read / load.
499  Location loc = linalgOp.getLoc();
500  Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
501  for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
502  BlockArgument bbarg = block->getArgument(opOperand->getOperandNumber());
503  if (linalgOp.isScalar(opOperand)) {
504  bvm.map(bbarg, opOperand->get());
505  continue;
506  }
507  VectorType readType;
508  AffineMap map;
509  // TODO: can we keep this simplification?
510  // if (linalgOp.getShape(opOperand).empty()) {
511  // readType = VectorType::get({}, bbarg.getType());
512  // } else {
513  if (opOperand->getOperandNumber() < linalgOp.getNumInputs()) {
515  linalgOp.getTiedIndexingMap(opOperand));
516  readType = VectorType::get(commonVectorShape,
517  getElementTypeOrSelf(opOperand->get()));
518  } else {
519  map = inversePermutation(
520  reindexIndexingMap(linalgOp.getTiedIndexingMap(opOperand)));
521  readType = VectorType::get(map.compose(linalgOp.getShape(opOperand)),
522  getElementTypeOrSelf(opOperand->get()));
523  }
524  // }
525 
526  auto shape = linalgOp.getShape(opOperand);
527  SmallVector<Value> indices(shape.size(), zero);
528  Value readValue = b.create<vector::TransferReadOp>(
529  loc, readType, opOperand->get(), indices, map);
530  // Not all ops support 0-d vectors, extract the scalar for now.
531  // TODO: remove this.
532  if (readValue.getType().cast<VectorType>().getRank() == 0)
533  readValue = b.create<vector::ExtractElementOp>(loc, readValue);
534 
535  LDBG("new vectorized bbarg(" << bbarg.getArgNumber() << "): " << readValue);
536  bvm.map(bbarg, readValue);
537  bvm.map(opOperand->get(), readValue);
538  }
539 
541  // 4a. Register CustomVectorizationHook for yieldOp.
542  CustomVectorizationHook vectorizeYield =
543  [&](Operation *op,
545  return vectorizeLinalgYield(b, op, bvm, linalgOp, newResults);
546  };
547  hooks.push_back(vectorizeYield);
548 
549  // 4b. Register CustomVectorizationHook for indexOp.
550  CustomVectorizationHook vectorizeIndex =
551  [&](Operation *op,
553  return vectorizeLinalgIndex(b, op, linalgOp);
554  };
555  hooks.push_back(vectorizeIndex);
556 
557  // 5. Iteratively call `vectorizeOneOp` to each op in the slice.
558  for (Operation &op : block->getOperations()) {
559  VectorizationResult result = vectorizeOneOp(b, linalgOp, &op, bvm, hooks);
560  if (result.status == VectorizationStatus::Failure) {
561  LDBG("failed to vectorize: " << op);
562  return failure();
563  }
564  if (result.status == VectorizationStatus::NewOp) {
565  LDBG("new vector op: " << *result.newOp;);
566  bvm.map(op.getResults(), result.newOp->getResults());
567  }
568  }
569 
570  return success();
571 }
572 
573 /// Helper function to vectorize a `linalgOp` with contraction semantics in a
574 /// generic fashion.
575 /// This helper is needed atm because the truly generic implementation requires
576 /// good vector.multi_reduce folding patterns that are currently NYI.
577 // TODO: drop reliance on a specific pattern.
578 static bool allIndexingsAreProjectedPermutation(LinalgOp op) {
579  return llvm::all_of(op.getIndexingMaps(), [](AffineMap m) {
580  return m.isProjectedPermutation(/*allowZeroInResults=*/true);
581  });
582 }
583 
584 // TODO: probably need some extra checks for reduction followed by consumer
585 // ops that may not commute (e.g. linear reduction + non-linear instructions).
587  if (llvm::none_of(op.iterator_types(), isReductionIterator)) {
588  LDBG("reduction precondition failed: no reduction iterator");
589  return failure();
590  }
591  for (OpOperand *opOperand : op.getOutputOperands()) {
592  Operation *reduceOp = matchLinalgReduction(opOperand);
593  if (!reduceOp || !getCombinerOpKind(reduceOp)) {
594  LDBG("reduction precondition failed: reduction detection failed");
595  return failure();
596  }
597  }
598  return success();
599 }
600 
602  if (isElementwise(op))
603  return success();
604  // TODO: isaConvolutionOpInterface that can also infer from generic features.
605  // But we will still need stride/dilation attributes that will be annoying to
606  // reverse-engineer...
607  if (isa<ConvolutionOpInterface>(op.getOperation()))
608  return success();
609  // TODO: the common vector shape is equal to the static loop sizes only when
610  // all indexing maps are projected permutations. For convs and stencils the
611  // logic will need to evolve.
613  LDBG("precondition failed: not projected permutations");
614  return failure();
615  }
616  if (failed(reductionPreconditions(op))) {
617  LDBG("precondition failed: reduction preconditions");
618  return failure();
619  }
620  return success();
621 }
622 
623 static LogicalResult vectorizeLinalgOpPrecondition(LinalgOp linalgOp) {
624  // All types must be static shape to go to vector.
625  if (linalgOp.hasDynamicShape()) {
626  LDBG("precondition failed: dynamic shape");
627  return failure();
628  }
629  return vectorizeStaticLinalgOpPrecondition(linalgOp);
630 }
631 
633  LinalgOp linalgOp) {
634  if (failed(vectorizeLinalgOpPrecondition(linalgOp)))
635  return failure();
636 
637  SmallVector<Value> results;
638  // TODO: isaConvolutionOpInterface that can also infer from generic
639  // features. Will require stride/dilation attributes inference.
640  FailureOr<Operation *> convOr = vectorizeConvolution(rewriter, linalgOp);
641  if (succeeded(convOr)) {
642  llvm::append_range(results, (*convOr)->getResults());
643  } else {
644  if (failed(vectorizeLinalgOpPrecondition(linalgOp)))
645  return failure();
646  LDBG("Vectorize generic by broadcasting to a common shape: " << linalgOp);
647  if (failed(vectorizeAsLinalgGeneric(rewriter, linalgOp, results)))
648  return failure();
649  }
650 
651  if (!results.empty())
652  rewriter.replaceOp(linalgOp, results);
653  else
654  rewriter.eraseOp(linalgOp);
655 
656  return success();
657 }
658 
659 //----------------------------------------------------------------------------//
660 // Misc. vectorization patterns.
661 //----------------------------------------------------------------------------//
662 
663 /// Helper function that retrieves the value of an IntegerAttr.
664 static int64_t getIntFromAttr(Attribute attr) {
665  return attr.cast<IntegerAttr>().getInt();
666 }
667 
668 /// Given an ArrayRef of OpFoldResults, return a vector of Values.
669 /// IntegerAttrs are converted to ConstantIndexOps. Other attribute types are
670 /// not supported.
672  ArrayRef<OpFoldResult> ofrs) {
673  SmallVector<Value> result;
674  llvm::for_each(ofrs, [&](auto o) {
675  if (auto val = o.template dyn_cast<Value>()) {
676  result.push_back(val);
677  } else {
678  result.push_back(builder.create<arith::ConstantIndexOp>(
679  loc, getIntFromAttr(o.template get<Attribute>())));
680  }
681  });
682  return result;
683 }
684 
685 /// Rewrite a tensor::PadOp into a sequence of InitTensorOp, FillOp and
686 /// InsertSliceOp. For now, only constant padding values are supported.
687 /// If there is enough static type information, TransferReadOps and
688 /// TransferWriteOps may be generated instead of InsertSliceOps.
691  PatternBenefit benefit = 1)
692  : GeneralizePadOpPattern(context, tryVectorizeCopy, benefit) {}
693  /// Vectorize the copying of a tensor::PadOp's source. This is possible if
694  /// each dimension size is statically know in the source type or the result
695  /// type (or both).
697  tensor::PadOp padOp, Value dest) {
698  auto sourceType = padOp.getSourceType();
699  auto resultType = padOp.getResultType();
700 
701  // Copy cannot be vectorized if pad value is non-constant and source shape
702  // is dynamic. In case of a dynamic source shape, padding must be appended
703  // by TransferReadOp, but TransferReadOp supports only constant padding.
704  auto padValue = padOp.getConstantPaddingValue();
705  if (!padValue) {
706  if (!sourceType.hasStaticShape())
707  return failure();
708  // Create dummy padding value.
709  auto elemType = sourceType.getElementType();
710  padValue = rewriter.create<arith::ConstantOp>(
711  padOp.getLoc(), elemType, rewriter.getZeroAttr(elemType));
712  }
713 
714  SmallVector<int64_t> vecShape;
715  SmallVector<bool> readInBounds;
716  SmallVector<bool> writeInBounds;
717  for (unsigned i = 0; i < sourceType.getRank(); ++i) {
718  if (!sourceType.isDynamicDim(i)) {
719  vecShape.push_back(sourceType.getDimSize(i));
720  // Source shape is statically known: Neither read nor write are
721  // out-of- bounds.
722  readInBounds.push_back(true);
723  writeInBounds.push_back(true);
724  } else if (!resultType.isDynamicDim(i)) {
725  // Source shape is not statically known, but result shape is.
726  // Vectorize with size of result shape. This may be larger than the
727  // source size.
728  vecShape.push_back(resultType.getDimSize(i));
729  // Read may be out-of-bounds because the result size could be larger
730  // than the source size.
731  readInBounds.push_back(false);
732  // Write is out-of-bounds if low padding > 0.
733  writeInBounds.push_back(
734  getConstantIntValue(padOp.getMixedLowPad()[i]) ==
735  static_cast<int64_t>(0));
736  } else {
737  // Neither source nor result dim of padOp is static. Cannot vectorize
738  // the copy.
739  return failure();
740  }
741  }
742  auto vecType = VectorType::get(vecShape, sourceType.getElementType());
743 
744  // Generate TransferReadOp.
745  SmallVector<Value> readIndices(
746  vecType.getRank(),
747  rewriter.create<arith::ConstantIndexOp>(padOp.getLoc(), 0));
748  auto read = rewriter.create<vector::TransferReadOp>(
749  padOp.getLoc(), vecType, padOp.source(), readIndices, padValue,
750  ArrayRef<bool>{readInBounds});
751 
752  // If `dest` is a FillOp and the TransferWriteOp would overwrite the
753  // entire tensor, write directly to the FillOp's operand.
754  if (llvm::equal(vecShape, resultType.getShape()) &&
755  llvm::all_of(writeInBounds, [](bool b) { return b; }))
756  if (auto fill = dest.getDefiningOp<FillOp>())
757  dest = fill.output();
758 
759  // Generate TransferWriteOp.
760  auto writeIndices =
761  ofrToIndexValues(rewriter, padOp.getLoc(), padOp.getMixedLowPad());
762  rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
763  padOp, read, dest, writeIndices, ArrayRef<bool>{writeInBounds});
764 
765  return success();
766  }
767 };
768 
769 /// Base pattern for rewriting tensor::PadOps whose result is consumed by a
770 /// given operation type OpTy.
771 template <typename OpTy>
772 struct VectorizePadOpUserPattern : public OpRewritePattern<tensor::PadOp> {
774 
775  LogicalResult matchAndRewrite(tensor::PadOp padOp,
776  PatternRewriter &rewriter) const final {
777  bool changed = false;
778  // Insert users in vector, because some users may be replaced/removed.
779  for (auto *user : llvm::to_vector<4>(padOp->getUsers()))
780  if (auto op = dyn_cast<OpTy>(user))
781  changed |= rewriteUser(rewriter, padOp, op).succeeded();
782  return success(changed);
783  }
784 
785 protected:
786  virtual LogicalResult rewriteUser(PatternRewriter &rewriter,
787  tensor::PadOp padOp, OpTy op) const = 0;
788 };
789 
790 /// Rewrite use of tensor::PadOp result in TransferReadOp. E.g.:
791 /// ```
792 /// %0 = linalg.pad_tensor %src ... : tensor<?x?xf32> to tensor<17x5xf32>
793 /// %r = vector.transfer_read %0[%c0, %c0], %cst
794 /// {in_bounds = [true, true]} : tensor<17x5xf32>, vector<17x5xf32>
795 /// ```
796 /// is rewritten to:
797 /// ```
798 /// %r = vector.transfer_read %src[%c0, %c0], %padding
799 /// {in_bounds = [true, true]}
800 /// : tensor<?x?xf32>, vector<17x5xf32>
801 /// ```
802 /// Note: By restricting this pattern to in-bounds TransferReadOps, we can be
803 /// sure that the original padding value %cst was never used.
804 ///
805 /// This rewrite is possible if:
806 /// - `xferOp` has no out-of-bounds dims or mask.
807 /// - Low padding is static 0.
808 /// - Single, scalar padding value.
810  : public VectorizePadOpUserPattern<vector::TransferReadOp> {
812  vector::TransferReadOp>::VectorizePadOpUserPattern;
813 
814  LogicalResult rewriteUser(PatternRewriter &rewriter, tensor::PadOp padOp,
815  vector::TransferReadOp xferOp) const override {
816  // Low padding must be static 0.
817  if (!padOp.hasZeroLowPad())
818  return failure();
819  // Pad value must be a constant.
820  auto padValue = padOp.getConstantPaddingValue();
821  if (!padValue)
822  return failure();
823  // Padding value of existing `xferOp` is unused.
824  if (xferOp.hasOutOfBoundsDim() || xferOp.mask())
825  return failure();
826 
827  rewriter.updateRootInPlace(xferOp, [&]() {
828  SmallVector<bool> inBounds(xferOp.getVectorType().getRank(), false);
829  xferOp->setAttr(xferOp.getInBoundsAttrName(),
830  rewriter.getBoolArrayAttr(inBounds));
831  xferOp.sourceMutable().assign(padOp.source());
832  xferOp.paddingMutable().assign(padValue);
833  });
834 
835  return success();
836  }
837 };
838 
839 /// Rewrite use of tensor::PadOp result in TransferWriteOp.
840 /// This pattern rewrites TransferWriteOps that write to a padded tensor
841 /// value, where the same amount of padding is immediately removed again after
842 /// the write. In such cases, the TransferWriteOp can write to the non-padded
843 /// tensor value and apply out-of-bounds masking. E.g.:
844 /// ```
845 /// %0 = tensor.extract_slice ...[...] [%s0, %s1] [1, 1]
846 /// : tensor<...> to tensor<?x?xf32>
847 /// %1 = linalg.pad_tensor %0 ... : tensor<?x?xf32> to tensor<17x5xf32>
848 /// %2 = vector.transfer_write %vec, %1[...]
849 /// : vector<17x5xf32>, tensor<17x5xf32>
850 /// %r = tensor.extract_slice %2[0, 0] [%s0, %s1] [1, 1]
851 /// : tensor<17x5xf32> to tensor<?x?xf32>
852 /// ```
853 /// is rewritten to:
854 /// ```
855 /// %0 = tensor.extract_slice ...[...] [%s0, %s1] [1, 1]
856 /// : tensor<...> to tensor<?x?xf32>
857 /// %r = vector.transfer_write %vec, %0[...] : vector<17x5xf32>,
858 /// tensor<?x?xf32>
859 /// ```
860 /// Note: It is important that the ExtractSliceOp %r resizes the result of the
861 /// TransferWriteOp to the same size as the input of the TensorPadOp (or an
862 /// even smaller size). Otherwise, %r's new (dynamic) dimensions would differ
863 /// from %r's old dimensions.
864 ///
865 /// This rewrite is possible if:
866 /// - Low padding is static 0.
867 /// - `xferOp` has exactly one use, which is an ExtractSliceOp. This
868 /// ExtractSliceOp trims the same amount of padding that was added
869 /// beforehand.
870 /// - Single, scalar padding value.
872  : public VectorizePadOpUserPattern<vector::TransferWriteOp> {
874  vector::TransferWriteOp>::VectorizePadOpUserPattern;
875 
876  LogicalResult rewriteUser(PatternRewriter &rewriter, tensor::PadOp padOp,
877  vector::TransferWriteOp xferOp) const override {
878  // TODO: support 0-d corner case.
879  if (xferOp.getTransferRank() == 0)
880  return failure();
881 
882  // Low padding must be static 0.
883  if (!padOp.hasZeroLowPad())
884  return failure();
885  // Pad value must be a constant.
886  auto padValue = padOp.getConstantPaddingValue();
887  if (!padValue)
888  return failure();
889  // TransferWriteOp result must be directly consumed by an ExtractSliceOp.
890  if (!xferOp->hasOneUse())
891  return failure();
892  auto trimPadding = dyn_cast<tensor::ExtractSliceOp>(*xferOp->user_begin());
893  if (!trimPadding)
894  return failure();
895  // Only static zero offsets supported when trimming padding.
896  if (!trimPadding.hasZeroOffset())
897  return failure();
898  // trimPadding must remove the amount of padding that was added earlier.
899  if (!hasSameTensorSize(padOp.source(), trimPadding))
900  return failure();
901 
902  // Insert the new TransferWriteOp at position of the old TransferWriteOp.
903  rewriter.setInsertionPoint(xferOp);
904 
905  SmallVector<bool> inBounds(xferOp.getVectorType().getRank(), false);
906  auto newXferOp = rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
907  xferOp, padOp.source().getType(), xferOp.vector(), padOp.source(),
908  xferOp.indices(), xferOp.permutation_mapAttr(), xferOp.mask(),
909  rewriter.getBoolArrayAttr(inBounds));
910  rewriter.replaceOp(trimPadding, newXferOp->getResult(0));
911 
912  return success();
913  }
914 
915  /// Check if `beforePadding` and `afterTrimming` have the same tensor size,
916  /// i.e., same dimensions.
917  ///
918  /// Dimensions may be static, dynamic or mix of both. In case of dynamic
919  /// dimensions, this function tries to infer the (static) tensor size by
920  /// looking at the defining op and utilizing op-specific knowledge.
921  ///
922  /// This is a conservative analysis. In case equal tensor sizes cannot be
923  /// proven statically, this analysis returns `false` even though the tensor
924  /// sizes may turn out to be equal at runtime.
925  bool hasSameTensorSize(Value beforePadding,
926  tensor::ExtractSliceOp afterTrimming) const {
927  // If the input to tensor::PadOp is a CastOp, try with with both CastOp
928  // result and CastOp operand.
929  if (auto castOp = beforePadding.getDefiningOp<tensor::CastOp>())
930  if (hasSameTensorSize(castOp.source(), afterTrimming))
931  return true;
932 
933  auto t1 = beforePadding.getType().dyn_cast<RankedTensorType>();
934  auto t2 = afterTrimming.getType().dyn_cast<RankedTensorType>();
935  // Only RankedTensorType supported.
936  if (!t1 || !t2)
937  return false;
938  // Rank of both values must be the same.
939  if (t1.getRank() != t2.getRank())
940  return false;
941 
942  // All static dimensions must be the same. Mixed cases (e.g., dimension
943  // static in `t1` but dynamic in `t2`) are not supported.
944  for (unsigned i = 0; i < t1.getRank(); ++i) {
945  if (t1.isDynamicDim(i) != t2.isDynamicDim(i))
946  return false;
947  if (!t1.isDynamicDim(i) && t1.getDimSize(i) != t2.getDimSize(i))
948  return false;
949  }
950 
951  // Nothing more to check if all dimensions are static.
952  if (t1.getNumDynamicDims() == 0)
953  return true;
954 
955  // All dynamic sizes must be the same. The only supported case at the
956  // moment is when `beforePadding` is an ExtractSliceOp (or a cast
957  // thereof).
958 
959  // Apart from CastOp, only ExtractSliceOp is supported.
960  auto beforeSlice = beforePadding.getDefiningOp<tensor::ExtractSliceOp>();
961  if (!beforeSlice)
962  return false;
963 
964  assert(static_cast<size_t>(t1.getRank()) ==
965  beforeSlice.getMixedSizes().size());
966  assert(static_cast<size_t>(t2.getRank()) ==
967  afterTrimming.getMixedSizes().size());
968 
969  for (unsigned i = 0; i < t1.getRank(); ++i) {
970  // Skip static dimensions.
971  if (!t1.isDynamicDim(i))
972  continue;
973  auto size1 = beforeSlice.getMixedSizes()[i];
974  auto size2 = afterTrimming.getMixedSizes()[i];
975 
976  // Case 1: Same value or same constant int.
977  if (isEqualConstantIntOrValue(size1, size2))
978  continue;
979 
980  // Other cases: Take a deeper look at defining ops of values.
981  auto v1 = size1.dyn_cast<Value>();
982  auto v2 = size2.dyn_cast<Value>();
983  if (!v1 || !v2)
984  return false;
985 
986  // Case 2: Both values are identical AffineMinOps. (Should not happen if
987  // CSE is run.)
988  auto minOp1 = v1.getDefiningOp<AffineMinOp>();
989  auto minOp2 = v2.getDefiningOp<AffineMinOp>();
990  if (minOp1 && minOp2 && minOp1.getAffineMap() == minOp2.getAffineMap() &&
991  minOp1.operands() == minOp2.operands())
992  continue;
993 
994  // Add additional cases as needed.
995  }
996 
997  // All tests passed.
998  return true;
999  }
1000 };
1001 
1002 /// Rewrite use of tensor::PadOp result in InsertSliceOp. E.g.:
1003 /// ```
1004 /// %0 = linalg.pad_tensor %src ... : tensor<?x?xf32> to tensor<17x5xf32>
1005 /// %r = tensor.insert_slice %0
1006 /// into %dest[%a, %b, 0, 0] [1, 1, 17, 5] [1, 1, 1, 1]
1007 /// : tensor<17x5xf32> into tensor<?x?x17x5xf32>
1008 /// ```
1009 /// is rewritten to:
1010 /// ```
1011 /// %0 = vector.transfer_read %src[%c0, %c0], %padding
1012 /// : tensor<?x?xf32>, vector<17x5xf32>
1013 /// %r = vector.transfer_write %0, %dest[%a, %b, %c0, %c0]
1014 /// {in_bounds = [true, true]} : vector<17x5xf32>, tensor<?x?x17x5xf32>
1015 /// ```
1016 ///
1017 /// This rewrite is possible if:
1018 /// - Low padding is static 0.
1019 /// - `padOp` result shape is static.
1020 /// - The entire padded tensor is inserted.
1021 /// (Implies that sizes of `insertOp` are all static.)
1022 /// - Only unit strides in `insertOp`.
1023 /// - Single, scalar padding value.
1024 /// - `padOp` result not used as destination.
1026  : public VectorizePadOpUserPattern<tensor::InsertSliceOp> {
1028  tensor::InsertSliceOp>::VectorizePadOpUserPattern;
1029 
1030  LogicalResult rewriteUser(PatternRewriter &rewriter, tensor::PadOp padOp,
1031  tensor::InsertSliceOp insertOp) const override {
1032  // Low padding must be static 0.
1033  if (!padOp.hasZeroLowPad())
1034  return failure();
1035  // Only unit stride supported.
1036  if (!insertOp.hasUnitStride())
1037  return failure();
1038  // Pad value must be a constant.
1039  auto padValue = padOp.getConstantPaddingValue();
1040  if (!padValue)
1041  return failure();
1042  // Dynamic shapes not supported.
1043  if (!padOp.result().getType().cast<ShapedType>().hasStaticShape())
1044  return failure();
1045  // Pad result not used as destination.
1046  if (insertOp.dest() == padOp.result())
1047  return failure();
1048 
1049  auto vecType = VectorType::get(padOp.getType().getShape(),
1050  padOp.getType().getElementType());
1051  unsigned vecRank = vecType.getRank();
1052  unsigned tensorRank = insertOp.getType().getRank();
1053 
1054  // Check if sizes match: Insert the entire tensor into most minor dims.
1055  // (No permutations allowed.)
1056  SmallVector<int64_t> expectedSizes(tensorRank - vecRank, 1);
1057  expectedSizes.append(vecType.getShape().begin(), vecType.getShape().end());
1058  if (!llvm::all_of(
1059  llvm::zip(insertOp.getMixedSizes(), expectedSizes), [](auto it) {
1060  return getConstantIntValue(std::get<0>(it)) == std::get<1>(it);
1061  }))
1062  return failure();
1063 
1064  // Insert the TransferReadOp and TransferWriteOp at the position of the
1065  // InsertSliceOp.
1066  rewriter.setInsertionPoint(insertOp);
1067 
1068  // Generate TransferReadOp: Read entire source tensor and add high
1069  // padding.
1070  SmallVector<Value> readIndices(
1071  vecRank, rewriter.create<arith::ConstantIndexOp>(padOp.getLoc(), 0));
1072  auto read = rewriter.create<vector::TransferReadOp>(
1073  padOp.getLoc(), vecType, padOp.source(), readIndices, padValue);
1074 
1075  // Generate TransferWriteOp: Write to InsertSliceOp's dest tensor at
1076  // specified offsets. Write is fully in-bounds because a InsertSliceOp's
1077  // source must fit into the destination at the specified offsets.
1078  auto writeIndices =
1079  ofrToIndexValues(rewriter, padOp.getLoc(), insertOp.getMixedOffsets());
1080  SmallVector<bool> inBounds(vecRank, true);
1081  rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
1082  insertOp, read, insertOp.dest(), writeIndices,
1083  ArrayRef<bool>{inBounds});
1084 
1085  return success();
1086  }
1087 };
1088 
1090  RewritePatternSet &patterns, PatternBenefit baseBenefit) {
1091  patterns.add<GenericPadOpVectorizationPattern>(patterns.getContext(),
1092  baseBenefit);
1093  // Try these specialized patterns first before resorting to the generic one.
1097  patterns.getContext(), baseBenefit.getBenefit() + 1);
1098 }
1099 
1100 //----------------------------------------------------------------------------//
1101 // Forwarding patterns
1102 //----------------------------------------------------------------------------//
1103 
1104 /// Check whether there is any interleaved use of any `values` between
1105 /// `firstOp` and `secondOp`. Conservatively return `true` if any op or value
1106 /// is in a different block.
1107 static bool mayExistInterleavedUses(Operation *firstOp, Operation *secondOp,
1108  ValueRange values) {
1109  if (firstOp->getBlock() != secondOp->getBlock() ||
1110  !firstOp->isBeforeInBlock(secondOp)) {
1111  LDBG("interleavedUses precondition failed, firstOp: "
1112  << *firstOp << ", second op: " << *secondOp);
1113  return true;
1114  }
1115  for (auto v : values) {
1116  for (auto &u : v.getUses()) {
1117  Operation *owner = u.getOwner();
1118  if (owner == firstOp || owner == secondOp)
1119  continue;
1120  // TODO: this is too conservative, use dominance info in the future.
1121  if (owner->getBlock() == firstOp->getBlock() &&
1122  (owner->isBeforeInBlock(firstOp) || secondOp->isBeforeInBlock(owner)))
1123  continue;
1124  LDBG(" found interleaved op " << *owner << ", firstOp: " << *firstOp
1125  << ", second op: " << *secondOp);
1126  return true;
1127  }
1128  }
1129  return false;
1130 }
1131 
1132 /// Return the unique subview use of `v` if it is indeed unique, null
1133 /// otherwise.
1134 static memref::SubViewOp getSubViewUseIfUnique(Value v) {
1135  memref::SubViewOp subViewOp;
1136  for (auto &u : v.getUses()) {
1137  if (auto newSubViewOp = dyn_cast<memref::SubViewOp>(u.getOwner())) {
1138  if (subViewOp)
1139  return memref::SubViewOp();
1140  subViewOp = newSubViewOp;
1141  }
1142  }
1143  return subViewOp;
1144 }
1145 
1146 /// TODO: use interfaces, side-effects and aliasing analysis as appropriate,
1147 /// when available.
1149  vector::TransferReadOp xferOp, PatternRewriter &rewriter) const {
1150 
1151  // TODO: support mask.
1152  if (xferOp.mask())
1153  return failure();
1154 
1155  // Transfer into `view`.
1156  Value viewOrAlloc = xferOp.source();
1157  if (!viewOrAlloc.getDefiningOp<memref::ViewOp>() &&
1158  !viewOrAlloc.getDefiningOp<memref::AllocOp>())
1159  return failure();
1160 
1161  LDBG(viewOrAlloc);
1162 
1163  // Ensure there is exactly one subview of `viewOrAlloc` defining `subView`.
1164  memref::SubViewOp subViewOp = getSubViewUseIfUnique(viewOrAlloc);
1165  if (!subViewOp)
1166  return failure();
1167  Value subView = subViewOp.getResult();
1168  LDBG("with subView " << subView);
1169 
1170  // Find the copy into `subView` without interleaved uses.
1171  CopyOp copyOp;
1172  for (auto &u : subView.getUses()) {
1173  if (auto newCopyOp = dyn_cast<CopyOp>(u.getOwner())) {
1174  assert(newCopyOp.output().getType().isa<MemRefType>());
1175  if (newCopyOp.output() != subView)
1176  continue;
1177  LDBG("copy candidate " << *newCopyOp);
1178  if (mayExistInterleavedUses(newCopyOp, xferOp, {viewOrAlloc, subView}))
1179  continue;
1180  copyOp = newCopyOp;
1181  break;
1182  }
1183  }
1184  if (!copyOp)
1185  return failure();
1186  LDBG("with copy " << *copyOp);
1187 
1188  // Find the fill into `viewOrAlloc` without interleaved uses before the
1189  // copy.
1190  FillOp maybeFillOp;
1191  for (auto &u : viewOrAlloc.getUses()) {
1192  if (auto newFillOp = dyn_cast<FillOp>(u.getOwner())) {
1193  assert(newFillOp.output().getType().isa<MemRefType>());
1194  if (newFillOp.output() != viewOrAlloc)
1195  continue;
1196  LDBG("fill candidate " << *newFillOp);
1197  if (mayExistInterleavedUses(newFillOp, copyOp, {viewOrAlloc, subView}))
1198  continue;
1199  maybeFillOp = newFillOp;
1200  break;
1201  }
1202  }
1203  // Ensure padding matches.
1204  if (maybeFillOp && xferOp.padding() != maybeFillOp.value())
1205  return failure();
1206  if (maybeFillOp)
1207  LDBG("with maybeFillOp " << *maybeFillOp);
1208 
1209  // `in` is the subview that linalg.copy reads. Replace it.
1210  Value in = copyOp.input();
1211 
1212  // linalg.copy + linalg.fill can be used to create a padded local buffer.
1213  // The `masked` attribute is only valid on this padded buffer.
1214  // When forwarding to vector.transfer_read, the attribute must be reset
1215  // conservatively.
1216  Value res = rewriter.create<vector::TransferReadOp>(
1217  xferOp.getLoc(), xferOp.getVectorType(), in, xferOp.indices(),
1218  xferOp.permutation_mapAttr(), xferOp.padding(), xferOp.mask(),
1219  // in_bounds is explicitly reset
1220  /*inBoundsAttr=*/ArrayAttr());
1221 
1222  if (maybeFillOp)
1223  rewriter.eraseOp(maybeFillOp);
1224  rewriter.eraseOp(copyOp);
1225  rewriter.replaceOp(xferOp, res);
1226 
1227  return success();
1228 }
1229 
1230 /// TODO: use interfaces, side-effects and aliasing analysis as appropriate,
1231 /// when available.
1233  vector::TransferWriteOp xferOp, PatternRewriter &rewriter) const {
1234  // TODO: support mask.
1235  if (xferOp.mask())
1236  return failure();
1237 
1238  // Transfer into `viewOrAlloc`.
1239  Value viewOrAlloc = xferOp.source();
1240  if (!viewOrAlloc.getDefiningOp<memref::ViewOp>() &&
1241  !viewOrAlloc.getDefiningOp<memref::AllocOp>())
1242  return failure();
1243 
1244  // Ensure there is exactly one subview of `viewOrAlloc` defining `subView`.
1245  memref::SubViewOp subViewOp = getSubViewUseIfUnique(viewOrAlloc);
1246  if (!subViewOp)
1247  return failure();
1248  Value subView = subViewOp.getResult();
1249 
1250  // Find the copy from `subView` without interleaved uses.
1251  CopyOp copyOp;
1252  for (auto &u : subViewOp.getResult().getUses()) {
1253  if (auto newCopyOp = dyn_cast<CopyOp>(u.getOwner())) {
1254  if (newCopyOp.getInputOperand(0)->get() != subView)
1255  continue;
1256  if (mayExistInterleavedUses(xferOp, newCopyOp, {viewOrAlloc, subView}))
1257  continue;
1258  copyOp = newCopyOp;
1259  break;
1260  }
1261  }
1262  if (!copyOp)
1263  return failure();
1264 
1265  // `out` is the subview copied into that we replace.
1266  assert(copyOp.output().getType().isa<MemRefType>());
1267  Value out = copyOp.output();
1268 
1269  // Forward vector.transfer into copy.
1270  // linalg.copy + linalg.fill can be used to create a padded local buffer.
1271  // The `masked` attribute is only valid on this padded buffer.
1272  // When forwarding to vector.transfer_write, the attribute must be reset
1273  // conservatively.
1274  rewriter.create<vector::TransferWriteOp>(
1275  xferOp.getLoc(), xferOp.vector(), out, xferOp.indices(),
1276  xferOp.permutation_mapAttr(), xferOp.mask(),
1277  // in_bounds is explicitly reset
1278  /*inBoundsAttr=*/ArrayAttr());
1279 
1280  rewriter.eraseOp(copyOp);
1281  rewriter.eraseOp(xferOp);
1282 
1283  return success();
1284 }
1285 
1286 //===----------------------------------------------------------------------===//
1287 // Convolution vectorization patterns
1288 //===----------------------------------------------------------------------===//
1289 
1290 template <int N>
1291 static void bindShapeDims(ShapedType shapedType) {}
1292 
1293 template <int N, typename IntTy, typename... IntTy2>
1294 static void bindShapeDims(ShapedType shapedType, IntTy &val, IntTy2 &...vals) {
1295  val = shapedType.getShape()[N];
1296  bindShapeDims<N + 1, IntTy2 &...>(shapedType, vals...);
1297 }
1298 
1299 /// Bind a pack of int& to the leading dimensions of shapedType.getShape().
1300 template <typename... IntTy>
1301 static void bindShapeDims(ShapedType shapedType, IntTy &...vals) {
1302  bindShapeDims<0>(shapedType, vals...);
1303 }
1304 
1305 namespace {
1306 /// Generate a vector implementation for either:
1307 /// ```
1308 /// Op def: ( n, w, c, kw, f )
1309 /// Iters: ({Par(), Par(), Par(), Red(), Red()})
1310 /// Layout: {{n, strideW * w + dilationW * kw, c}, {kw, c, f}, {n, w, f}}
1311 /// ```
1312 /// kw is unrolled, w is unrolled iff dilationW > 1.
1313 ///
1314 /// or
1315 ///
1316 /// ```
1317 /// Op def: ( n, w, c, kw )
1318 /// Iters: ({Par(), Par(), Par(), Red()})
1319 /// Layout: {{n, strideW * w + dilationW * kw, c}, {kw, c}, {n, w, c}}
1320 /// ```
1321 /// kw is unrolled, w is unrolled iff dilationW > 1.
1322 struct Conv1DNwcGenerator : public StructuredGenerator<LinalgOp> {
1323  Conv1DNwcGenerator(OpBuilder &builder, LinalgOp linalgOp, int strideW,
1324  int dilationW)
1325  : StructuredGenerator<LinalgOp>(builder, linalgOp), valid(false),
1326  strideW(strideW), dilationW(dilationW) {
1327  // Determine whether `linalgOp` can be generated with this generator
1328  if (linalgOp.getNumInputs() != 2 || linalgOp.getNumOutputs() != 1)
1329  return;
1330  lhsShaped = linalgOp.inputs()[0];
1331  rhsShaped = linalgOp.inputs()[1];
1332  resShaped = linalgOp.outputs()[0];
1333  lhsShapedType = lhsShaped.getType().dyn_cast<ShapedType>();
1334  rhsShapedType = rhsShaped.getType().dyn_cast<ShapedType>();
1335  resShapedType = resShaped.getType().dyn_cast<ShapedType>();
1336  if (!lhsShapedType || !rhsShapedType || !resShapedType)
1337  return;
1338  if (lhsShapedType.getRank() != 3 ||
1339  (rhsShapedType.getRank() != 2 && rhsShapedType.getRank() != 3) ||
1340  resShapedType.getRank() != 3)
1341  return;
1342 
1343  // Check for reduction `add` preceded by `mul`.
1344  Operation *reduceOp = matchLinalgReduction(linalgOp.getOutputOperand(0));
1345  if (!reduceOp)
1346  return;
1348  maybeKind = getCombinerOpKind(reduceOp);
1349  if (!maybeKind || *maybeKind != vector::CombiningKind::ADD)
1350  return;
1351  maybeKind = getCombinerOpKind(&(linalgOp->getRegion(0).front().front()));
1352  if (!maybeKind || *maybeKind != vector::CombiningKind::MUL)
1353  return;
1354 
1355  // The op is now known to be valid.
1356  valid = true;
1357  }
1358 
1359  /// Generate a vector implementation for:
1360  /// ```
1361  /// Op def: ( n, w, c, kw, f )
1362  /// Iters: ({Par(), Par(), Par(), Red(), Red()})
1363  /// Layout: {{n, strideW * w + dilationW * kw, c}, {kw, c, f}, {n, w, f}}
1364  /// ```
1365  /// kw is always unrolled.
1366  /// TODO: w (resp. kw) is unrolled when the strideW ( resp. dilationW) is
1367  /// > 1.
1368  FailureOr<Operation *> conv() {
1369  if (!valid)
1370  return failure();
1371 
1372  int64_t nSize, wSize, cSize, kwSize, fSize;
1373  // kernel{kw, c, f}
1374  bindShapeDims(rhsShapedType, kwSize, cSize, fSize);
1375  // out{n, w, f}
1376  bindShapeDims(resShapedType, nSize, wSize);
1377 
1378  vector::TransferWriteOp write;
1379  Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
1380 
1381  // w is unrolled (i.e. wSizeStep == 1) iff strideW > 1.
1382  // When strideW == 1, we can batch the contiguous loads and avoid
1383  // unrolling
1384  int64_t wSizeStep = strideW == 1 ? wSize : 1;
1385 
1386  Type lhsEltType = lhsShapedType.getElementType();
1387  Type rhsEltType = rhsShapedType.getElementType();
1388  Type resEltType = resShapedType.getElementType();
1389  VectorType lhsType = VectorType::get(
1390  {nSize,
1391  // iw = ow * sw + kw * dw - 1
1392  // (i.e. 16 convolved with 3 (@stride 1 dilation 1) -> 14)
1393  // Perform the proper inclusive -> exclusive -> inclusive.
1394  ((wSize - 1) * strideW + 1) + ((kwSize - 1) * dilationW + 1) - 1,
1395  cSize},
1396  lhsEltType);
1397  VectorType rhsType = VectorType::get({kwSize, cSize, fSize}, rhsEltType);
1398  VectorType resType = VectorType::get({nSize, wSize, fSize}, resEltType);
1399 
1400  // Read lhs slice of size {w * strideW + kw * dilationW, c, f} @ [0, 0,
1401  // 0].
1402  Value lhs = builder.create<vector::TransferReadOp>(
1403  loc, lhsType, lhsShaped, ValueRange{zero, zero, zero});
1404  // Read rhs slice of size {kw, c, f} @ [0, 0, 0].
1405  Value rhs = builder.create<vector::TransferReadOp>(
1406  loc, rhsType, rhsShaped, ValueRange{zero, zero, zero});
1407  // Read res slice of size {n, w, f} @ [0, 0, 0].
1408  Value res = builder.create<vector::TransferReadOp>(
1409  loc, resType, resShaped, ValueRange{zero, zero, zero});
1410 
1411  //===------------------------------------------------------------------===//
1412  // Begin vector-only rewrite part
1413  //===------------------------------------------------------------------===//
1414  // Unroll along kw and read slices of lhs and rhs.
1415  SmallVector<Value> lhsVals, rhsVals, resVals;
1416  // Extract lhs slice of size {n, wSizeStep, c} @ [0, sw * w + dw * kw, 0].
1417  for (int64_t kw = 0; kw < kwSize; ++kw) {
1418  for (int64_t w = 0; w < wSize; w += wSizeStep) {
1419  lhsVals.push_back(builder.create<vector::ExtractStridedSliceOp>(
1420  loc, lhs,
1421  /*offsets=*/ArrayRef<int64_t>{0, w * strideW + kw * dilationW, 0},
1422  /*sizes=*/ArrayRef<int64_t>{nSize, wSizeStep, cSize},
1423  /*strides=*/ArrayRef<int64_t>{1, 1, 1}));
1424  }
1425  }
1426  // Extract rhs slice of size {c, f} @ [kw].
1427  for (int64_t kw = 0; kw < kwSize; ++kw) {
1428  rhsVals.push_back(builder.create<vector::ExtractOp>(
1429  loc, rhs, /*offsets=*/ArrayRef<int64_t>{kw}));
1430  }
1431  // Extract res slice: {n, wSizeStep, f} @ [0, w, 0].
1432  for (int64_t w = 0; w < wSize; w += wSizeStep) {
1433  resVals.push_back(builder.create<vector::ExtractStridedSliceOp>(
1434  loc, res,
1435  /*offsets=*/ArrayRef<int64_t>{0, w, 0},
1436  /*sizes=*/ArrayRef<int64_t>{nSize, wSizeStep, fSize},
1437  /*strides=*/ArrayRef<int64_t>{1, 1, 1}));
1438  }
1439 
1440  auto linearIndex = [&](int64_t kw, int64_t w) {
1441  return kw * (wSize / wSizeStep) + w;
1442  };
1443 
1444  // Compute contraction: O{n, w, f} += I{n, sw * w + dw * kw, c} * F{c, f}
1445  for (int64_t kw = 0; kw < kwSize; ++kw) {
1446  for (int64_t w = 0; w < wSize; w += wSizeStep) {
1447  resVals[w] = conv1dSliceAsContraction(
1448  builder, loc, lhsVals[linearIndex(kw, w)], rhsVals[kw], resVals[w]);
1449  }
1450  }
1451 
1452  // Write back res slice: {n, wSizeStep, f} @ [0, w, 0].
1453  // This does not depend on kw.
1454  for (int64_t w = 0; w < wSize; w += wSizeStep) {
1455  res = builder.create<vector::InsertStridedSliceOp>(
1456  loc, resVals[w], res,
1457  /*offsets=*/ArrayRef<int64_t>{0, w, 0},
1458  /*strides=*/ArrayRef<int64_t>{1, 1, 1});
1459  }
1460  //===------------------------------------------------------------------===//
1461  // End vector-only rewrite part
1462  //===------------------------------------------------------------------===//
1463 
1464  // Write back res slice of size {n, w, f} @ [0, 0, 0].
1465  return builder
1466  .create<vector::TransferWriteOp>(loc, res, resShaped,
1467  ValueRange{zero, zero, zero})
1468  .getOperation();
1469  }
1470 
1471  // Create a contraction: lhs{n, w, c} * rhs{c, f} -> res{n, w, f}
1472  Value conv1dSliceAsContraction(OpBuilder &b, Location loc, Value lhs,
1473  Value rhs, Value res) {
1474  StringRef par = Par().strRef, red = Red().strRef;
1475  AffineExpr n, w, f, c;
1476  bindDims(ctx, n, w, f, c);
1477  return builder.create<vector::ContractionOp>(
1478  loc, lhs, rhs, res,
1479  /*indexingMaps=*/MapList{{n, w, c}, {c, f}, {n, w, f}},
1480  /*iteratorTypes=*/ArrayRef<StringRef>{par, par, par, red});
1481  }
1482 
1483  /// Generate a vector implementation for:
1484  /// ```
1485  /// Op def: ( n, w, c, kw)
1486  /// Iters: ({Par(), Par(), Par(), Red()})
1487  /// Layout: {{n, strideW * w + dilationW * kw, c}, {kw, c}, {n, w, c}}
1488  /// ```
1489  /// kw is always unrolled.
1490  /// TODO: w (resp. kw) is unrolled when the strideW ( resp. dilationW) is
1491  /// > 1.
1492  FailureOr<Operation *> depthwiseConv() {
1493  if (!valid)
1494  return failure();
1495 
1496  int64_t nSize, wSize, cSize, kwSize;
1497  // kernel{kw, c}
1498  bindShapeDims(rhsShapedType, kwSize, cSize);
1499  // out{n, w, c}
1500  bindShapeDims(resShapedType, nSize, wSize);
1501 
1502  vector::TransferWriteOp write;
1503  Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
1504 
1505  // w is unrolled (i.e. wSizeStep == 1) iff strideW > 1.
1506  // When strideW == 1, we can batch the contiguous loads and avoid
1507  // unrolling
1508  int64_t wSizeStep = strideW == 1 ? wSize : 1;
1509 
1510  Type lhsEltType = lhsShapedType.getElementType();
1511  Type rhsEltType = rhsShapedType.getElementType();
1512  Type resEltType = resShapedType.getElementType();
1513  VectorType lhsType = VectorType::get(
1514  {nSize,
1515  // iw = ow * sw + kw * dw - 1
1516  // (i.e. 16 convolved with 3 (@stride 1 dilation 1) -> 14)
1517  ((wSize - 1) * strideW + 1) + ((kwSize - 1) * dilationW + 1) - 1,
1518  cSize},
1519  lhsEltType);
1520  VectorType rhsType = VectorType::get({kwSize, cSize}, rhsEltType);
1521  VectorType resType = VectorType::get({nSize, wSize, cSize}, resEltType);
1522 
1523  // Read lhs slice of size {n, w * strideW + kw * dilationW, c} @ [0, 0,
1524  // 0].
1525  Value lhs = builder.create<vector::TransferReadOp>(
1526  loc, lhsType, lhsShaped, ValueRange{zero, zero, zero});
1527  // Read rhs slice of size {kw, c} @ [0, 0].
1528  Value rhs = builder.create<vector::TransferReadOp>(loc, rhsType, rhsShaped,
1529  ValueRange{zero, zero});
1530  // Read res slice of size {n, w, c} @ [0, 0, 0].
1531  Value res = builder.create<vector::TransferReadOp>(
1532  loc, resType, resShaped, ValueRange{zero, zero, zero});
1533 
1534  //===------------------------------------------------------------------===//
1535  // Begin vector-only rewrite part
1536  //===------------------------------------------------------------------===//
1537  // Unroll along kw and read slices of lhs and rhs.
1538  SmallVector<Value> lhsVals, rhsVals, resVals;
1539  // Extract lhs slice of size {n, wSizeStep, c}
1540  // @ [0, sw * w + dw * kw, 0].
1541  for (int64_t kw = 0; kw < kwSize; ++kw) {
1542  for (int64_t w = 0; w < wSize; w += wSizeStep) {
1543  lhsVals.push_back(builder.create<vector::ExtractStridedSliceOp>(
1544  loc, lhs,
1545  /*offsets=*/ArrayRef<int64_t>{0, w * strideW + kw * dilationW, 0},
1546  /*sizes=*/ArrayRef<int64_t>{nSize, wSizeStep, cSize},
1547  /*strides=*/ArrayRef<int64_t>{1, 1, 1}));
1548  }
1549  }
1550  // Extract rhs slice of size {c} @ [kw].
1551  for (int64_t kw = 0; kw < kwSize; ++kw) {
1552  rhsVals.push_back(builder.create<vector::ExtractOp>(
1553  loc, rhs, /*offsets=*/ArrayRef<int64_t>{kw}));
1554  }
1555  // Extract res slice: {n, wSizeStep, c} @ [0, w, 0].
1556  for (int64_t w = 0; w < wSize; w += wSizeStep) {
1557  resVals.push_back(builder.create<vector::ExtractStridedSliceOp>(
1558  loc, res,
1559  /*offsets=*/ArrayRef<int64_t>{0, w, 0},
1560  /*sizes=*/ArrayRef<int64_t>{nSize, wSizeStep, cSize},
1561  /*strides=*/ArrayRef<int64_t>{1, 1, 1}));
1562  }
1563 
1564  auto linearIndex = [&](int64_t kw, int64_t w) {
1565  return kw * (wSize / wSizeStep) + w;
1566  };
1567 
1568  // Compute contraction: O{n, w, c} += I{n, sw * w + dw * kw, c} * F{c}
1569  for (int64_t kw = 0; kw < kwSize; ++kw) {
1570  for (int64_t w = 0; w < wSize; w += wSizeStep) {
1571  resVals[w] = depthwiseConv1dSliceAsFma(
1572  builder, loc, lhsVals[linearIndex(kw, w)], rhsVals[kw], resVals[w]);
1573  }
1574  }
1575 
1576  // Write back res slice: {n, wSizeStep, c} @ [0, w, 0].
1577  // This does not depend on kw.
1578  for (int64_t w = 0; w < wSize; w += wSizeStep) {
1579  res = builder.create<vector::InsertStridedSliceOp>(
1580  loc, resVals[w], res,
1581  /*offsets=*/ArrayRef<int64_t>{0, w, 0},
1582  /*strides=*/ArrayRef<int64_t>{1, 1, 1});
1583  }
1584  //===------------------------------------------------------------------===//
1585  // End vector-only rewrite part
1586  //===------------------------------------------------------------------===//
1587 
1588  // Write back res slice of size {n, w, c} @ [0, 0, 0].
1589  return builder
1590  .create<vector::TransferWriteOp>(loc, res, resShaped,
1591  ValueRange{zero, zero, zero})
1592  .getOperation();
1593  }
1594 
1595  /// Lower lhs{n, w, c} * rhs{c} -> res{n, w, c} to fma.
1596  Value depthwiseConv1dSliceAsFma(OpBuilder &b, Location loc, Value lhs,
1597  Value rhs, Value res) {
1598  Value bcast = builder.create<vector::BroadcastOp>(loc, res.getType(), rhs);
1599  return b.create<vector::FMAOp>(loc, lhs, bcast, res);
1600  }
1601 
1602  /// Entry point that transposes into the common form:
1603  /// {{n, strideW * w + dilationW * kw, c}, {kw, c, f}, {n, w, f}}
1604  FailureOr<Operation *> generateConv() {
1605  AffineExpr n, w, f, kw, c;
1606  bindDims(ctx, n, w, f, kw, c);
1607  if (!iters({Par(), Par(), Par(), Red(), Red()}))
1608  return failure();
1609 
1610  // No transposition needed.
1611  if (layout({/*lhsIndex*/ {n, strideW * w + dilationW * kw, c},
1612  /*rhsIndex*/ {kw, c, f},
1613  /*resIndex*/ {n, w, f}}))
1614  return conv();
1615  return failure();
1616  }
1617 
1618  /// Entry point that transposes into the common form:
1619  /// {{n, strideW * w + dilationW * kw, c}, {kw, c}, {n, w, c}}
1620  FailureOr<Operation *> generateDilatedConv() {
1621  AffineExpr n, w, c, kw;
1622  bindDims(ctx, n, w, c, kw);
1623  if (!iters({Par(), Par(), Par(), Red()}))
1624  return failure();
1625 
1626  // No transposition needed.
1627  if (layout({/*lhsIndex*/ {n, strideW * w + dilationW * kw, c},
1628  /*rhsIndex*/ {kw, c},
1629  /*resIndex*/ {n, w, c}}))
1630  return depthwiseConv();
1631  return failure();
1632  }
1633 
1634 private:
1635  bool valid;
1636  int strideW, dilationW;
1637  Value lhsShaped, rhsShaped, resShaped;
1638  ShapedType lhsShapedType, rhsShapedType, resShapedType;
1639 };
1640 } // namespace
1641 
1642 /// Helper function to vectorize a LinalgOp with convolution semantics.
1643 // TODO: extend the generic vectorization to support windows and drop this.
1645  // The ConvolutionOpInterface gives us guarantees of existence for
1646  // strides/dilations. However, we do not need to rely on those, we can simply
1647  // use them if present, otherwise use the default and let the generic conv.
1648  // matcher in the ConvGenerator succeed or fail.
1649  auto strides = op->getAttrOfType<DenseIntElementsAttr>("strides");
1650  auto dilations = op->getAttrOfType<DenseIntElementsAttr>("dilations");
1651  auto stride = strides ? *strides.getValues<uint64_t>().begin() : 1;
1652  auto dilation = dilations ? *dilations.getValues<uint64_t>().begin() : 1;
1653  Conv1DNwcGenerator e(b, op, stride, dilation);
1654  auto res = e.generateConv();
1655  if (succeeded(res))
1656  return res;
1657  return e.generateDilatedConv();
1658 }
1659 
1662 
1664  PatternRewriter &rewriter) const override {
1665  FailureOr<Operation *> resultOrFail = vectorizeConvolution(rewriter, op);
1666  if (failed(resultOrFail))
1667  return failure();
1668  Operation *newOp = *resultOrFail;
1669  if (newOp->getNumResults() == 0) {
1670  rewriter.eraseOp(op.getOperation());
1671  return success();
1672  }
1673  assert(newOp->getNumResults() == 1 && "expected single result");
1674  rewriter.replaceOp(op.getOperation(), newOp->getResult(0));
1675  return success();
1676  }
1677 };
1678 
1680  RewritePatternSet &patterns, PatternBenefit benefit) {
1681  patterns.add<VectorizeConvolution>(patterns.getContext(), benefit);
1682 }
static memref::SubViewOp getSubViewUseIfUnique(Value v)
Return the unique subview use of v if it is indeed unique, null otherwise.
enum VectorizationStatus status
Return status from vectorizing the current op.
Include the generated interface declarations.
static Operation * createVectorizedOp(OpBuilder &b, Operation *op, ValueRange newOperands, ArrayRef< Type > types)
Create a new vectorized verstion of op with the given operands and types.
OpTy create(Location location, Args &&...args)
Create an operation of specific op type at the current insertion point.
Definition: Builders.h:430
This class contains a list of basic blocks and a link to the parent operation it is attached to...
Definition: Region.h:26
Rewrite use of tensor::PadOp result in TransferReadOp.
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
Definition: AffineMap.cpp:673
static int64_t getIntFromAttr(Attribute attr)
Helper function that retrieves the value of an IntegerAttr.
U cast() const
Definition: Attributes.h:123
void createOrFold(SmallVectorImpl< Value > &results, Location location, Args &&...args)
Create an operation of specific op type at the current insertion point, and immediately try to fold i...
Definition: Builders.h:444
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:881
AffineMap compose(AffineMap map) const
Returns the AffineMap resulting from composing this with map.
Definition: AffineMap.cpp:444
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:28
static bool hasOnlyScalarElementwiseOp(Region &r)
Detect whether r has only ConstantOp, ElementwiseMappable and YieldOp.
Attribute getZeroAttr(Type type)
Definition: Builders.cpp:264
operand_range getOperands()
Returns an iterator on the underlying Value&#39;s.
Definition: Operation.h:247
Base pattern for rewriting tensor::PadOps whose result is consumed by a given operation type OpTy...
Block represents an ordered list of Operations.
Definition: Block.h:29
Block & front()
Definition: Region.h:65
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:329
Rewrite a tensor::PadOp into a sequence of InitTensorOp, FillOp and InsertSliceOp.
Definition: Transforms.h:1177
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
Operation * clone(Operation &op, BlockAndValueMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
Definition: Builders.cpp:457
ArrayRef< NamedAttribute > getAttrs()
Return all of the attributes on this operation.
Definition: Operation.h:308
bool hasElementwiseMappableTraits(Operation *op)
Together, Elementwise, Scalarizable, Vectorizable, and Tensorizable provide an easy way for scalar op...
Definition: Operation.cpp:1122
OpListType & getOperations()
Definition: Block.h:128
static OpType getSingleOpOfType(Block &block)
Return the unique instance of OpType in block if it is indeed unique.
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Definition: LogicalResult.h:72
static bool isElementwise(Operation *op)
OpInterfaceRewritePattern(MLIRContext *context, PatternBenefit benefit=1)
Definition: PatternMatch.h:372
static Value buildVectorWrite(OpBuilder &b, Value value, OpOperand *outputOperand)
Build a vector.transfer_write of value into outputOperand at indices set to all 0; where outputOperan...
bool isBeforeInBlock(Operation *other)
Given an operation &#39;other&#39; that is within the same parent block, return whether the current operation...
Definition: Operation.cpp:271
static bool mayExistInterleavedUses(Operation *firstOp, Operation *secondOp, ValueRange values)
Check whether there is any interleaved use of any values between firstOp and secondOp.
LogicalResult matchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const override
bool succeeded(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a success value...
Definition: LogicalResult.h:68
LogicalResult matchAndRewrite(vector::TransferWriteOp xferOp, PatternRewriter &rewriter) const override
TODO: use interfaces, side-effects and aliasing analysis as appropriate, when available.
AffineMap compressUnusedDims(AffineMap map)
Drop the dims that are not used.
Definition: AffineMap.cpp:568
T lookup(T from) const
Lookup a mapped value within the map.
unsigned getArgNumber() const
Returns the number of this argument.
Definition: Value.h:310
Block * getBlock()
Returns the operation block that contains this operation.
Definition: Operation.h:96
BlockArgument getArgument(unsigned i)
Definition: Block.h:120
static constexpr const bool value
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:48
LogicalResult matchAndRewrite(vector::TransferReadOp xferOp, PatternRewriter &rewriter) const override
TODO: use interfaces, side-effects and aliasing analysis as appropriate, when available.
BroadcastableToResult isBroadcastableTo(Type srcType, VectorType dstVectorType, std::pair< int, int > *mismatchingDims=nullptr)
Definition: VectorOps.cpp:1561
void map(Block *from, Block *to)
Inserts a new mapping for &#39;from&#39; to &#39;to&#39;.
unsigned short getBenefit() const
If the corresponding pattern can match, return its benefit. If the.
static void bindShapeDims(ShapedType shapedType)
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
unsigned getOperandNumber()
Return which operand this is in the OpOperand list of the Operation.
Definition: Value.cpp:212
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
virtual void replaceOp(Operation *op, ValueRange newValues)
This method replaces the results of the operation with the specified list of values.
This class provides support for representing a failure result, or a valid value of type T...
Definition: LogicalResult.h:77
Op vectorized and custom function took care of replacement logic.
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
llvm::Optional< vector::CombiningKind > getCombinerOpKind(Operation *combinerOp)
Return vector::CombiningKind for the given op.
static SmallVector< bool > getReductionMask(LinalgOp linalgOp)
void getUsedValuesDefinedAbove(Region &region, Region &limit, SetVector< Value > &values)
Fill values with a list of values defined at the ancestors of the limit region and used within region...
Definition: RegionUtils.cpp:59
Op vectorized into a new Op whose results will replace original Op&#39;s results.
static LogicalResult tryVectorizeCopy(PatternRewriter &rewriter, tensor::PadOp padOp, Value dest)
Vectorize the copying of a tensor::PadOp&#39;s source.
void populateConvolutionVectorizationPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Populate patterns for vectorizing low-D convolution ops.
SmallVector< T > applyPermutationMap(AffineMap map, llvm::ArrayRef< T > source)
Apply a permutation from map to source and return the result.
Definition: AffineMap.h:530
U dyn_cast() const
Definition: Types.h:244
bool isProjectedPermutation(bool allowZeroInResults=false) const
Returns true if the AffineMap represents a subset (i.e.
Definition: AffineMap.cpp:483
Attributes are known-constant values of operations.
Definition: Attributes.h:24
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
Definition: PatternMatch.h:32
U dyn_cast() const
Definition: Value.h:99
Operation * createOperation(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:380
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:206
static VectorizationResult vectorizeLinalgYield(OpBuilder &b, Operation *op, const BlockAndValueMapping &bvm, LinalgOp linalgOp, SmallVectorImpl< Value > &newResults)
Helper function to vectorize the terminator of a linalgOp.
Rewrite a tensor::PadOp into a sequence of InitTensorOp, FillOp and InsertSliceOp.
void addAttributes(ArrayRef< NamedAttribute > newAttributes)
Add an array of named attributes.
Base type for affine expression.
Definition: AffineExpr.h:68
OpResult getResult(unsigned idx)
Get the &#39;idx&#39;th result of this operation.
Definition: Operation.h:276
LogicalResult rewriteUser(PatternRewriter &rewriter, tensor::PadOp padOp, vector::TransferReadOp xferOp) const override
Rewrite use of tensor::PadOp result in InsertSliceOp.
static WalkResult advance()
Definition: Visitors.h:51
Value matchReduction(ArrayRef< BlockArgument > iterCarriedArgs, unsigned redPos, SmallVectorImpl< Operation *> &combinerOps)
Utility to match a generic reduction given a list of iteration-carried arguments, iterCarriedArgs and...
void updateRootInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around a root update of an operation.
Definition: PatternMatch.h:789
static LogicalResult vectorizeStaticLinalgOpPrecondition(linalg::LinalgOp op)
bool hasSameTensorSize(Value beforePadding, tensor::ExtractSliceOp afterTrimming) const
Check if beforePadding and afterTrimming have the same tensor size, i.e., same dimensions.
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:106
IRValueT get() const
Return the current value being used by this operand.
Definition: UseDefLists.h:133
LogicalResult matchAndRewrite(tensor::PadOp padOp, PatternRewriter &rewriter) const final
This represents an operation in an abstracted form, suitable for use with the builder APIs...
A multi-dimensional affine map Affine map&#39;s are immutable like Type&#39;s, and they are uniqued...
Definition: AffineMap.h:38
static Operation * reduceIfNeeded(OpBuilder &b, LinalgOp linalgOp, Operation *op, Value reduceValue, Value initialValue, const BlockAndValueMapping &bvm)
Emit reduction operations if the shapes of the value to reduce is different that the result shape...
static WalkResult interrupt()
Definition: Visitors.h:50
RetT walk(FnT &&callback)
Walk the operations in this block.
Definition: Block.h:273
VectorizationStatus
Helper data structure to represent the result of vectorization.
This class represents an argument of a Block.
Definition: Value.h:298
Location getLoc() const
Return the location of this value.
Definition: Value.cpp:26
ArrayAttr getBoolArrayAttr(ArrayRef< bool > values)
Definition: Builders.cpp:209
LogicalResult rewriteUser(PatternRewriter &rewriter, tensor::PadOp padOp, tensor::InsertSliceOp insertOp) const override
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:72
bool isReductionIterator(Attribute attr)
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:84
GenericPadOpVectorizationPattern(MLIRContext *context, PatternBenefit benefit=1)
static Value buildMultiDimReduce(OpBuilder &b, Operation *reduceOp, Value valueToReduce, const SmallVector< bool > &reductionMask)
Create MultiDimReductionOp to compute the reduction for reductionOp.
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:355
static VectorizationResult vectorizeOneOp(OpBuilder &b, LinalgOp linalgOp, Operation *op, const BlockAndValueMapping &bvm, ArrayRef< CustomVectorizationHook > customVectorizationHooks)
Generic vectorization for a single operation op, given already vectorized operands carried by bvm...
static VectorizationResult vectorizeLinalgIndex(OpBuilder &b, Operation *op, LinalgOp linalgOp)
Helper function to vectorize the index operations of a linalgOp.
AffineMap inverseAndBroadcastProjectedPermuation(AffineMap map)
Return the reverse map of a projected permutation where the projected dimensions are transformed into...
Definition: AffineMap.cpp:697
static LogicalResult vectorizeLinalgOpPrecondition(LinalgOp linalgOp)
OpTy replaceOpWithNewOp(Operation *op, Args &&... args)
Replaces the result op with a new op that is created without verification.
Definition: PatternMatch.h:741
Type getType() const
Return the type of this value.
Definition: Value.h:117
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&... args)
Add an instance of each of the pattern types &#39;Ts&#39; to the pattern list with the given arguments...
Definition: PatternMatch.h:930
IndexType getIndexType()
Definition: Builders.cpp:48
static LogicalResult vectorizeAsLinalgGeneric(OpBuilder &b, LinalgOp linalgOp, SmallVectorImpl< Value > &newResults)
Generic vectorization function that rewrites the body of a linalgOp into vector form.
static Operation * matchLinalgReduction(OpOperand *outputOperand)
Check whether outputOperand is a reduction with a single combiner operation.
bool succeeded() const
Returns true if the provided LogicalResult corresponds to a success value.
Definition: LogicalResult.h:41
Op failed to vectorize.
static bool allIndexingsAreProjectedPermutation(LinalgOp op)
Helper function to vectorize a linalgOp with contraction semantics in a generic fashion.
Specialization of arith.constant op that returns an integer of index type.
Definition: Arithmetic.h:78
static VectorType vectorType(CodeGen &codegen, Type etp)
Constructs vector type.
Operation * getOwner() const
Return the owner of this operand.
Definition: UseDefLists.h:37
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:55
This class represents an operand of an operation.
Definition: Value.h:249
LogicalResult vectorize(RewriterBase &builder, LinalgOp linalgOp)
Emit a suitable vector form for a Linalg op with fully static shape.
static Value broadcastIfNeeded(OpBuilder &b, Value value, ArrayRef< int64_t > shape)
Broadcast value to a vector of shape if possible.
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:273
static SmallVector< Value > ofrToIndexValues(OpBuilder &builder, Location loc, ArrayRef< OpFoldResult > ofrs)
Given an ArrayRef of OpFoldResults, return a vector of Values.
static FailureOr< Operation * > vectorizeConvolution(OpBuilder &b, LinalgOp convOp)
Try to vectorize convOp as a convolution.
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
Definition: AffineExpr.h:328
static AffineMap reindexIndexingMap(AffineMap map)
Given an indexing map coming from a LinalgOp indexing, restricted to a projectedPermutation, compress the unused dimensions to serve as a permutation_map for a vector transfer operation.
Block::iterator getInsertionPoint() const
Returns the current insertion point of the builder.
Definition: Builders.h:376
OperationName getName()
The name of an operation is the key identifier for it.
Definition: Operation.h:57
std::function< VectorizationResult(Operation *, const BlockAndValueMapping &)> CustomVectorizationHook
void populatePadOpVectorizationPatterns(RewritePatternSet &patterns, PatternBenefit baseBenefit=1)
Populates patterns with patterns that vectorize linalg.pad_tensor.
bool isa() const
Definition: Types.h:234
Optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
Rewrite use of tensor::PadOp result in TransferWriteOp.
result_range getResults()
Definition: Operation.h:284
This class helps build Operations.
Definition: Builders.h:177
This class provides an abstraction over the different types of ranges over Values.
use_range getUses() const
Returns a range of all uses, which is useful for iterating over all uses.
Definition: Value.h:196
static LogicalResult reductionPreconditions(LinalgOp op)
result_type_range getResultTypes()
Definition: Operation.h:297
Helper StructuredGenerator class to manipulate and rewrite ops with StructuredOpInterface.
DenseIntElementsAttr getIndexVectorAttr(ArrayRef< int64_t > values)
Definition: Builders.cpp:121
MLIRContext * getContext() const
Definition: PatternMatch.h:906
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:688
bool isEqualConstantIntOrValue(OpFoldResult ofr1, OpFoldResult ofr2)
Return true if ofr1 and ofr2 are the same integer constant attribute values or the same SSA value...
OpInterfaceRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting a...
Definition: PatternMatch.h:370
#define LDBG(X)
Operation * newOp
New vectorized operation to replace the current op.
An attribute that represents a reference to a dense integer vector or tensor object.
LogicalResult rewriteUser(PatternRewriter &rewriter, tensor::PadOp padOp, vector::TransferWriteOp xferOp) const override