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