MLIR  22.0.0git
LinalgInterfaces.cpp
Go to the documentation of this file.
1 //===- LinalgInterfaces.cpp - Linalg interfaces implementation ------------===//
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 
10 
16 #include "mlir/IR/AffineExpr.h"
18 #include "mlir/IR/AffineMap.h"
20 #include "mlir/IR/MLIRContext.h"
21 #include "mlir/IR/TypeUtilities.h"
22 #include "llvm/ADT/STLExtras.h"
23 #include "llvm/ADT/SetOperations.h"
24 #include "llvm/ADT/SmallBitVector.h"
25 #include "llvm/ADT/SmallVector.h"
26 #include "llvm/Support/Casting.h"
27 #include "llvm/Support/raw_ostream.h"
28 #include <optional>
29 
30 using namespace mlir;
31 using namespace mlir::linalg;
32 
33 /// Include the definitions of the copy operation interface.
34 #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.cpp.inc"
35 
36 //===----------------------------------------------------------------------===//
37 // Interface utility functions
38 //===----------------------------------------------------------------------===//
39 
41  linalg::LinalgOp linalgOp, ArrayRef<OpOperand *> droppedOperands) {
42  SmallVector<AffineMap> indexingMaps;
43  for (auto &opOperand : linalgOp->getOpOperands()) {
44  if (llvm::is_contained(droppedOperands, &opOperand))
45  continue;
46  indexingMaps.push_back(linalgOp.getMatchingIndexingMap(&opOperand));
47  }
48  if (indexingMaps.empty()) {
49  // If there are no indexing maps, the operand can only be dropped
50  // if the op has no loops.
51  return linalgOp.getNumLoops() == 0;
52  }
54  indexingMaps, linalgOp.getContext())) != AffineMap();
55 }
56 
57 //===----------------------------------------------------------------------===//
58 // CopyOpInterface implementation
59 //===----------------------------------------------------------------------===//
60 
61 bool linalg::isaCopyOpInterface(LinalgOp op) {
62  // Check all loops are parallel and linalgOp is single input and output.
63  if (!op.isAllParallelLoops() || !op.isSingleInputOutput())
64  return false;
65 
66  auto mapRange = op.getIndexingMapsArray();
67  if (mapRange.size() != 2 || !mapRange.front().isIdentity() ||
68  !mapRange.back().isIdentity()) {
69  return false;
70  }
71  // Check yield first block argument.
72  Block *body = op.getBlock();
73  if (body->getOperations().size() != 1)
74  return false;
75  auto yieldOp = dyn_cast<linalg::YieldOp>(body->back());
76  if (!yieldOp || yieldOp.getNumOperands() != 1)
77  return false;
78  return yieldOp->getOperand(0) == body->getArgument(0);
79 }
80 
81 //===----------------------------------------------------------------------===//
82 // FillOpInterface implementation
83 //===----------------------------------------------------------------------===//
84 /// Detects if a linalg.generic operation represents a fill with an inlined
85 /// constant. If so, returns the constant value. Otherwise, returns
86 /// std::nullopt.
87 static std::optional<Value> isaInlinedFillOp(GenericOp op) {
88  if (!op.isAllParallelLoops() || op.getNumDpsInits() != 1 ||
89  op.getNumDpsInputs() != 0)
90  return std::nullopt;
91 
92  // Init should not be referenced.
93  if (op.payloadUsesValueFromOperand(op.getDpsInitOperand(0)))
94  return std::nullopt;
95 
96  Block *body = op.getBody();
97  if (body->getOperations().size() != 1)
98  return std::nullopt;
99 
100  auto yieldOp = dyn_cast<linalg::YieldOp>(body->back());
101  if (!yieldOp || yieldOp.getNumOperands() != 1)
102  return std::nullopt;
103 
104  Value yieldOperand = yieldOp->getOperand(0);
105  if (!yieldOperand.getDefiningOp<arith::ConstantOp>() &&
106  !yieldOperand.getDefiningOp<complex::ConstantOp>())
107  return std::nullopt;
108 
109  return yieldOperand;
110 }
111 
112 /// Detects if a linalg.generic operation represents an external scalar input.
113 /// If so, returns the constant value. Otherwise, returns std::nullopt.
114 static std::optional<Value> isaExternalFillOp(GenericOp op) {
115  // Structural.
116  if (!op.isAllParallelLoops() || !op.isSingleInputOutput() ||
117  !op.isSingleYieldOp())
118  return std::nullopt;
119 
120  // Input should be referenced and init should not.
121  if (!op.payloadUsesValueFromOperand(op.getDpsInputOperand(0)) ||
122  op.payloadUsesValueFromOperand(op.getDpsInitOperand(0)))
123  return std::nullopt;
124 
125  OpOperand *value = op.getDpsInputOperand(0);
126  if (!op.isScalar(value))
127  return std::nullopt;
128  return value->get();
129 }
130 
131 std::optional<Value> linalg::isaFillOpInterface(GenericOp op) {
132  if (auto fillVal = isaInlinedFillOp(op))
133  return fillVal;
134  return isaExternalFillOp(op);
135 }
136 
137 //===----------------------------------------------------------------------===//
138 // BroadcastOpInterface implementation
139 //===----------------------------------------------------------------------===//
140 std::optional<SmallVector<int64_t>>
142  // Structural.
143  if (!op.isAllParallelLoops() || !op.isSingleInputOutput() ||
144  !op.isSingleYieldOp())
145  return std::nullopt;
146 
147  auto srcTy = op.getDpsInputOperand(0)->get().getType();
148  auto dstTy = op.getDpsInitOperand(0)->get().getType();
149  if (!isa<MemRefType, RankedTensorType>(srcTy) ||
150  !isa<MemRefType, RankedTensorType>(dstTy))
151  return std::nullopt;
152 
153  // Check output is identity map. Broadcast could additionally be
154  // employing permutation of indices and that would be expressible
155  // in linalg.generic but is not expressible for named broadcast op.
156  auto dstMap = op.getIndexingMapsArray()[1];
157  if (!dstMap.isIdentity())
158  return std::nullopt;
159 
160  SmallVector<int64_t> position;
161  auto srcMap = op.getIndexingMapsArray()[0];
162 
163  if (srcMap.getResults().size() >= dstMap.getResults().size())
164  return std::nullopt;
165 
166  // Check input map is monotonically increasing DimIds.
167  for (unsigned i = 0; i < srcMap.getNumResults(); ++i) {
168  auto expr = llvm::dyn_cast<AffineDimExpr>(srcMap.getResults()[i]);
169  if (!expr)
170  return std::nullopt;
171  int64_t pos = expr.getPosition();
172  if (i > 0 && pos <= position[i - 1])
173  return std::nullopt;
174  position.push_back(expr.getPosition());
175  }
176 
177  SmallVector<int64_t> broadcastedDims;
178  auto numDims = srcMap.getNumDims();
179  // This is quadratic but number of items is generally small.
180  for (auto dim : llvm::seq<int64_t>(0, numDims)) {
181  if (!llvm::is_contained(position, dim))
182  broadcastedDims.push_back(dim);
183  }
184  return broadcastedDims;
185 }
186 
187 //===----------------------------------------------------------------------===//
188 // TransposeOpInterface implementation
189 //===----------------------------------------------------------------------===//
190 std::optional<SmallVector<int64_t>>
192  // To specialize as a transpose op, the genericOp must be
193  // all parallel loops, single input, single output, and its body
194  // should be just a yield op, yielding input as output as is (no compute).
195  if (!op.isAllParallelLoops() || !op.isSingleInputOutput() ||
196  !op.isSingleYieldOp())
197  return std::nullopt;
198 
199  auto mapRange = op.getIndexingMapsArray();
200  if (mapRange.size() != 2)
201  return std::nullopt;
202 
203  auto mapOfInput = mapRange.front();
204  auto mapOfResult = mapRange.back();
205 
206  // linalg.transpose permutes the dimensions of input using this
207  // rule: dim(result, i) = dim(input, permutation[i])
208  if (!mapOfResult.isIdentity() || !mapOfInput.isPermutation())
209  return std::nullopt;
210 
211  SmallVector<int64_t> permutation(mapOfInput.getNumDims());
212  for (unsigned i = 0; i < mapOfInput.getNumDims(); ++i) {
213  auto expr = llvm::cast<AffineDimExpr>(mapOfInput.getResults()[i]);
214  permutation[expr.getPosition()] = i;
215  }
216  return permutation;
217 }
218 
219 //===----------------------------------------------------------------------===//
220 // Elementwise Single Unary/Binary-OpInterface implementation
221 //===----------------------------------------------------------------------===//
222 static bool isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op,
223  unsigned arity) {
224  // Check all loops are parallel.
225  if (!op.isAllParallelLoops() || op.getNumLoops() < 1)
226  return false;
227 
228  // Check there are arity-inputs, 1-output and all are identity-maps.
229  if (op.getNumDpsInputs() != arity || op.getNumDpsInits() != 1 ||
230  !llvm::all_of(op.getIndexingMapsArray(),
231  [](AffineMap map) { return map.isIdentity(); }))
232  return false;
233 
234  // Init should not be referenced for elementwise operations.
235  if (op.payloadUsesValueFromOperand(op.getDpsInitOperand(0)))
236  return false;
237 
238  // A linalg.generic could be series of elementwise ops e.g. exp(neg(x)) such
239  // as resulting from producer-consumer fusion. Here, we restrict to two ops in
240  // the body, where the first is the elementwise single op and the second a
241  // yield.
242  Block *body = op.getBody();
243  if (body->getOperations().size() != 2)
244  return false;
245 
246  Operation *oper = &body->front();
247  if (oper->getNumOperands() != arity || oper->getNumResults() != 1)
248  return false;
249 
250  auto yieldOp = dyn_cast<linalg::YieldOp>(body->back());
251  return !(!yieldOp || yieldOp.getNumOperands() != 1 ||
252  yieldOp->getOperand(0).getDefiningOp() != oper);
253 }
254 
255 bool linalg::isaElemwiseSingleUnaryOpInterface(linalg::GenericOp op) {
256  // All basic elemwise checks.
258  return false;
259 
260  // Check input is actully used.
261  if (!op.payloadUsesValueFromOperand(op.getDpsInputOperand(0)))
262  return false;
263  return true;
264 }
265 
266 bool linalg::isaElemwiseSingleBinaryOpInterface(linalg::GenericOp op) {
268  return false;
269 
270  // Check both inputs are used (elementwise).
271  OpOperand *inputOpOperand0 = op.getDpsInputOperand(0);
272  OpOperand *inputOpOperand1 = op.getDpsInputOperand(1);
273  return !(!op.payloadUsesValueFromOperand(inputOpOperand0) ||
274  !op.payloadUsesValueFromOperand(inputOpOperand1));
275 }
276 
277 //===----------------------------------------------------------------------===//
278 // ContractionOpInterface implementation
279 //===----------------------------------------------------------------------===//
280 
281 /// If the value is defined by a chain of unary side effect-free, go up the
282 /// use-def chain until the first value that isn't defined by such an op.
283 // TODO: relax to multi-operands with constants, which are technically unary ops
284 // as needed (e.g. add5).
286  Operation *op = value.getDefiningOp();
287  while (op && op->getNumOperands() == 1) {
288  auto iface = dyn_cast<MemoryEffectOpInterface>(op);
289  if (!iface || !iface.hasNoEffect())
290  break;
291  value = op->getOperand(0);
292  op = value.getDefiningOp();
293  }
294  return value;
295 }
296 
298  Block &block, function_ref<bool(Operation *, Operation *)> isaPair,
299  llvm::raw_ostream &errs) {
300  if (block.empty() || !block.back().mightHaveTrait<OpTrait::IsTerminator>()) {
301  errs << "no terminator in the block";
302  return false;
303  }
304 
305  if (block.getNumArguments() != 3) {
306  errs << "expected block with 3 arguments";
307  return false;
308  }
309 
310  Operation *terminator = block.getTerminator();
311  if (terminator->getNumOperands() != 1) {
312  errs << "expected terminator with 1 operand";
313  return false;
314  }
315 
316  Value yielded = getSourceSkipUnary(terminator->getOperand(0));
317  Operation *reductionOp = yielded.getDefiningOp();
318  if (reductionOp->getNumResults() != 1 || reductionOp->getNumOperands() != 2) {
319  errs << "expected reduction op to be binary";
320  return false;
321  }
322 
323  Value reductionLHS = getSourceSkipUnary(reductionOp->getOperand(0));
324  Value reductionRHS = getSourceSkipUnary(reductionOp->getOperand(1));
325 
326  if (reductionLHS != block.getArgument(2) &&
327  reductionRHS != block.getArgument(2)) {
328  errs << "expected reduction to take block argument #2 as one of the "
329  "operands (modulo unary casts)";
330  return false;
331  }
332 
333  Value contributed = getSourceSkipUnary(
334  isa<BlockArgument>(reductionLHS) ? reductionRHS : reductionLHS);
335  Operation *elementwiseOp = contributed.getDefiningOp();
336  if (!elementwiseOp || elementwiseOp->getNumResults() != 1 ||
337  elementwiseOp->getNumOperands() != 2) {
338  errs << "expected elementwise op to be binary";
339  return false;
340  }
341 
342  if (!isaPair(elementwiseOp, reductionOp)) {
343  errs << "expected reduction/elementwise op kind not satisfied";
344  return false;
345  }
346 
347  Value elementwiseLHS = getSourceSkipUnary(elementwiseOp->getOperand(0));
348  Value elementwiseRHS = getSourceSkipUnary(elementwiseOp->getOperand(1));
349  if ((elementwiseLHS == block.getArgument(0) &&
350  elementwiseRHS == block.getArgument(1)) ||
351  (elementwiseLHS == block.getArgument(1) &&
352  elementwiseRHS == block.getArgument(0))) {
353  return true;
354  }
355 
356  errs << "expected elementwise op to apply to block arguments (modulo unary "
357  "casts)";
358  return false;
359 }
360 
361 /// Returns true if the two operations are of the kinds specified by a pair of
362 /// consecutive template arguments.
363 template <typename AddOpTy, typename MulOpTy, typename... Args>
365  static_assert(sizeof...(Args) % 2 == 0,
366  "expected an even number of template arguments");
367  if (isa<AddOpTy>(add) && isa<MulOpTy>(mul))
368  return true;
369 
370  if constexpr (sizeof...(Args) > 0)
371  return isPairTemplateImpl<Args...>(add, mul);
372  else
373  return false;
374 }
375 
376 /// Returns true if the block is a body of a contraction with the kinds of
377 /// operations given pairwise by template arguments.
378 template <typename... Args>
379 static bool isContractionBody(Block &block) {
380  return linalg::detail::isContractionBody(block, &isPairTemplateImpl<Args...>);
381 }
382 
383 /// Given an `indexingMap` and its corresponding `iterators`, returns
384 /// the positions of the iterators of type `iter` that are indexed by
385 /// the `indexingMap` as a permutation. This is useful to infer various
386 /// subcomputations on a `LinalgOp`. This is performed by looking up
387 /// each result in the `indexingMap` and determining whether:
388 /// - It is a single AffineDimExpr.
389 /// - It is the only result involving this AffineDimExpr.
390 static llvm::SmallDenseSet<int64_t>
393  utils::IteratorType iter) {
394  assert(iterators.size() == indexingMap.getNumDims());
395  llvm::SmallDenseSet<int64_t> res;
396  for (AffineExpr e : indexingMap.getResults()) {
397  if (auto d = dyn_cast<AffineDimExpr>(e)) {
398  if (iterators[d.getPosition()] == iter &&
399  llvm::count_if(indexingMap.getResults(), [d](AffineExpr e) {
400  return e.isFunctionOfDim(d.getPosition());
401  }) == 1)
402  res.insert(d.getPosition());
403  }
404  }
405  return res;
406 }
407 
408 namespace {
409 auto par = utils::IteratorType::parallel;
410 auto red = utils::IteratorType::reduction;
411 } // namespace
412 
413 /// Infer the iterator types from the init affine map. This looks at which dims
414 /// are present in the map results, and returns an iterator types array with
415 /// parallel types for dims that are present, and reduction types for dims that
416 /// are not present.
417 static FailureOr<SmallVector<utils::IteratorType>>
419  if (!map.isProjectedPermutation())
420  return failure();
421  SmallVector<utils::IteratorType> iterators(map.getNumDims(), red);
422  for (auto expr : map.getResults())
423  if (auto dim = dyn_cast<AffineDimExpr>(expr))
424  iterators[dim.getPosition()] = par;
425  return iterators;
426 }
427 
428 /// Find 2 parallel (m and n) and 1 reduction (k) dimension candidates that form
429 /// a matmul subcomputation within `linalgOp`. These dimensions are such that:
430 /// 1. The m dimension is involved in an outer-product along LHS
431 /// (i.e. it is a permutation on RES and LHS and does not appear in RHS).
432 /// 2. The n dimension is involved in an outer-product along RHS
433 /// (i.e. it is a permutation on RES and RHS and does not appear in LHS).
434 /// 3. The k dimension appears as a permutation on LHS and RHS.
435 /// 4. m, n and k appear only once in any given indexing.
436 /// 5. Optional batch dimensions that appear in all operands are captured.
437 /// This allows e.g. detecting that some contraction is embedded within
438 /// `linalgOp` with some orthogonal heuristic.
439 static FailureOr<ContractionDimensions>
441  ArrayRef<utils::IteratorType> iterators) {
442  llvm::SmallDenseSet<int64_t> a =
443  findPermutationsIndexingOperand(indexingMaps[0], iterators, par);
444  llvm::SmallDenseSet<int64_t> b =
445  findPermutationsIndexingOperand(indexingMaps[1], iterators, par);
446  llvm::SmallDenseSet<int64_t> c =
447  findPermutationsIndexingOperand(indexingMaps[2], iterators, par);
448 
449  // A & C - B are the iterators involved in an outer-product along A (the LHS).
450  llvm::SmallDenseSet<int64_t> ac = a;
451  llvm::set_intersect(ac, c);
452  llvm::set_subtract(ac, b);
453  // B & C - A are the iterators involved in an outer-product along B (the RHS).
454  llvm::SmallDenseSet<int64_t> bc = b;
455  llvm::set_intersect(bc, c);
456  llvm::set_subtract(bc, a);
457  // A & B & C are the "batch" dimensions.
458  llvm::SmallDenseSet<int64_t> batches = a;
459  llvm::set_intersect(batches, b);
460  llvm::set_intersect(batches, c);
461 
462  // A & B red are the reduction dimensions.
463  llvm::SmallDenseSet<int64_t> ra =
464  findPermutationsIndexingOperand(indexingMaps[0], iterators, red);
465  llvm::SmallDenseSet<int64_t> rb =
466  findPermutationsIndexingOperand(indexingMaps[1], iterators, red);
467  llvm::set_intersect(ra, rb);
468 
469  // Return each set in sorted order.
470  ContractionDimensions dimensions{
471  SmallVector<unsigned, 2>(batches.begin(), batches.end()),
472  SmallVector<unsigned, 2>(ac.begin(), ac.end()),
473  SmallVector<unsigned, 2>(bc.begin(), bc.end()),
474  SmallVector<unsigned, 2>(ra.begin(), ra.end())};
475  llvm::sort(dimensions.batch);
476  llvm::sort(dimensions.m);
477  llvm::sort(dimensions.n);
478  llvm::sort(dimensions.k);
479  return dimensions;
480 }
481 
482 FailureOr<ContractionDimensions>
484  if (linalgOp.getNumDpsInits() != 1 || linalgOp.getNumDpsInputs() != 2)
485  return failure();
486  return inferContractionDimsImpl(linalgOp.getIndexingMapsArray(),
487  linalgOp.getIteratorTypesArray());
488 }
489 
490 FailureOr<ContractionDimensions>
492  if (indexingMaps.size() != 3)
493  return failure();
494  auto iterators = inferIteratorsFromOutMap(indexingMaps[2]);
495  if (failed(iterators))
496  return failure();
497  return inferContractionDimsImpl(indexingMaps, iterators.value());
498 }
499 
500 namespace mlir::linalg::detail {
502  Success = 0,
503  NotLinalgOp,
505  NoReduction,
507  NotAddMul
508 };
509 } // namespace mlir::linalg::detail
510 
514  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
515  if (!linalgOp)
516  return MatchContractionResult::NotLinalgOp;
517  if (linalgOp.getNumDpsInputs() != 2 || linalgOp.getNumDpsInits() != 1)
518  return MatchContractionResult::WrongNumOperands;
519  auto mapRange = linalgOp.getIndexingMapsArray();
520  if (linalgOp.getNumReductionLoops() == 0)
521  return MatchContractionResult::NoReduction;
522  if (llvm::any_of(mapRange,
523  [](AffineMap m) { return !m.isProjectedPermutation(); }))
524  return MatchContractionResult::NotProjectedPermutations;
525  // TODO: more fields than add/mul.
526  // clang-format off
527  if (!::isContractionBody<
528  arith::MulFOp, arith::AddFOp,
529  arith::MulIOp, arith::AddIOp,
530  complex::MulOp, complex::AddOp,
531  arith::AndIOp, arith::OrIOp>(
532  *linalgOp.getBlock())) {
533  return MatchContractionResult::NotAddMul;
534  }
535  // clang-format on
536 
537  if (dimensions) {
538  FailureOr<ContractionDimensions> res = inferContractionDims(linalgOp);
539  assert(succeeded(res) && "unexpected failure to infer contraction dims");
540  *dimensions = *res;
541  }
542  return MatchContractionResult::Success;
543 }
544 
545 StringRef
547  switch (res) {
548  case MatchContractionResult::NotLinalgOp:
549  return "expected a LinalgOp";
550  case MatchContractionResult::WrongNumOperands:
551  return "expected op with 2 inputs and 1 output";
552  case MatchContractionResult::NoReduction:
553  return "expected at least 1 reduction";
554  case MatchContractionResult::NotProjectedPermutations:
555  return "expected indexing maps to be projected permutations";
556  case MatchContractionResult::NotAddMul:
557  return "expected add/mul op in the body";
558  case MatchContractionResult::Success:
559  return "";
560  }
561  llvm_unreachable("unhandled MatchContractionResult case");
562 }
563 
565  if (!linalgOp)
566  return false;
567  Operation *op = linalgOp.getOperation();
568  return isa<ContractionOpInterface>(op) ||
571 }
572 
573 /// Verify that a LinalgOp `op` is a contraction.
574 /// A Linalg contraction is defined in general terms:
575 /// 1. Has 2 input and 1 output shapes.
576 /// 2. Has at least one reduction dimension.
577 /// 3. Has only projected permutation indexing maps.
578 /// 4. its body computes `u5(u1(c) + u2(u3(a) * u4(b)))` on some field
579 /// (AddOpType, MulOpType), where u1, u2, u3, u4 and u5 represent scalar unary
580 /// operations that may change the type (e.g. for mixed-precision).
581 /// As a consequence, when vectorization of such an op occurs, the only special
582 /// behavior is that the (unique) MulOpType is vectorized into a
583 /// `vector.contract`. All other ops are handled in a generic fashion.
584 /// In the future, we may wish to allow more input arguments and elementwise and
585 /// constant operations that do not involve the reduction dimension(s).
587  auto res = isContractionInterfaceImpl(op);
588  if (res != MatchContractionResult::Success)
589  return op->emitError(getMatchContractionMessage(res));
590  return success();
591 }
592 
593 //===----------------------------------------------------------------------===//
594 // ConvolutionOpInterface implementation
595 //===----------------------------------------------------------------------===//
596 
597 /// Of the given two expressions returns one that is of type T (`lhs` gets
598 /// preference over `rhs`)
599 template <typename T>
601  return isa<T>(lhs) ? cast<T>(lhs) : (isa<T>(rhs) ? cast<T>(rhs) : nullptr);
602 }
603 
604 namespace {
605 /// Walk the indexing expressions for input of a convolution operation to verify
606 /// its of the right form, either
607 /// - AffineDimExpr
608 /// - AffineDimExpr (`*` (AffineSymbolExpr | AffineConstantExpr))?
609 /// (`+` AffineDimExpr (`*` (AffineSymbolExpr | AffineConstantExpr))?)*
610 ///
611 /// classifies the AffineDimExpr as convolved dimensions or unconvolved
612 /// dimensions and verifies each dimension occurs only once.
613 struct ConvAccessExprWalker
614  : public AffineExprVisitor<ConvAccessExprWalker, LogicalResult> {
615  // Stores dimensions used in expressions of the above form.
616  llvm::SmallDenseSet<int64_t> convolvedDims;
617  // Stores the dual mapping between LHS and RHS of convolution exprs.
618  llvm::SmallDenseMap<int64_t, int64_t> convolvedDimMapping;
619  // Stores single use dimensions used by an AffineDimExpr.
620  llvm::SmallDenseSet<int64_t> unConvolvedDims;
621  // Stores a mapping from convolved dims to their coefficient.
622  llvm::SmallDenseMap<int64_t, AffineExpr> strideAndDilationMapping;
623 
624  // Removes dims with multiple uses in the source input map from dimension
625  // sets tracked by this walker.
626  void clearMultiUseDims(AffineMap map) {
627  for (int dimPos = 0, e = map.getNumDims(); dimPos < e; ++dimPos) {
628  if (llvm::count_if(map.getResults(), [dimPos](AffineExpr e) {
629  return e.isFunctionOfDim(dimPos);
630  }) > 1) {
631  convolvedDims.erase(dimPos);
632  unConvolvedDims.erase(dimPos);
633  // If a duplicate dim is marked as convolved, the pair of the duplicate
634  // dim must be removed from the map as well.
635  auto it = convolvedDimMapping.find(dimPos);
636  if (it != convolvedDimMapping.end()) {
637  int64_t pairedDim = it->second;
638  convolvedDims.erase(pairedDim);
639  unConvolvedDims.erase(pairedDim);
640  strideAndDilationMapping.erase(pairedDim);
641  convolvedDimMapping.erase(dimPos);
642  convolvedDimMapping.erase(pairedDim);
643  }
644  }
645  }
646  }
647 
648  LogicalResult visitDimExpr(AffineDimExpr dimExpr) {
649  unsigned position = dimExpr.getPosition();
650  if (unConvolvedDims.count(position) || convolvedDims.count(position)) {
651  return failure();
652  }
653  unConvolvedDims.insert(position);
654  return success();
655  }
656 
657  LogicalResult visitSymbolExpr(AffineSymbolExpr expr) { return failure(); }
658 
659  LogicalResult visitConstantExpr(AffineConstantExpr expr) { return failure(); }
660 
661  LogicalResult visitAffineBinaryOpExpr(AffineBinaryOpExpr binaryExpr) {
662  // In pre-order visit, top level op has to be an add op.
663  if (binaryExpr.getKind() != AffineExprKind::Add)
664  return failure();
665  auto lhsDimPos = getDimExprOrMulExprDimPos(binaryExpr.getLHS());
666  auto rhsDimPos = getDimExprOrMulExprDimPos(binaryExpr.getRHS());
667  if (failed(lhsDimPos) || failed(rhsDimPos))
668  return failure();
669  convolvedDimMapping[*lhsDimPos] = *rhsDimPos;
670  convolvedDimMapping[*rhsDimPos] = *lhsDimPos;
671  return success();
672  }
673 
674  FailureOr<int64_t> getDimExprOrMulExprDimPos(AffineExpr expr) {
675  if (auto dimExpr = dyn_cast<AffineDimExpr>(expr)) {
676  int64_t dim = dimExpr.getPosition();
677  if (convolvedDims.count(dim) || unConvolvedDims.count(dim))
678  return failure();
679  // Stride/dilation for this dim is implicitly 1.
680  strideAndDilationMapping[dim] =
682  convolvedDims.insert(dim);
683  return dim;
684  }
685  if (auto symbolMulExpr = dyn_cast<AffineBinaryOpExpr>(expr)) {
686  if (symbolMulExpr.getKind() != AffineExprKind::Mul)
687  return failure();
688  auto lhsExpr = symbolMulExpr.getLHS();
689  auto rhsExpr = symbolMulExpr.getRHS();
690  // Check for symbol expression.
691  AffineExpr mulExpr =
692  getAffineExprOfType<AffineSymbolExpr>(lhsExpr, rhsExpr);
693  // If there was no symbol expr, check for constant expression.
694  if (!mulExpr) {
695  mulExpr = getAffineExprOfType<AffineConstantExpr>(lhsExpr, rhsExpr);
696  }
697  auto dimExpr = getAffineExprOfType<AffineDimExpr>(lhsExpr, rhsExpr);
698  if (!mulExpr || !dimExpr)
699  return failure();
700  int64_t dim = dimExpr.getPosition();
701  if (convolvedDims.count(dim) || unConvolvedDims.count(dim))
702  return failure();
703  strideAndDilationMapping[dim] = mulExpr;
704  convolvedDims.insert(dim);
705  return dim;
706  }
707  return failure();
708  }
709 };
710 } // namespace
711 
712 static llvm::SmallDenseSet<int64_t> getPreservedDims(AffineMap map) {
713  assert(map.isProjectedPermutation() &&
714  "expected map to have projected permutations");
715  llvm::SmallDenseSet<int64_t> preservedDims;
716  for (auto expr : map.getResults())
717  preservedDims.insert(cast<AffineDimExpr>(expr).getPosition());
718  return preservedDims;
719 }
720 
724  for (auto e : exprs) {
725  auto constantExpr = dyn_cast<AffineConstantExpr>(e);
726  assert(constantExpr && "Found non-constant stride/dilation");
727  vals.push_back(constantExpr.getValue());
728  }
729  return vals;
730 }
731 
732 /// Classifies dimensions in the `linalgOp` used by a convolution
733 /// subcomputation, as captured by `inputExprWalker`. If
734 /// `allowEmptyConvolvedDims` is not set this this will fail if there is not
735 /// at least convolved dimension pair (output image + filter loop). Convolution
736 /// dimensions are specified in sorted order, and strides match the order of
737 /// the filter loop dimensions, while the dilations match the order of the
738 /// output image dimensions.
739 static FailureOr<ConvolutionDimensions>
740 inferConvolutionDimsImpl(LinalgOp linalgOp,
741  ConvAccessExprWalker &inputExprWalker,
742  bool allowEmptyConvolvedDims) {
743  auto filterMap =
744  linalgOp.getMatchingIndexingMap(linalgOp.getDpsInputOperand(1));
745  auto outputMap =
746  linalgOp.getMatchingIndexingMap(linalgOp.getDpsInitOperand(0));
747  llvm::SmallDenseSet<int64_t> filterDims = findPermutationsIndexingOperand(
748  filterMap, linalgOp.getIteratorTypesArray(), par);
749  llvm::SmallDenseSet<int64_t> outputDims = findPermutationsIndexingOperand(
750  outputMap, linalgOp.getIteratorTypesArray(), par);
751 
752  // unConvolvedDims & outputDims - filterDims are the batch iterators.
753  llvm::SmallDenseSet<int64_t> batch = inputExprWalker.unConvolvedDims;
754  llvm::set_intersect(batch, outputDims);
755  llvm::set_subtract(batch, filterDims);
756 
757  // convolvedDims & outputDims are the output image iterators.
758  llvm::SmallDenseSet<int64_t> oi = inputExprWalker.convolvedDims;
759  llvm::set_intersect(oi, outputDims);
760 
761  // filterDims & outputDims - unConvolvedDims are the output channel iterators.
762  llvm::SmallDenseSet<int64_t> oc = filterDims;
763  llvm::set_intersect(oc, outputDims);
764  llvm::set_subtract(oc, inputExprWalker.unConvolvedDims);
765 
766  // filterDims & outputDims & unConvolvedDims are the depth iterators.
767  llvm::SmallDenseSet<int64_t> depth = filterDims;
768  llvm::set_intersect(depth, outputDims);
769  llvm::set_intersect(depth, inputExprWalker.unConvolvedDims);
770 
771  llvm::SmallDenseSet<int64_t> filterReducedDims =
773  linalgOp.getIteratorTypesArray(), red);
774 
775  // convolvedDims & filterReducedDims are the filter loop iterators.
776  llvm::SmallDenseSet<int64_t> fl = inputExprWalker.convolvedDims;
777  llvm::set_intersect(fl, filterReducedDims);
778 
779  // unConvolvedDims & filterReducedDims are the input channel iterators.
780  llvm::SmallDenseSet<int64_t> ic = inputExprWalker.unConvolvedDims;
781  llvm::set_intersect(ic, filterReducedDims);
782 
783  if (oi.empty() && !allowEmptyConvolvedDims)
784  return failure();
785 
786  // Return each set in sorted order.
787  ConvolutionDimensions dimensions{
788  SmallVector<unsigned, 2>(batch.begin(), batch.end()),
789  SmallVector<unsigned, 2>(oi.begin(), oi.end()),
790  SmallVector<unsigned, 2>(oc.begin(), oc.end()),
791  SmallVector<unsigned, 2>(fl.begin(), fl.end()),
792  SmallVector<unsigned, 2>(ic.begin(), ic.end()),
793  SmallVector<unsigned, 2>(depth.begin(), depth.end()),
794  /*strides=*/SmallVector<int64_t, 2>{},
795  /*dilations=*/SmallVector<int64_t, 2>{}};
796  llvm::sort(dimensions.batch);
797  llvm::sort(dimensions.outputImage);
798  llvm::sort(dimensions.outputChannel);
799  llvm::sort(dimensions.filterLoop);
800  llvm::sort(dimensions.inputChannel);
801  llvm::sort(dimensions.depth);
802 
803  // Use the op carried strides/dilations attribute if present.
804  auto nativeStrides = linalgOp->getAttrOfType<DenseIntElementsAttr>("strides");
805  if (!nativeStrides) {
806  SmallVector<AffineExpr, 2> strideExprs;
807  for (unsigned oiDim : dimensions.outputImage)
808  strideExprs.push_back(inputExprWalker.strideAndDilationMapping[oiDim]);
809  dimensions.strides = getConstantsFromExprList(strideExprs);
810  } else {
811  dimensions.strides = llvm::to_vector<2>(nativeStrides.getValues<int64_t>());
812  }
813  auto nativeDilations =
814  linalgOp->getAttrOfType<DenseIntElementsAttr>("dilations");
815  if (!nativeDilations) {
816  SmallVector<AffineExpr, 2> dilationExprs;
817  for (unsigned flDim : dimensions.filterLoop)
818  dilationExprs.push_back(inputExprWalker.strideAndDilationMapping[flDim]);
819  dimensions.dilations = getConstantsFromExprList(dilationExprs);
820  } else {
821  dimensions.dilations =
822  llvm::to_vector<2>(nativeDilations.getValues<int64_t>());
823  }
824  return dimensions;
825 }
826 
827 /// Find at least 1 parallel (output_image) and reduction (filter_loop)
828 /// dimension candidates that form a convolution subcomputation within
829 /// `linalgOp`. The LHS is assumed to be the convolution input while the
830 /// RHS is assumed as the filter.
831 /// These dimensions are such that:
832 /// 1. Optional batch dimensions that appear in the input and filter.
833 /// 2. The output_image dimension is involved in a cross-correlation along LHS
834 /// (i.e. it is a permutation on RES and LHS and has an associated
835 /// filter_loop in RHS).
836 /// 3. Optional output_channel dimension is involved in an outer-product along
837 /// RHS (i.e. it is a permutation on RES and RHS and does not appear in
838 /// LHS).
839 /// 4. Optional input_channel dimension appears as a permutation on LHS and
840 /// RHS.
841 /// 5. The filter_loop dimension appears as a permutation on the RHS and
842 /// represents the shape of the kernel cross-correlated along a
843 /// corresponding output_image dim.
844 /// 6. The input_channel dimension appears as a permutation on LHS and RHS.
845 /// 7. All dimensions appear only once in any given indexing map.
846 /// This allows e.g. detecting that some convolution is embedded within
847 /// `linalgOp` with some orthogonal heuristic.
848 /// When multiple dimension occurrences exist that match any classification
849 /// indices are returned in sorted order.
850 /// Returns a failure if `output_image` (and implicitly `filter_loop`) is empty.
851 FailureOr<ConvolutionDimensions>
853  if (linalgOp.getNumDpsInits() != 1 || linalgOp.getNumDpsInputs() != 2)
854  return failure();
855 
856  auto indexingMaps = linalgOp.getIndexingMapsArray();
857 
858  // Check the input indexing map has the right form.
859  ConvAccessExprWalker inputExprWalker;
860  for (AffineExpr expr : indexingMaps[0].getResults())
861  (void)inputExprWalker.visit(expr);
862  inputExprWalker.clearMultiUseDims(indexingMaps[0]);
863 
864  return inferConvolutionDimsImpl(linalgOp, inputExprWalker,
865  /*allowEmptyConvolvedDims=*/false);
866 }
867 
868 namespace mlir::linalg::detail {
870  Success = 0,
871  NotLinalgOp,
879 };
880 } // namespace mlir::linalg::detail
881 
884  Operation *op, ConvolutionDimensions *dimensions,
885  bool allowEmptyConvolvedDims) {
886  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
887  if (!linalgOp)
888  return MatchConvolutionResult::NotLinalgOp;
889  if (linalgOp.getNumDpsInputs() < 2 || linalgOp.getNumDpsInits() != 1)
890  return MatchConvolutionResult::WrongNumOperands;
891 
892  auto indexingMaps = linalgOp.getIndexingMapsArray();
893 
894  // Check the input indexing map has the right form.
895  ConvAccessExprWalker inputExprWalker;
896  if (llvm::any_of(indexingMaps[0].getResults(),
897  [&inputExprWalker](AffineExpr expr) {
898  return failed(inputExprWalker.visit(expr));
899  })) {
900  return MatchConvolutionResult::WrongInputIndexingMap;
901  }
902 
903  // Filter and output maps must be projected permutation.
904  if (!indexingMaps[1].isProjectedPermutation() ||
905  !indexingMaps.back().isProjectedPermutation())
906  return MatchConvolutionResult::NotProjectedPermutations;
907 
908  auto iteratorTypes = linalgOp.getIteratorTypesArray();
909 
910  llvm::SmallDenseSet<int64_t> outputDims =
911  getPreservedDims(indexingMaps.back());
912  llvm::SmallDenseSet<int64_t> filterDims = getPreservedDims(indexingMaps[1]);
913  // Make sure all loops are characterized as one of:
914  // - Batch loop : present in output, as non-convolved in input, not present in
915  // filter.
916  // - Output image dimension : present in output, convolved dims in input, not
917  // present in filter.
918  // - Output channel dimension : present in output, not present in input,
919  // present in filter.
920  // - Filter loop dimension : present in filter, convolved in input, not
921  // present in output.
922  // - Input channel dimension : unconvolved in input, not present in output,
923  // present in filter.
924  // - Depth multiplier : unconvolved in input, present in output, present in
925  // filter.
926  llvm::SmallDenseSet<int64_t> allLoopDims;
927  for (auto outputExpr : indexingMaps.back().getResults()) {
928  int64_t outputDim = cast<AffineDimExpr>(outputExpr).getPosition();
929  if (inputExprWalker.unConvolvedDims.count(outputDim) &&
930  !filterDims.count(outputDim)) {
931  // Batch dimension.
932  if (iteratorTypes[outputDim] != utils::IteratorType::parallel)
933  return MatchConvolutionResult::OutputDimsNotParallel;
934  allLoopDims.insert(outputDim);
935  continue;
936  }
937  if (inputExprWalker.convolvedDims.count(outputDim) &&
938  !filterDims.count(outputDim)) {
939  // Output image Loop dimension.
940  if (iteratorTypes[outputDim] != utils::IteratorType::parallel)
941  return MatchConvolutionResult::OutputDimsNotParallel;
942  allLoopDims.insert(outputDim);
943  continue;
944  }
945  if (!inputExprWalker.convolvedDims.count(outputDim) &&
946  !inputExprWalker.unConvolvedDims.count(outputDim) &&
947  filterDims.count(outputDim)) {
948  // Output channel dimension.
949  if (iteratorTypes[outputDim] != utils::IteratorType::parallel)
950  return MatchConvolutionResult::OutputDimsNotParallel;
951  allLoopDims.insert(outputDim);
952  continue;
953  }
954  if (inputExprWalker.unConvolvedDims.count(outputDim) &&
955  filterDims.count(outputDim)) {
956  // Depth multiplier.
957  if (iteratorTypes[outputDim] != utils::IteratorType::parallel)
958  return MatchConvolutionResult::OutputDimsNotParallel;
959  allLoopDims.insert(outputDim);
960  continue;
961  }
962  return MatchConvolutionResult::NonConvolutionLoop;
963  }
964  for (auto filterExpr : indexingMaps[1].getResults()) {
965  int64_t filterDim = cast<AffineDimExpr>(filterExpr).getPosition();
966  if (outputDims.count(filterDim) &&
967  !inputExprWalker.unConvolvedDims.count(filterDim) &&
968  !inputExprWalker.convolvedDims.count(filterDim)) {
969  // Output channel dimension. This is already seen, continue;
970  continue;
971  }
972  if (inputExprWalker.convolvedDims.count(filterDim) &&
973  !outputDims.count(filterDim)) {
974  // Filter loop dimension.
975  if (iteratorTypes[filterDim] != utils::IteratorType::reduction)
976  return MatchConvolutionResult::NonOutputDimNotReduction;
977  if (allLoopDims.count(filterDim))
978  return MatchConvolutionResult::NonConvolutionLoop;
979  allLoopDims.insert(filterDim);
980  continue;
981  }
982  if (inputExprWalker.unConvolvedDims.count(filterDim) &&
983  !outputDims.count(filterDim)) {
984  // Input channel dimension.
985  if (iteratorTypes[filterDim] != utils::IteratorType::reduction)
986  return MatchConvolutionResult::NonOutputDimNotReduction;
987  if (allLoopDims.count(filterDim))
988  return MatchConvolutionResult::NonConvolutionLoop;
989  allLoopDims.insert(filterDim);
990  continue;
991  }
992  if (inputExprWalker.unConvolvedDims.count(filterDim) &&
993  outputDims.count(filterDim)) {
994  // Depthwise loop. Already seen.
995  continue;
996  }
997  return MatchConvolutionResult::NonConvolutionLoop;
998  }
999  // All loops must be covered now.
1000  if (allLoopDims.size() != linalgOp.getNumLoops())
1001  return MatchConvolutionResult::NonConvolutionLoop;
1002 
1003  if (!allowEmptyConvolvedDims && inputExprWalker.convolvedDims.empty())
1004  return MatchConvolutionResult::EmptyConvolvedDims;
1005 
1006  if (dimensions) {
1007  FailureOr<ConvolutionDimensions> res = inferConvolutionDimsImpl(
1008  linalgOp, inputExprWalker, allowEmptyConvolvedDims);
1009  assert(succeeded(res) && "unexpected failure to infer convolution dims");
1010  *dimensions = *res;
1011  }
1012 
1013  return MatchConvolutionResult::Success;
1014 }
1015 
1016 StringRef
1018  switch (res) {
1019  case MatchConvolutionResult::NotLinalgOp:
1020  return "expected a LinalgOp";
1021  case MatchConvolutionResult::WrongNumOperands:
1022  return "expected op with 2 inputs and 1 output";
1023  case MatchConvolutionResult::WrongInputIndexingMap:
1024  return "unexpected input index map for convolutions";
1025  case MatchConvolutionResult::NotProjectedPermutations:
1026  return "expected output/filter indexing maps to be projected permutations";
1027  case MatchConvolutionResult::NonConvolutionLoop:
1028  return "unexpected loop dimension for convolution op";
1029  case MatchConvolutionResult::OutputDimsNotParallel:
1030  return "expected all iterators used to access outputs to be parallel";
1031  case MatchConvolutionResult::NonOutputDimNotReduction:
1032  return "expected all iterators not used to access outputs to be reduction";
1033  case MatchConvolutionResult::EmptyConvolvedDims:
1034  return "expected convolved dim to be non-empty";
1035  case MatchConvolutionResult::Success:
1036  return "";
1037  }
1038  llvm_unreachable("unhandled MatchConvolutionResult case");
1039 }
1040 
1042  bool allowEmptyConvolvedDims) {
1044  linalgOp.getOperation(), nullptr, allowEmptyConvolvedDims) ==
1046 }
1047 
1050  if (res != MatchConvolutionResult::Success)
1051  return op->emitError(getMatchConvolutionMessage(res));
1052  return success();
1053 }
1054 
1055 //===----------------------------------------------------------------------===//
1056 // FillOpInterface implementation
1057 //===----------------------------------------------------------------------===//
1058 
1059 enum class MatchFillResult {
1060  Success = 0,
1061  NotLinalgOp,
1062  WrongNumOperands,
1064 };
1065 
1067  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
1068  if (!linalgOp)
1070  if (linalgOp.getNumDpsInputs() != 1 || linalgOp.getNumDpsInits() != 1)
1072 
1073  OpOperand *value = linalgOp.getDpsInputOperand(0);
1074  if (!linalgOp.isScalar(value))
1076 
1077  return MatchFillResult::Success;
1078 }
1079 
1081  auto res = isFillInterfaceImpl(op);
1082  if (res == MatchFillResult::NotLinalgOp)
1083  return op->emitError("expected a LinalgOp");
1085  return op->emitError("expected op with 1 input and 1 output");
1087  return op->emitError("expected op with scalar input");
1088 
1089  return success();
1090 }
1091 
1092 //===----------------------------------------------------------------------===//
1093 // StructuredOpInterface implementation
1094 //===----------------------------------------------------------------------===//
1095 
1096 SmallVector<OpFoldResult> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,
1097  Location loc) {
1099  for (OpOperand &opOperand : getOperation()->getOpOperands()) {
1100  for (int64_t i = 0, e = getRank(&opOperand); i < e; ++i)
1101  res.push_back(createFoldedDimOp(b, loc, opOperand.get(), i));
1102  }
1103  return res;
1104 }
1105 
1106 SmallVector<int64_t, 4> LinalgOp::createFlatListOfOperandStaticDims() {
1108  assert(!hasDynamicShape() && "expected operands to have static shapes");
1109  for (OpOperand &opOperand : getOperation()->getOpOperands())
1110  llvm::append_range(res, getShape(&opOperand));
1111  return res;
1112 }
1113 
1114 SmallVector<Range, 4> LinalgOp::createLoopRanges(OpBuilder &b, Location loc) {
1115  AffineMap map = getLoopsToShapesMap();
1116  unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
1117  auto viewSizes = createFlatListOfOperandDims(b, loc);
1118  SmallVector<Range, 4> res(numDims);
1119  for (unsigned idx = 0; idx < numRes; ++idx) {
1120  auto result = map.getResult(idx);
1121  if (auto d = dyn_cast<AffineDimExpr>(result)) {
1122  if (res[d.getPosition()].offset)
1123  continue;
1124  res[d.getPosition()] =
1125  Range{b.getIndexAttr(0), viewSizes[idx], b.getIndexAttr(1)};
1126  }
1127  }
1128  return res;
1129 }
1130 
1131 /// Visitor to check if any of the given set of positions from AffineDimExprs
1132 /// are used within an AffineExpr.
1134  : public AffineExprVisitor<HasAffineDimExprVisitor, bool> {
1135  HasAffineDimExprVisitor(llvm::SmallBitVector positions)
1136  : positions(std::move(positions)) {}
1137 
1139  return visit(binaryOpExpr.getLHS()) || visit(binaryOpExpr.getRHS());
1140  }
1141 
1142  bool visitDimExpr(AffineDimExpr dimExpr) {
1143  return positions.test(dimExpr.getPosition());
1144  }
1145 
1146  bool visitConstantExpr(AffineConstantExpr constExpr) { return false; }
1147 
1148  bool visitSymbolExpr(AffineSymbolExpr symbolExpr) { return false; }
1149 
1150 private:
1151  llvm::SmallBitVector positions;
1152 };
1153 
1154 static std::pair<int64_t, int64_t>
1156  int64_t inputRankSum = 0;
1157  int64_t outputRankSum = 0;
1158  for (OpOperand *input : op.getDpsInputOperands())
1159  inputRankSum += op.getRank(input);
1160  for (OpOperand &output : op.getDpsInitsMutable())
1161  outputRankSum += op.getRank(&output);
1162  return {inputRankSum, inputRankSum + outputRankSum};
1163 }
1164 
1165 LogicalResult
1167  ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
1168  // An example that helps understand the logic below.
1169  // Consider the following expression O(i+j, j) += A(i,k) * B(k, j)
1170  // We want to express the shape of dim 0 of O in terms of shape of the inputs.
1171  // This is achieved as follows.
1172  // loopsToShapesMap = (d0, d1, d2) -> (d0, d2, d2, d1, d0 + d1, d1)
1173  // subMapOfResultShapes = (d0, d1, d2) -> (d0 + d1, d1)
1174  // shapesToLoopsMap = (d0, d2, d2, d3, d4, d5) -> (d0, d3, d2)
1175  // resultShapesFromInputShapes = subMapOfResultDim.compose(shapesToLoopMap)
1176  // = (d0, d1, d2, d3, d4, d5) -> (d0 + d1, d1)
1177  AffineMap loopsToShapesMap = getLoopsToShapesMap();
1178 
1179  // Find the position in the above map that represents the shape of the
1180  // result:dim being inferred.
1181  auto resultShapesSubMapPos = getResultsPositionInLoopsToShapeMap(*this);
1182 
1183  /// From loopsToShapesMap extract the submap that represents the shape of the
1184  /// (resultIdx, dim) needed.
1185  AffineMap loopToResultsShapeMap = loopsToShapesMap.getSliceMap(
1186  resultShapesSubMapPos.first,
1187  resultShapesSubMapPos.second - resultShapesSubMapPos.first);
1188  AffineMap resultShapesFromInputShapesMap =
1189  loopToResultsShapeMap.compose(getShapesToLoopsMap());
1190 
1191  // Check that the result dim map does not contain the positions corresponding
1192  // to the outputs.
1193  llvm::SmallBitVector outputDims(resultShapesFromInputShapesMap.getNumDims());
1194  outputDims.set(resultShapesSubMapPos.first, resultShapesSubMapPos.second);
1195  HasAffineDimExprVisitor checkDimExpr(std::move(outputDims));
1196  Location loc = getOperation()->getLoc();
1197  IRRewriter rewriter(b);
1198  SmallVector<OpFoldResult> allResultDimValues =
1200  rewriter, loc, resultShapesFromInputShapesMap,
1201  createFlatListOfOperandDims(b, loc));
1202  int64_t pos = 0;
1203  ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults();
1204  for (OpOperand &opOperand : getDpsInitsMutable()) {
1206  for (int64_t dim : llvm::seq<int64_t>(0, getRank(&opOperand))) {
1207  auto shapedType = llvm::cast<ShapedType>(opOperand.get().getType());
1208  if (!shapedType.isDynamicDim(dim)) {
1209  // Static dim: Return IntegerAttr.
1210  shapes.push_back(b.getIndexAttr(shapedType.getDimSize(dim)));
1211  } else {
1212  // Dynamic dim: Return Value.
1213  OpFoldResult ofr = checkDimExpr.visit(shapeExprs[pos])
1214  ? createOrFoldDimOp(b, loc, opOperand.get(), dim)
1215  : allResultDimValues[pos];
1216  shapes.push_back(getValueOrCreateConstantIndexOp(b, loc, ofr));
1217  }
1218  pos++;
1219  }
1220  reifiedReturnShapes.emplace_back(std::move(shapes));
1221  }
1222  return success();
1223 }
1224 
1225 /// Return the index in the indexingMaps vector that corresponds to this
1226 /// `opOperand`.
1227 int64_t LinalgOp::getIndexingMapIndex(OpOperand *opOperand) {
1228  auto operandNumber = opOperand->getOperandNumber();
1229  auto dpsIface = cast<DestinationStyleOpInterface>(*this->getOperation());
1230  if (!dpsIface.isDpsInput(opOperand))
1231  return operandNumber;
1232  unsigned start = dpsIface.getDpsInits().getBeginOperandIndex();
1233  assert(!dpsIface.isDpsInit(opOperand));
1234  // Account for potential inputs that are not DPS and may not appear in
1235  // `indexingMaps`.
1236  return cast<DestinationStyleOpInterface>(*this->getOperation())
1237  .getNumDpsInputs() +
1238  operandNumber - start;
1239 }
1240 
1242  LinalgOp linalgOp = cast<LinalgOp>(op);
1243  // Mixed tensor/buffer operands are not allowed.
1244  if (!linalgOp.hasPureTensorSemantics() &&
1245  !linalgOp.hasPureBufferSemantics() && op->getNumOperands() > 0)
1246  return op->emitOpError("expected to have pure tensor or buffer semantics");
1247 
1248  // Before checking indexing maps, we need to make sure the attributes
1249  // referenced by it are valid.
1250  if (linalgOp.hasDynamicIndexingMaps())
1251  if (failed(linalgOp.verifyIndexingMapRequiredAttributes()))
1252  return failure();
1253 
1254  // Delayed calling of IndexingMapOpInterface::verifyImpl.
1255  if (failed(cast<IndexingMapOpInterface>(op).verifyImpl()))
1256  return failure();
1257 
1258  // Set this flag if this op has user defined maps. This is required to guard
1259  // the below error condition which assume default indexing maps.
1260  for (OpOperand &opOperand : linalgOp->getOpOperands()) {
1261  AffineMap indexingMap = linalgOp.getMatchingIndexingMap(&opOperand);
1262  // Domain must be consistent.
1263  unsigned numLoops = linalgOp.getNumLoops();
1264  if (indexingMap.getNumDims() != numLoops)
1265  return op->emitOpError("expected indexing_map #")
1266  << opOperand.getOperandNumber() << " to have " << numLoops
1267  << " dim(s) to match the number of loops";
1268  }
1269  SmallVector<unsigned> redDims;
1270  linalgOp.getReductionDims(redDims);
1271 
1272  if (!linalgOp.getShapesToLoopsMap())
1273  return op->emitOpError("expected the shape-to-loops map to be non-null");
1274 
1275  // Check the region has exactly one block.
1276  if (linalgOp->getNumRegions() != 1 || !linalgOp->getRegion(0).hasOneBlock())
1277  return op->emitOpError("expects to have 1 region with 1 block");
1278 
1279  // Simplifying assumption: bbargs match 1-1 with shape operands elemental
1280  // types.
1281  // TODO: once ranked shape types are plugged in, we may want to drop the
1282  // corresponding bbargs, that can never be read from. This will be subject to
1283  // consistency discussions (i.e. what to do with output tensors whose bbarg is
1284  // not used).
1285  Block &block = linalgOp->getRegion(0).front();
1286 
1287  if (linalgOp.getOpOperandsMatchingBBargs().size() != block.getNumArguments())
1288  return op->emitOpError("expected as many non-induction variable region "
1289  "arguments as the number of input/output operands");
1290 
1291  for (OpOperand *opOperand : linalgOp.getOpOperandsMatchingBBargs()) {
1292  Type elementType = opOperand->get().getType();
1293  if (isa<MemRefType, RankedTensorType>(elementType))
1294  elementType = getElementTypeOrSelf(opOperand->get().getType());
1295  Type argType = block.getArgument(opOperand->getOperandNumber()).getType();
1296  if (elementType != argType)
1297  return op->emitOpError("expected type of bb argument #")
1298  << opOperand->getOperandNumber() << " (" << argType << ")"
1299  << " to match element or self type of the corresponding operand ("
1300  << elementType << ")";
1301  }
1302 
1303  return success();
1304 }
static void visit(Operation *op, DenseSet< Operation * > &visited)
Visits all the pdl.operand(s), pdl.result(s), and pdl.operation(s) connected to the given operation.
Definition: PDL.cpp:62
static FailureOr< ConvolutionDimensions > inferConvolutionDimsImpl(LinalgOp linalgOp, ConvAccessExprWalker &inputExprWalker, bool allowEmptyConvolvedDims)
Classifies dimensions in the linalgOp used by a convolution subcomputation, as captured by inputExprW...
static std::optional< Value > isaExternalFillOp(GenericOp op)
Detects if a linalg.generic operation represents an external scalar input.
static Value getSourceSkipUnary(Value value)
If the value is defined by a chain of unary side effect-free, go up the use-def chain until the first...
static T getAffineExprOfType(AffineExpr lhs, AffineExpr rhs)
Of the given two expressions returns one that is of type T (lhs gets preference over rhs)
static bool isPairTemplateImpl(Operation *add, Operation *mul)
Returns true if the two operations are of the kinds specified by a pair of consecutive template argum...
static SmallVector< int64_t, 2 > getConstantsFromExprList(const SmallVector< AffineExpr, 2 > &exprs)
static MatchFillResult isFillInterfaceImpl(Operation *op)
static FailureOr< ContractionDimensions > inferContractionDimsImpl(ArrayRef< AffineMap > indexingMaps, ArrayRef< utils::IteratorType > iterators)
Find 2 parallel (m and n) and 1 reduction (k) dimension candidates that form a matmul subcomputation ...
static bool isContractionBody(Block &block)
Returns true if the block is a body of a contraction with the kinds of operations given pairwise by t...
static std::optional< Value > isaInlinedFillOp(GenericOp op)
Detects if a linalg.generic operation represents a fill with an inlined constant.
static llvm::SmallDenseSet< int64_t > getPreservedDims(AffineMap map)
static bool isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op, unsigned arity)
MatchFillResult
static llvm::SmallDenseSet< int64_t > findPermutationsIndexingOperand(AffineMap indexingMap, ArrayRef< utils::IteratorType > iterators, utils::IteratorType iter)
Given an indexingMap and its corresponding iterators, returns the positions of the iterators of type ...
static FailureOr< SmallVector< utils::IteratorType > > inferIteratorsFromOutMap(AffineMap map)
Infer the iterator types from the init affine map.
static std::pair< int64_t, int64_t > getResultsPositionInLoopsToShapeMap(LinalgOp &op)
static ArrayRef< int64_t > getShape(Type type)
Returns the shape of the given type.
Definition: Traits.cpp:117
Affine binary operation expression.
Definition: AffineExpr.h:214
AffineExpr getLHS() const
Definition: AffineExpr.cpp:338
AffineExpr getRHS() const
Definition: AffineExpr.cpp:341
An integer constant appearing in affine expression.
Definition: AffineExpr.h:239
A dimensional identifier appearing in an affine expression.
Definition: AffineExpr.h:223
unsigned getPosition() const
Definition: AffineExpr.cpp:346
See documentation for AffineExprVisitorBase.
Base type for affine expression.
Definition: AffineExpr.h:68
AffineExprKind getKind() const
Return the classification for this type.
Definition: AffineExpr.cpp:33
MLIRContext * getContext() const
Definition: AffineExpr.cpp:31
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
AffineMap getSliceMap(unsigned start, unsigned length) const
Returns the map consisting of length expressions starting from start.
Definition: AffineMap.cpp:655
bool isProjectedPermutation(bool allowZeroInResults=false) const
Returns true if the AffineMap represents a subset (i.e.
Definition: AffineMap.cpp:611
unsigned getNumDims() const
Definition: AffineMap.cpp:390
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:403
unsigned getNumResults() const
Definition: AffineMap.cpp:398
AffineExpr getResult(unsigned idx) const
Definition: AffineMap.cpp:407
AffineMap compose(AffineMap map) const
Returns the AffineMap resulting from composing this with map.
Definition: AffineMap.cpp:552
A symbolic identifier appearing in an affine expression.
Definition: AffineExpr.h:231
Block represents an ordered list of Operations.
Definition: Block.h:33
bool empty()
Definition: Block.h:148
BlockArgument getArgument(unsigned i)
Definition: Block.h:129
unsigned getNumArguments()
Definition: Block.h:128
Operation & back()
Definition: Block.h:152
Operation * getTerminator()
Get the terminator operation of this block.
Definition: Block.cpp:244
OpListType & getOperations()
Definition: Block.h:137
Operation & front()
Definition: Block.h:153
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:107
An attribute that represents a reference to a dense integer vector or tensor object.
IRValueT get() const
Return the current value being used by this operand.
Definition: UseDefLists.h:160
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
Definition: PatternMatch.h:774
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
This class helps build Operations.
Definition: Builders.h:207
This class represents a single result from folding an operation.
Definition: OpDefinition.h:272
This class represents an operand of an operation.
Definition: Value.h:257
unsigned getOperandNumber()
Return which operand this is in the OpOperand list of the Operation.
Definition: Value.cpp:226
This class provides the API for ops that are known to be terminators.
Definition: OpDefinition.h:773
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
Value getOperand(unsigned idx)
Definition: Operation.h:350
bool mightHaveTrait()
Returns true if the operation might have the provided trait.
Definition: Operation.h:757
unsigned getNumOperands()
Definition: Operation.h:346
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:267
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "'dim' op " which is convenient for verifiers.
Definition: Operation.cpp:672
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:404
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:105
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:18
SmallVector< OpFoldResult > makeComposedFoldedMultiResultAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands, bool composeAffineMin=false)
Variant of makeComposedFoldedAffineApply suitable for multi-result maps.
Definition: AffineOps.cpp:1374
MatchConvolutionResult isConvolutionInterfaceImpl(Operation *op, ConvolutionDimensions *dimensions=nullptr, bool allowEmptyConvolvedDims=false)
Checks whether op conforms to ConvolutionOpInterface and populates dimensions with indexes of the dif...
bool isContractionBody(Block &block, function_ref< bool(Operation *, Operation *)> isaPair, llvm::raw_ostream &errs=mlir::thread_safe_nulls())
Returns true if the block contains a contraction of the following form:
StringRef getMatchConvolutionMessage(MatchConvolutionResult res)
Returns the error message corresponding to the convolution checking return code.
bool canOpOperandsBeDroppedImpl(linalg::LinalgOp linalgOp, ArrayRef< OpOperand * > droppedOperands)
Implementation of the method that check if given operands can be dropped, i.e.
MatchContractionResult isContractionInterfaceImpl(Operation *op, ContractionDimensions *dimensions=nullptr)
Checks whether op conforms to ContractionOpInterface and populates dimensions with indexes of the dif...
LogicalResult verifyContractionInterface(Operation *op)
Verify that op conforms to ContractionOpInterface.
LogicalResult verifyFillInterface(Operation *op)
Verify that op conforms to the FillOpInterface.
StringRef getMatchContractionMessage(MatchContractionResult res)
Returns the error message corresponding to the contraction checking return code.
LogicalResult verifyStructuredOpInterface(Operation *op)
Verify that op conforms to the invariants of StructuredOpInterface.
LogicalResult verifyConvolutionInterface(Operation *op)
Verify that op conforms to the ConvolutionOpInterface.
std::optional< SmallVector< int64_t > > isaTransposeOpInterface(GenericOp genericOp)
Checks whether genericOp is semantically equivalent to a linalg.transpose.
bool isaElemwiseSingleUnaryOpInterface(GenericOp genericOp)
Checks whether a given genericOp is semantically equivalent to a single linalgelementwise unary op.
bool isaCopyOpInterface(LinalgOp linalgOp)
Checks whether linalgOp is semantically equivalent to a linalg.copyOp.
FailureOr< ConvolutionDimensions > inferConvolutionDims(LinalgOp linalgOp)
Find at least 1 parallel (output_image) and reduction (filter_loop) dimension candidates that form a ...
OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value val, int64_t dim)
Create one memref::DimOp or tensor::DimOp depending on the type of val.
Definition: LinalgOps.cpp:104
bool isaConvolutionOpInterface(LinalgOp linalgOp, bool allowEmptyConvolvedDims=false)
Checks whether linalgOp conforms to ConvolutionOpInterface.
std::optional< SmallVector< int64_t > > isaBroadcastOpInterface(GenericOp genericOp)
Checks whether genericOp is semantically equivalent to a linalg.broadcast.
FailureOr< ContractionDimensions > inferContractionDims(LinalgOp linalgOp)
Find at least 2 parallel (m and n) and 1 reduction (k) dimension candidates that form a matmul subcom...
Value createOrFoldDimOp(OpBuilder &b, Location loc, Value val, int64_t dim)
Create one memref::DimOp or tensor::DimOp depending on the type of val.
Definition: LinalgOps.cpp:95
bool isaContractionOpInterface(LinalgOp linalgOp)
Checks whether linalgOp conforms to ContractionOpInterface.
std::optional< Value > isaFillOpInterface(GenericOp genericOp)
Checks whether genericOp is semantically equivalent to a linalg.fill.
bool isaElemwiseSingleBinaryOpInterface(GenericOp genericOp)
Checks whether genericOp is semantically equivalent to a single linalg elementwise binary op e....
detail::InFlightRemark failed(Location loc, RemarkOpts opts)
Report an optimization remark that failed.
Definition: Remarks.h:491
detail::LazyTextBuild add(const char *fmt, Ts &&...ts)
Create a Remark with llvm::formatv formatting.
Definition: Remarks.h:463
Include the generated interface declarations.
AffineMap concatAffineMaps(ArrayRef< AffineMap > maps, MLIRContext *context)
Concatenates a list of maps into a single AffineMap, stepping over potentially empty maps.
Definition: AffineMap.cpp:829
LogicalResult reifyResultShapes(OpBuilder &b, Operation *op, ReifiedRankedShapedTypeDims &reifiedReturnShapes)
Reify the shape of the result of an operation (typically in terms of the shape of its operands).
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
Definition: AffineMap.cpp:784
@ Mul
RHS of mul is always a constant or a symbolic expression.
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
Definition: Utils.cpp:111
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:643
Visitor to check if any of the given set of positions from AffineDimExprs are used within an AffineEx...
HasAffineDimExprVisitor(llvm::SmallBitVector positions)
bool visitDimExpr(AffineDimExpr dimExpr)
bool visitAffineBinaryOpExpr(AffineBinaryOpExpr binaryOpExpr)
bool visitSymbolExpr(AffineSymbolExpr symbolExpr)
bool visitConstantExpr(AffineConstantExpr constExpr)
Represents a range (offset, size, and stride) where each element of the triple may be dynamic or stat...
Positions of a Linalg op loops that correspond to different kinds of a contraction dimension.
Positions of a Linalg op loops that correspond to different kinds of a convolution dimension.