MLIR  20.0.0git
DropUnitDims.cpp
Go to the documentation of this file.
1 //===- DropUnitDims.cpp - Pass to drop use of unit-extent for broadcasting ===//
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 patterns/pass to remove usage of unit-extent dimensions
10 // to specify broadcasting in favor of more canonical representation of the
11 // computation
12 //
13 //===----------------------------------------------------------------------===//
14 
16 
27 #include "mlir/IR/AffineExpr.h"
28 #include "mlir/IR/AffineMap.h"
29 #include "mlir/IR/BuiltinTypes.h"
32 #include "llvm/ADT/SetVector.h"
33 #include "llvm/Support/CommandLine.h"
34 #include "llvm/Support/Debug.h"
35 
36 namespace mlir {
37 #define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS
38 #include "mlir/Dialect/Linalg/Passes.h.inc"
39 } // namespace mlir
40 
41 #define DEBUG_TYPE "linalg-drop-unit-dims"
42 
43 using namespace mlir;
44 using namespace mlir::linalg;
45 
46 namespace {
47 /// Pattern to move init operands to ins when all the loops are parallel and
48 /// blockArgument corresponding to init is used in the region. This is a fix-up
49 /// when unit reduction dimensions are all folded away. In this context, it
50 /// becomes a elementwise generic op. E.g., it converts
51 ///
52 /// %0 = tensor.empty() : tensor<1x1xf32>
53 /// %1 = linalg.fill
54 /// ins(%cst : f32)
55 /// outs(%0 : tensor<1x1xf32>) -> tensor<1x1xf32>
56 /// %2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (0, d0, 0, 0)>,
57 /// affine_map<(d0) -> (0, d0)>],
58 /// iterator_types = ["parallel"]}
59 /// ins(%arg0 : tensor<1x?x1x1xf32>)
60 /// outs(%1 : tensor<1x1xf32>) {
61 /// ^bb0(%in: f32, %out: f32):
62 /// %3 = arith.addf %in, %out : f32
63 /// linalg.yield %3 : f32
64 /// } -> tensor<1x1xf32>
65 ///
66 /// into
67 ///
68 /// %0 = tensor.empty() : tensor<1x1xf32>
69 /// %1 = linalg.fill
70 /// ins(%cst : f32)
71 /// outs(%0 : tensor<1x1xf32>) -> tensor<1x1xf32>
72 /// %2 = tensor.empty() : tensor<1x1xf32>
73 /// %3 = linalg.generic {indexing_maps = [affine_map<(d0) -> (0, d0, 0, 0)>,
74 /// affine_map<(d0) -> (0, d0)>,
75 /// affine_map<(d0) -> (0, d0)>],
76 /// iterator_types = ["parallel"]}
77 /// ins(%arg0, %1 : tensor<1x?x1x1xf32>, tensor<1x1xf32>)
78 /// outs(%2 : tensor<1x1xf32>) {
79 /// ^bb0(%in: f32, %in_0: f32, %out: f32):
80 /// %4 = arith.addf %in, %in_0 : f32
81 /// linalg.yield %4 : f32
82 /// } -> tensor<1x1xf32>
83 struct MoveInitOperandsToInput : public OpRewritePattern<GenericOp> {
85  LogicalResult matchAndRewrite(GenericOp genericOp,
86  PatternRewriter &rewriter) const override {
87  if (!genericOp.hasPureTensorSemantics())
88  return failure();
89  if (genericOp.getNumParallelLoops() != genericOp.getNumLoops())
90  return failure();
91 
92  auto outputOperands = genericOp.getDpsInitsMutable();
93  SetVector<OpOperand *> candidates;
94  for (OpOperand &op : outputOperands) {
95  if (genericOp.getMatchingBlockArgument(&op).use_empty())
96  continue;
97  candidates.insert(&op);
98  }
99 
100  if (candidates.empty())
101  return failure();
102 
103  // Compute the modified indexing maps.
104  int64_t origNumInput = genericOp.getNumDpsInputs();
105  SmallVector<Value> newInputOperands = genericOp.getDpsInputs();
106  SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
107  SmallVector<AffineMap> newIndexingMaps;
108  newIndexingMaps.append(indexingMaps.begin(),
109  std::next(indexingMaps.begin(), origNumInput));
110  for (OpOperand *op : candidates) {
111  newInputOperands.push_back(op->get());
112  newIndexingMaps.push_back(genericOp.getMatchingIndexingMap(op));
113  }
114  newIndexingMaps.append(std::next(indexingMaps.begin(), origNumInput),
115  indexingMaps.end());
116 
117  Location loc = genericOp.getLoc();
118  SmallVector<Value> newOutputOperands =
119  llvm::to_vector(genericOp.getDpsInits());
120  for (OpOperand *op : candidates) {
121  OpBuilder::InsertionGuard guard(rewriter);
122  rewriter.setInsertionPointAfterValue(op->get());
123  auto elemType = cast<ShapedType>(op->get().getType()).getElementType();
124  auto empty = rewriter.create<tensor::EmptyOp>(
125  loc, tensor::getMixedSizes(rewriter, loc, op->get()), elemType);
126 
127  unsigned start = genericOp.getDpsInits().getBeginOperandIndex();
128  newOutputOperands[op->getOperandNumber() - start] = empty.getResult();
129  }
130 
131  auto newOp = rewriter.create<GenericOp>(
132  loc, genericOp.getResultTypes(), newInputOperands, newOutputOperands,
133  newIndexingMaps, genericOp.getIteratorTypesArray(),
134  /*bodyBuild=*/nullptr, linalg::getPrunedAttributeList(genericOp));
135 
136  OpBuilder::InsertionGuard guard(rewriter);
137  Region &region = newOp.getRegion();
138  Block *block = rewriter.createBlock(&region);
139  IRMapping mapper;
140  for (auto bbarg : genericOp.getRegionInputArgs())
141  mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));
142 
143  for (OpOperand *op : candidates) {
144  BlockArgument bbarg = genericOp.getMatchingBlockArgument(op);
145  mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));
146  }
147 
148  for (OpOperand &op : outputOperands) {
149  BlockArgument bbarg = genericOp.getMatchingBlockArgument(&op);
150  if (candidates.count(&op))
151  block->addArgument(bbarg.getType(), loc);
152  else
153  mapper.map(bbarg, block->addArgument(bbarg.getType(), loc));
154  }
155 
156  for (auto &op : genericOp.getBody()->getOperations()) {
157  rewriter.clone(op, mapper);
158  }
159  rewriter.replaceOp(genericOp, newOp.getResults());
160 
161  return success();
162  }
163 };
164 } // namespace
165 
166 //===---------------------------------------------------------------------===//
167 // Drop loops that are unit-extents within Linalg operations.
168 //===---------------------------------------------------------------------===//
169 
170 /// Implements a pass that canonicalizes the uses of unit-extent dimensions for
171 /// broadcasting. For example,
172 ///
173 /// ```mlir
174 /// #accesses = [
175 /// affine_map<(d0, d1) -> (0, d1)>,
176 /// affine_map<(d0, d1) -> (d0, 0)>,
177 /// affine_map<(d0, d1) -> (d0, d1)>
178 /// ]
179 ///
180 /// #trait = {
181 /// indexing_maps = #accesses,
182 /// iterator_types = ["parallel", "parallel"],
183 /// library_call = "some_external_fn"
184 /// }
185 ///
186 /// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
187 /// tensor<5x5xf32>
188 /// {
189 /// %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1) -> (d0, d1)>] :
190 /// tensor<5xf32> into tensor<1x5xf32>
191 /// %1 = linalg.tensor_reshape %arg1 [affine_map<(d0, d1) -> (d0, d1)>] :
192 /// tensor<5xf32> into tensor<5x1xf32>
193 /// %2 = linalg.generic #trait %0, %1 {
194 /// ^bb0(%arg2: f32, %arg3: f32):
195 /// %3 = arith.addf %arg2, %arg3 : f32
196 /// linalg.yield %3 : f32
197 /// } : tensor<1x5xf32>, tensor<5x1xf32> -> tensor<5x5xf32>
198 /// return %2 : tensor<5x5xf32>
199 /// }
200 ///
201 /// would canonicalize to
202 ///
203 /// ```mlir
204 /// #accesses = [
205 /// affine_map<(d0, d1) -> (d1)>,
206 /// affine_map<(d0, d1) -> (d0)>,
207 /// affine_map<(d0, d1) -> (d0, d1)>
208 /// ]
209 ///
210 /// #trait = {
211 /// indexing_maps = #accesses,
212 /// iterator_types = ["parallel", "parallel"],
213 /// library_call = "some_external_fn"
214 /// }
215 ///
216 /// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
217 /// tensor<5x5xf32>
218 /// {
219 /// %0 = linalg.generic #trait %arg0, %arg1 {
220 /// ^bb0(%arg2: f32, %arg3: f32):
221 /// %3 = arith.addf %arg2, %arg3 : f32
222 /// linalg.yield %3 : f32
223 /// } : tensor<5xf32>, tensor<5xf32> -> tensor<5x5xf32>
224 /// return %0 : tensor<5x5xf32>
225 /// }
226 
227 /// Update the index accesses of linalg operations having index semantics.
228 static void
229 replaceUnitDimIndexOps(GenericOp genericOp,
230  const llvm::SmallDenseSet<unsigned> &unitDims,
231  RewriterBase &rewriter) {
232  for (IndexOp indexOp :
233  llvm::make_early_inc_range(genericOp.getBody()->getOps<IndexOp>())) {
234  OpBuilder::InsertionGuard guard(rewriter);
235  rewriter.setInsertionPoint(indexOp);
236  if (unitDims.count(indexOp.getDim()) != 0) {
237  rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(indexOp, 0);
238  } else {
239  // Update the dimension of the index operation if needed.
240  unsigned droppedDims = llvm::count_if(
241  unitDims, [&](unsigned dim) { return dim < indexOp.getDim(); });
242  if (droppedDims != 0)
243  rewriter.replaceOpWithNewOp<IndexOp>(indexOp,
244  indexOp.getDim() - droppedDims);
245  }
246  }
247 }
248 
249 /// Expand the given `value` so that the type matches the type of `origDest`.
250 /// The `reassociation` is used when `rankReductionStrategy` is set to
251 /// `RankReductionStrategy::ReassociativeReshape`.
252 static Value
253 expandValue(RewriterBase &rewriter, Location loc, Value result, Value origDest,
254  ArrayRef<ReassociationIndices> reassociation,
255  ControlDropUnitDims::RankReductionStrategy rankReductionStrategy) {
256  // There are no results for memref outputs.
257  auto origResultType = cast<RankedTensorType>(origDest.getType());
258  if (rankReductionStrategy ==
260  unsigned rank = origResultType.getRank();
261  SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
263  tensor::getMixedSizes(rewriter, loc, origDest);
264  SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
265  return rewriter.createOrFold<tensor::InsertSliceOp>(
266  loc, result, origDest, offsets, sizes, strides);
267  }
268 
269  assert(rankReductionStrategy ==
271  "unknown rank reduction strategy");
272  return rewriter
273  .create<tensor::ExpandShapeOp>(loc, origResultType, result, reassociation)
274  .getResult();
275 }
276 
277 /// Collapse the given `value` so that the type matches the type of
278 /// `origOutput`. The `reassociation` is used when `rankReductionStrategy` is
279 /// set to `RankReductionStrategy::ReassociativeReshape`.
281  RewriterBase &rewriter, Location loc, Value operand,
282  ArrayRef<int64_t> targetShape, ArrayRef<ReassociationIndices> reassociation,
283  ControlDropUnitDims::RankReductionStrategy rankReductionStrategy) {
284  if (auto memrefType = dyn_cast<MemRefType>(operand.getType())) {
285  if (rankReductionStrategy ==
287  FailureOr<Value> rankReducingExtract =
288  memref::SubViewOp::rankReduceIfNeeded(rewriter, loc, operand,
289  targetShape);
290  assert(succeeded(rankReducingExtract) && "not a unit-extent collapse");
291  return *rankReducingExtract;
292  }
293 
294  assert(
295  rankReductionStrategy ==
297  "unknown rank reduction strategy");
298  MemRefLayoutAttrInterface layout;
299  auto targetType = MemRefType::get(targetShape, memrefType.getElementType(),
300  layout, memrefType.getMemorySpace());
301  return rewriter.create<memref::CollapseShapeOp>(loc, targetType, operand,
302  reassociation);
303  }
304  if (auto tensorType = dyn_cast<RankedTensorType>(operand.getType())) {
305  if (rankReductionStrategy ==
307  FailureOr<Value> rankReducingExtract =
308  tensor::ExtractSliceOp::rankReduceIfNeeded(rewriter, loc, operand,
309  targetShape);
310  assert(succeeded(rankReducingExtract) && "not a unit-extent collapse");
311  return *rankReducingExtract;
312  }
313 
314  assert(
315  rankReductionStrategy ==
317  "unknown rank reduction strategy");
318  auto targetType =
319  RankedTensorType::get(targetShape, tensorType.getElementType());
320  return rewriter.create<tensor::CollapseShapeOp>(loc, targetType, operand,
321  reassociation);
322  }
323  llvm_unreachable("unsupported operand type");
324 }
325 
326 /// Compute the modified metadata for an operands of operation
327 /// whose unit dims are being dropped. Return the new indexing map
328 /// to use, the shape of the operand in the replacement op
329 /// and the `reassocation` to use to go from original operand shape
330 /// to modified operand shape.
335 };
337  MLIRContext *context, GenericOp genericOp, OpOperand *opOperand,
338  llvm::SmallDenseMap<unsigned, unsigned> &oldDimsToNewDimsMap,
339  ArrayRef<AffineExpr> dimReplacements) {
341  ReassociationIndices reassociationGroup;
342  SmallVector<AffineExpr> newIndexExprs;
343  AffineMap indexingMap = genericOp.getMatchingIndexingMap(opOperand);
344  ArrayRef<int64_t> operandShape = genericOp.getShape(opOperand);
345  ArrayRef<AffineExpr> exprs = indexingMap.getResults();
346 
347  auto isUnitDim = [&](unsigned dim) {
348  if (auto dimExpr = dyn_cast<AffineDimExpr>(exprs[dim])) {
349  unsigned oldPosition = dimExpr.getPosition();
350  return !oldDimsToNewDimsMap.count(oldPosition) &&
351  (operandShape[dim] == 1);
352  }
353  // Handle the other case where the shape is 1, and is accessed using a
354  // constant 0.
355  if (operandShape[dim] == 1) {
356  auto constAffineExpr = dyn_cast<AffineConstantExpr>(exprs[dim]);
357  return constAffineExpr && constAffineExpr.getValue() == 0;
358  }
359  return false;
360  };
361 
362  unsigned dim = 0;
363  while (dim < operandShape.size() && isUnitDim(dim))
364  reassociationGroup.push_back(dim++);
365  while (dim < operandShape.size()) {
366  assert(!isUnitDim(dim) && "expected non unit-extent");
367  reassociationGroup.push_back(dim);
368  AffineExpr newExpr = exprs[dim].replaceDims(dimReplacements);
369  newIndexExprs.push_back(newExpr);
370  info.targetShape.push_back(operandShape[dim]);
371  ++dim;
372  // Fold all following dimensions that are unit-extent.
373  while (dim < operandShape.size() && isUnitDim(dim)) {
374  reassociationGroup.push_back(dim++);
375  }
376  info.reassociation.push_back(reassociationGroup);
377  reassociationGroup.clear();
378  }
379  info.indexMap =
380  AffineMap::get(oldDimsToNewDimsMap.size(), indexingMap.getNumSymbols(),
381  newIndexExprs, context);
382  return info;
383 }
384 
385 FailureOr<DropUnitDimsResult>
386 linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
387  const ControlDropUnitDims &options) {
388  SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
389  if (indexingMaps.empty())
390  return failure();
391 
392  // 1. Check if any of the iteration dimensions are unit-trip count. They will
393  // end up being unit-trip count if they are used to index into a unit-dim
394  // tensor/memref.
395  AffineMap invertedMap =
396  inversePermutation(concatAffineMaps(indexingMaps, rewriter.getContext()));
397  if (!invertedMap) {
398  return rewriter.notifyMatchFailure(genericOp,
399  "invalid indexing maps for operation");
400  }
401  SmallVector<int64_t> dims = genericOp.getStaticShape();
402 
403  // 1a. Get the allowed list of dimensions to drop from the `options`.
404  SmallVector<unsigned> allowedUnitDims = options.controlFn(genericOp);
405  if (allowedUnitDims.empty()) {
406  return rewriter.notifyMatchFailure(
407  genericOp, "control function returns no allowed unit dims to prune");
408  }
409  llvm::SmallDenseSet<unsigned> unitDimsFilter(allowedUnitDims.begin(),
410  allowedUnitDims.end());
411  llvm::SmallDenseSet<unsigned> unitDims;
412  for (const auto &expr : enumerate(invertedMap.getResults())) {
413  if (AffineDimExpr dimExpr = dyn_cast<AffineDimExpr>(expr.value())) {
414  if (dims[dimExpr.getPosition()] == 1 &&
415  unitDimsFilter.count(expr.index()))
416  unitDims.insert(expr.index());
417  }
418  }
419 
420  // 2. Compute the iterator types of the modified op by dropping the one-trip
421  // count loops.
422  SmallVector<utils::IteratorType> newIteratorTypes;
423  llvm::SmallDenseMap<unsigned, unsigned> oldDimToNewDimMap;
424  SmallVector<AffineExpr> dimReplacements;
425  unsigned newDims = 0;
426  for (auto [index, attr] :
427  llvm::enumerate(genericOp.getIteratorTypesArray())) {
428  if (unitDims.count(index)) {
429  dimReplacements.push_back(
430  getAffineConstantExpr(0, rewriter.getContext()));
431  } else {
432  newIteratorTypes.push_back(attr);
433  oldDimToNewDimMap[index] = newDims;
434  dimReplacements.push_back(
435  getAffineDimExpr(newDims, rewriter.getContext()));
436  newDims++;
437  }
438  }
439 
440  // 3. For each of the operands, find the
441  // - modified affine map to use.
442  // - shape of the operands after the unit-dims are dropped.
443  // - the reassociation indices used to convert from the original
444  // operand type to modified operand (needed only when using reshapes
445  // for rank reduction strategy)
446  // Note that the indexing maps might need changing even if there are no
447  // unit dimensions that are dropped to handle cases where `0` is used to
448  // access a unit-extent tensor. Consider moving this out of this specific
449  // transformation as a stand-alone transformation. Kept here right now due
450  // to legacy.
451  SmallVector<AffineMap> newIndexingMaps;
453  SmallVector<SmallVector<int64_t>> targetShapes;
454  SmallVector<bool> collapsed;
455  auto hasCollapsibleType = [](OpOperand &operand) {
456  Type operandType = operand.get().getType();
457  if (auto memrefOperandType = dyn_cast_or_null<MemRefType>(operandType)) {
458  return memrefOperandType.getLayout().isIdentity();
459  }
460  if (auto tensorOperandType = dyn_cast<RankedTensorType>(operandType)) {
461  return tensorOperandType.getEncoding() == nullptr;
462  }
463  return false;
464  };
465  for (OpOperand &opOperand : genericOp->getOpOperands()) {
466  auto indexingMap = genericOp.getMatchingIndexingMap(&opOperand);
467  ArrayRef<int64_t> shape = genericOp.getShape(&opOperand);
468  if (!hasCollapsibleType(opOperand)) {
469  AffineMap newIndexingMap = indexingMap.replaceDimsAndSymbols(
470  dimReplacements, ArrayRef<AffineExpr>{}, oldDimToNewDimMap.size(), 0);
471  newIndexingMaps.push_back(newIndexingMap);
472  targetShapes.push_back(llvm::to_vector(shape));
473  collapsed.push_back(false);
474  reassociations.push_back({});
475  continue;
476  }
477  auto replacementInfo = dropUnitExtentFromOperandMetadata(
478  rewriter.getContext(), genericOp, &opOperand, oldDimToNewDimMap,
479  dimReplacements);
480  reassociations.push_back(replacementInfo.reassociation);
481  newIndexingMaps.push_back(replacementInfo.indexMap);
482  targetShapes.push_back(replacementInfo.targetShape);
483  collapsed.push_back(!(replacementInfo.indexMap.getNumResults() ==
484  indexingMap.getNumResults()));
485  }
486 
487  // Abort if the indexing maps of the result operation are not invertible
488  // (i.e. not legal) or if no dimension was reduced.
489  if (newIndexingMaps == indexingMaps ||
491  concatAffineMaps(newIndexingMaps, rewriter.getContext())))
492  return failure();
493 
494  Location loc = genericOp.getLoc();
495  // 4. For each of the operands, collapse the operand to convert
496  // from original shape to shape in the modified operation if needed,
497  // either through use of reshapes or rank-reducing slices as
498  // specified in `options`.
499  SmallVector<Value> newOperands;
500  for (OpOperand &opOperand : genericOp->getOpOperands()) {
501  int64_t idx = opOperand.getOperandNumber();
502  if (!collapsed[idx]) {
503  newOperands.push_back(opOperand.get());
504  continue;
505  }
506  newOperands.push_back(collapseValue(rewriter, loc, opOperand.get(),
507  targetShapes[idx], reassociations[idx],
508  options.rankReductionStrategy));
509  }
510 
511  // 5. Create the `linalg.generic` operation with the new operands,
512  // indexing maps, iterator types and result types.
513  ArrayRef<Value> newInputs =
514  ArrayRef<Value>(newOperands).take_front(genericOp.getNumDpsInputs());
515  ArrayRef<Value> newOutputs =
516  ArrayRef<Value>(newOperands).take_back(genericOp.getNumDpsInits());
517  SmallVector<Type> resultTypes;
518  resultTypes.reserve(genericOp.getNumResults());
519  for (unsigned i : llvm::seq<unsigned>(0, genericOp.getNumResults()))
520  resultTypes.push_back(newOutputs[i].getType());
521  GenericOp replacementOp =
522  rewriter.create<GenericOp>(loc, resultTypes, newInputs, newOutputs,
523  newIndexingMaps, newIteratorTypes);
524  rewriter.inlineRegionBefore(genericOp.getRegion(), replacementOp.getRegion(),
525  replacementOp.getRegion().begin());
526  // 5a. Replace `linalg.index` operations that refer to the dropped unit
527  // dimensions.
528  replaceUnitDimIndexOps(replacementOp, unitDims, rewriter);
529 
530  // 6. If any result type changes, insert a reshape/slice to convert from the
531  // original type to the new type.
532  SmallVector<Value> resultReplacements;
533  for (auto [index, result] : llvm::enumerate(replacementOp.getResults())) {
534  unsigned opOperandIndex = index + replacementOp.getNumDpsInputs();
535  Value origDest = genericOp.getDpsInitOperand(index)->get();
536  if (!collapsed[opOperandIndex]) {
537  resultReplacements.push_back(result);
538  continue;
539  }
540  Value expandedValue = expandValue(rewriter, loc, result, origDest,
541  reassociations[opOperandIndex],
542  options.rankReductionStrategy);
543  resultReplacements.push_back(expandedValue);
544  }
545 
546  return DropUnitDimsResult{replacementOp, resultReplacements};
547 }
548 
549 namespace {
550 struct DropUnitDims : public OpRewritePattern<GenericOp> {
551  DropUnitDims(MLIRContext *context, ControlDropUnitDims options = {},
552  PatternBenefit benefit = 1)
553  : OpRewritePattern(context, benefit), options(std::move(options)) {}
554 
555  LogicalResult matchAndRewrite(GenericOp genericOp,
556  PatternRewriter &rewriter) const override {
557  FailureOr<DropUnitDimsResult> result =
558  dropUnitDims(rewriter, genericOp, options);
559  if (failed(result)) {
560  return failure();
561  }
562  rewriter.replaceOp(genericOp, result->replacements);
563  return success();
564  }
565 
566 private:
568 };
569 } // namespace
570 
571 //===---------------------------------------------------------------------===//
572 // Drop dimensions that are unit-extents within tensor operations.
573 //===---------------------------------------------------------------------===//
574 
575 namespace {
576 struct DropPadUnitDims : public OpRewritePattern<tensor::PadOp> {
577  DropPadUnitDims(MLIRContext *context, ControlDropUnitDims options = {},
578  PatternBenefit benefit = 1)
579  : OpRewritePattern(context, benefit), options(std::move(options)) {}
580 
581  LogicalResult matchAndRewrite(tensor::PadOp padOp,
582  PatternRewriter &rewriter) const override {
583  // 1a. Get the allowed list of dimensions to drop from the `options`.
584  SmallVector<unsigned> allowedUnitDims = options.controlFn(padOp);
585  if (allowedUnitDims.empty()) {
586  return rewriter.notifyMatchFailure(
587  padOp, "control function returns no allowed unit dims to prune");
588  }
589 
590  if (padOp.getSourceType().getEncoding()) {
591  return rewriter.notifyMatchFailure(
592  padOp, "cannot collapse dims of tensor with encoding");
593  }
594 
595  // Fail for non-constant padding values. The body of the pad could
596  // depend on the padding indices and/or properties of the padded
597  // tensor so for now we fail.
598  // TODO: Support non-constant padding values.
599  Value paddingVal = padOp.getConstantPaddingValue();
600  if (!paddingVal) {
601  return rewriter.notifyMatchFailure(
602  padOp, "unimplemented: non-constant padding value");
603  }
604 
605  ArrayRef<int64_t> sourceShape = padOp.getSourceType().getShape();
606  int64_t padRank = sourceShape.size();
607 
608  auto isStaticZero = [](OpFoldResult f) {
609  std::optional<int64_t> maybeInt = getConstantIntValue(f);
610  return maybeInt && *maybeInt == 0;
611  };
612 
613  llvm::SmallDenseSet<unsigned> unitDimsFilter(allowedUnitDims.begin(),
614  allowedUnitDims.end());
615  llvm::SmallDenseSet<unsigned> unitDims;
616  SmallVector<int64_t> newShape;
617  SmallVector<OpFoldResult> newLowPad;
618  SmallVector<OpFoldResult> newHighPad;
619  for (const auto [dim, size, low, high] :
620  zip_equal(llvm::seq(static_cast<int64_t>(0), padRank), sourceShape,
621  padOp.getMixedLowPad(), padOp.getMixedHighPad())) {
622  if (unitDimsFilter.contains(dim) && size == 1 && isStaticZero(low) &&
623  isStaticZero(high)) {
624  unitDims.insert(dim);
625  } else {
626  newShape.push_back(size);
627  newLowPad.push_back(low);
628  newHighPad.push_back(high);
629  }
630  }
631 
632  if (unitDims.empty()) {
633  return rewriter.notifyMatchFailure(padOp, "no unit dims to collapse");
634  }
635 
636  ReassociationIndices reassociationGroup;
637  SmallVector<ReassociationIndices> reassociationMap;
638  int64_t dim = 0;
639  while (dim < padRank && unitDims.contains(dim))
640  reassociationGroup.push_back(dim++);
641  while (dim < padRank) {
642  assert(!unitDims.contains(dim) && "expected non unit-extent");
643  reassociationGroup.push_back(dim);
644  dim++;
645  // Fold all following dimensions that are unit-extent.
646  while (dim < padRank && unitDims.contains(dim))
647  reassociationGroup.push_back(dim++);
648  reassociationMap.push_back(reassociationGroup);
649  reassociationGroup.clear();
650  }
651 
652  Value collapsedSource =
653  collapseValue(rewriter, padOp.getLoc(), padOp.getSource(), newShape,
654  reassociationMap, options.rankReductionStrategy);
655 
656  auto newPadOp = rewriter.create<tensor::PadOp>(
657  padOp.getLoc(), /*result=*/Type(), collapsedSource, newLowPad,
658  newHighPad, paddingVal, padOp.getNofold());
659 
660  Value dest = padOp.getResult();
661  if (options.rankReductionStrategy ==
663  SmallVector<OpFoldResult> expandedSizes;
664  int64_t numUnitDims = 0;
665  for (auto dim : llvm::seq(static_cast<int64_t>(0), padRank)) {
666  if (unitDims.contains(dim)) {
667  expandedSizes.push_back(rewriter.getIndexAttr(1));
668  numUnitDims++;
669  continue;
670  }
671  expandedSizes.push_back(tensor::getMixedSize(
672  rewriter, padOp.getLoc(), newPadOp, dim - numUnitDims));
673  }
674  dest = rewriter.create<tensor::EmptyOp>(
675  padOp.getLoc(), expandedSizes,
676  padOp.getResultType().getElementType());
677  }
678 
679  Value expandedValue =
680  expandValue(rewriter, padOp.getLoc(), newPadOp.getResult(), dest,
681  reassociationMap, options.rankReductionStrategy);
682  rewriter.replaceOp(padOp, expandedValue);
683  return success();
684  }
685 
686 private:
688 };
689 } // namespace
690 
691 namespace {
692 /// Convert `extract_slice` operations to rank-reduced versions.
693 struct RankReducedExtractSliceOp
694  : public OpRewritePattern<tensor::ExtractSliceOp> {
696 
697  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
698  PatternRewriter &rewriter) const override {
699  RankedTensorType resultType = sliceOp.getType();
700  SmallVector<OpFoldResult> targetShape;
701  for (auto size : resultType.getShape())
702  targetShape.push_back(rewriter.getIndexAttr(size));
703  auto reassociation = getReassociationMapForFoldingUnitDims(targetShape);
704  if (!reassociation ||
705  reassociation->size() == static_cast<size_t>(resultType.getRank()))
706  return failure();
707 
708  SmallVector<OpFoldResult> offsets = sliceOp.getMixedOffsets();
709  SmallVector<OpFoldResult> strides = sliceOp.getMixedStrides();
710  SmallVector<OpFoldResult> sizes = sliceOp.getMixedSizes();
711  auto rankReducedType = cast<RankedTensorType>(
712  tensor::ExtractSliceOp::inferCanonicalRankReducedResultType(
713  reassociation->size(), sliceOp.getSourceType(), offsets, sizes,
714  strides));
715 
716  Location loc = sliceOp.getLoc();
717  Value newSlice = rewriter.create<tensor::ExtractSliceOp>(
718  loc, rankReducedType, sliceOp.getSource(), offsets, sizes, strides);
719  rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
720  sliceOp, resultType, newSlice, *reassociation);
721  return success();
722  }
723 };
724 
725 /// Convert `insert_slice` operations to rank-reduced versions.
726 /// This patterns works with both InsertSliceOp and ParallelInsertSliceOp.
727 template <typename InsertOpTy>
728 struct RankReducedInsertSliceOp : public OpRewritePattern<InsertOpTy> {
730 
731  LogicalResult matchAndRewrite(InsertOpTy insertSliceOp,
732  PatternRewriter &rewriter) const override {
733  RankedTensorType sourceType = insertSliceOp.getSourceType();
734  SmallVector<OpFoldResult> targetShape;
735  for (auto size : sourceType.getShape())
736  targetShape.push_back(rewriter.getIndexAttr(size));
737  auto reassociation = getReassociationMapForFoldingUnitDims(targetShape);
738  if (!reassociation ||
739  reassociation->size() == static_cast<size_t>(sourceType.getRank()))
740  return failure();
741 
742  Location loc = insertSliceOp.getLoc();
743  tensor::CollapseShapeOp reshapedSource;
744  {
745  OpBuilder::InsertionGuard g(rewriter);
746  // The only difference between InsertSliceOp and ParallelInsertSliceOp
747  // is the insertion point is just before the ParallelCombiningOp in the
748  // parallel case.
749  if (std::is_same<InsertOpTy, tensor::ParallelInsertSliceOp>::value)
750  rewriter.setInsertionPoint(insertSliceOp->getParentOp());
751  reshapedSource = rewriter.create<tensor::CollapseShapeOp>(
752  loc, insertSliceOp.getSource(), *reassociation);
753  }
754  rewriter.replaceOpWithNewOp<InsertOpTy>(
755  insertSliceOp, reshapedSource, insertSliceOp.getDest(),
756  insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
757  insertSliceOp.getMixedStrides());
758  return success();
759  }
760 };
761 } // namespace
762 
763 /// Patterns that are used to canonicalize the use of unit-extent dims for
764 /// broadcasting.
765 static void
768  auto *context = patterns.getContext();
769  patterns.add<DropUnitDims>(context, options);
770  patterns.add<DropPadUnitDims>(context, options);
771  // TODO: Patterns unrelated to unit dim folding should be factored out.
772  patterns.add<RankReducedExtractSliceOp,
773  RankReducedInsertSliceOp<tensor::InsertSliceOp>,
774  RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>(
775  context);
776  linalg::FillOp::getCanonicalizationPatterns(patterns, context);
777  tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
778  tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
779  tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
783 }
784 
785 static void
788  auto *context = patterns.getContext();
789  patterns.add<DropUnitDims>(context, options);
790  patterns.add<DropPadUnitDims>(context, options);
791  // TODO: Patterns unrelated to unit dim folding should be factored out.
792  linalg::FillOp::getCanonicalizationPatterns(patterns, context);
793  tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
797 }
798 
801  if (options.rankReductionStrategy ==
804  } else if (options.rankReductionStrategy ==
806  ReassociativeReshape) {
808  }
809 }
810 
813  patterns.add<MoveInitOperandsToInput>(patterns.getContext());
814 }
815 
816 namespace {
817 /// Pass that removes unit-extent dims within generic ops.
818 struct LinalgFoldUnitExtentDimsPass
819  : public impl::LinalgFoldUnitExtentDimsPassBase<
820  LinalgFoldUnitExtentDimsPass> {
821  using impl::LinalgFoldUnitExtentDimsPassBase<
822  LinalgFoldUnitExtentDimsPass>::LinalgFoldUnitExtentDimsPassBase;
823  void runOnOperation() override {
824  Operation *op = getOperation();
825  MLIRContext *context = op->getContext();
826  RewritePatternSet patterns(context);
828  if (useRankReducingSlices) {
829  options.rankReductionStrategy = linalg::ControlDropUnitDims::
831  }
834  (void)applyPatternsGreedily(op, std::move(patterns));
835  }
836 };
837 
838 } // namespace
839 
840 namespace {
841 
842 /// Returns reassociation indices for collapsing/expanding a
843 /// tensor of rank `rank` at position `pos`.
845 getReassociationForReshapeAtDim(int64_t rank, int64_t pos) {
846  SmallVector<ReassociationIndices> reassociation(rank - 1, {0, 1});
847  bool lastDim = pos == rank - 1;
848  if (rank > 2) {
849  for (int64_t i = 0; i < rank - 1; i++) {
850  if (i == pos || (lastDim && i == pos - 1))
851  reassociation[i] = ReassociationIndices{i, i + 1};
852  else if (i < pos)
853  reassociation[i] = ReassociationIndices{i};
854  else
855  reassociation[i] = ReassociationIndices{i + 1};
856  }
857  }
858  return reassociation;
859 }
860 
861 /// Returns a collapsed `val` where the collapsing occurs at dim `pos`.
862 /// If `pos < 0`, then don't collapse.
863 static Value collapseSingletonDimAt(PatternRewriter &rewriter, Value val,
864  int64_t pos) {
865  if (pos < 0)
866  return val;
867  auto valType = cast<ShapedType>(val.getType());
868  SmallVector<int64_t> collapsedShape(valType.getShape());
869  collapsedShape.erase(collapsedShape.begin() + pos);
870  return collapseValue(
871  rewriter, val.getLoc(), val, collapsedShape,
872  getReassociationForReshapeAtDim(valType.getRank(), pos),
874 }
875 
876 /// Base class for all rank reduction patterns for contraction ops
877 /// with unit dimensions. All patterns should convert one named op
878 /// to another named op. Intended to reduce only one iteration space dim
879 /// at a time.
880 /// Reducing multiple dims will happen with recusive application of
881 /// pattern rewrites.
882 template <typename FromOpTy, typename ToOpTy>
883 struct RankReduceContractionOps : OpRewritePattern<FromOpTy> {
885 
886  /// Collapse all collapsable operands.
888  collapseOperands(PatternRewriter &rewriter, ArrayRef<Value> operands,
889  ArrayRef<int64_t> operandCollapseDims) const {
890  assert(operandCollapseDims.size() == 3 && operands.size() == 3 &&
891  "expected 3 operands and dims");
892  return llvm::map_to_vector(
893  llvm::zip(operands, operandCollapseDims), [&](auto pair) {
894  return collapseSingletonDimAt(rewriter, std::get<0>(pair),
895  std::get<1>(pair));
896  });
897  }
898 
899  /// Expand result tensor.
900  Value expandResult(PatternRewriter &rewriter, Value result,
901  RankedTensorType expandedType, int64_t dim) const {
902  return rewriter.create<tensor::ExpandShapeOp>(
903  result.getLoc(), expandedType, result,
904  getReassociationForReshapeAtDim(expandedType.getRank(), dim));
905  }
906 
907  LogicalResult matchAndRewrite(FromOpTy contractionOp,
908  PatternRewriter &rewriter) const override {
909 
910  auto loc = contractionOp.getLoc();
911  auto inputs = contractionOp.getDpsInputs();
912  auto inits = contractionOp.getDpsInits();
913  if (inputs.size() != 2 || inits.size() != 1)
914  return rewriter.notifyMatchFailure(contractionOp,
915  "expected 2 inputs and 1 init");
916  auto lhs = inputs[0];
917  auto rhs = inputs[1];
918  auto init = inits[0];
919  SmallVector<Value> operands{lhs, rhs, init};
920 
921  SmallVector<int64_t> operandUnitDims;
922  if (failed(getOperandUnitDims(contractionOp, operandUnitDims)))
923  return rewriter.notifyMatchFailure(contractionOp,
924  "no reducable dims found");
925 
926  SmallVector<Value> collapsedOperands =
927  collapseOperands(rewriter, operands, operandUnitDims);
928  Value collapsedLhs = collapsedOperands[0];
929  Value collapsedRhs = collapsedOperands[1];
930  Value collapsedInit = collapsedOperands[2];
931  SmallVector<Type, 1> collapsedResultTy;
932  if (isa<RankedTensorType>(collapsedInit.getType()))
933  collapsedResultTy.push_back(collapsedInit.getType());
934  auto collapsedOp = rewriter.create<ToOpTy>(
935  loc, collapsedResultTy, ValueRange{collapsedLhs, collapsedRhs},
936  ValueRange{collapsedInit});
937  for (auto attr : contractionOp->getAttrs()) {
938  if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName)
939  continue;
940  collapsedOp->setAttr(attr.getName(), attr.getValue());
941  }
942 
943  auto results = contractionOp.getResults();
944  assert(results.size() < 2 && "expected at most one result");
945  if (results.empty()) {
946  rewriter.replaceOp(contractionOp, collapsedOp);
947  } else {
948  rewriter.replaceOp(
949  contractionOp,
950  expandResult(rewriter, collapsedOp.getResultTensors()[0],
951  cast<RankedTensorType>(results[0].getType()),
952  operandUnitDims[2]));
953  }
954 
955  return success();
956  }
957 
958  /// Populate `operandUnitDims` with 3 indices indicating the unit dim
959  /// for each operand that should be collapsed in this pattern. If an
960  /// operand shouldn't be collapsed, the index should be negative.
961  virtual LogicalResult
962  getOperandUnitDims(LinalgOp op,
963  SmallVectorImpl<int64_t> &operandUnitDims) const = 0;
964 };
965 
966 /// Patterns for unbatching batched contraction ops
967 template <typename FromOpTy, typename ToOpTy>
968 struct RankReduceToUnBatched : RankReduceContractionOps<FromOpTy, ToOpTy> {
969  using RankReduceContractionOps<FromOpTy, ToOpTy>::RankReduceContractionOps;
970 
971  /// Look for unit batch dims to collapse.
972  LogicalResult
973  getOperandUnitDims(LinalgOp op,
974  SmallVectorImpl<int64_t> &operandUnitDims) const override {
975  FailureOr<ContractionDimensions> maybeContractionDims =
977  if (failed(maybeContractionDims)) {
978  LLVM_DEBUG(llvm::dbgs() << "could not infer contraction dims");
979  return failure();
980  }
981  ContractionDimensions contractionDims = maybeContractionDims.value();
982 
983  if (contractionDims.batch.size() != 1)
984  return failure();
985  auto batchDim = contractionDims.batch[0];
987  op.mapIterationSpaceDimToAllOperandDims(batchDim, bOperands);
988  if (bOperands.size() != 3 || llvm::any_of(bOperands, [](auto pair) {
989  return cast<ShapedType>(std::get<0>(pair).getType())
990  .getShape()[std::get<1>(pair)] != 1;
991  })) {
992  LLVM_DEBUG(llvm::dbgs() << "specified unit dims not found");
993  return failure();
994  }
995 
996  operandUnitDims = SmallVector<int64_t>{std::get<1>(bOperands[0]),
997  std::get<1>(bOperands[1]),
998  std::get<1>(bOperands[2])};
999  return success();
1000  }
1001 };
1002 
1003 /// Patterns for reducing non-batch dimensions
1004 template <typename FromOpTy, typename ToOpTy>
1005 struct RankReduceMatmul : RankReduceContractionOps<FromOpTy, ToOpTy> {
1006  using RankReduceContractionOps<FromOpTy, ToOpTy>::RankReduceContractionOps;
1007 
1008  /// Helper for determining whether the lhs/init or rhs/init are reduced.
1009  static bool constexpr reduceLeft =
1010  (std::is_same_v<FromOpTy, BatchMatmulOp> &&
1011  std::is_same_v<ToOpTy, BatchVecmatOp>) ||
1012  (std::is_same_v<FromOpTy, BatchMatmulTransposeAOp> &&
1013  std::is_same_v<ToOpTy, BatchVecmatOp>) ||
1014  (std::is_same_v<FromOpTy, MatmulOp> &&
1015  std::is_same_v<ToOpTy, VecmatOp>) ||
1016  (std::is_same_v<FromOpTy, MatmulTransposeAOp> &&
1017  std::is_same_v<ToOpTy, VecmatOp>) ||
1018  (std::is_same_v<FromOpTy, MatvecOp> && std::is_same_v<ToOpTy, DotOp>);
1019 
1020  /// Look for non-batch spatial dims to collapse.
1021  LogicalResult
1022  getOperandUnitDims(LinalgOp op,
1023  SmallVectorImpl<int64_t> &operandUnitDims) const override {
1024  FailureOr<ContractionDimensions> maybeContractionDims =
1026  if (failed(maybeContractionDims)) {
1027  LLVM_DEBUG(llvm::dbgs() << "could not infer contraction dims");
1028  return failure();
1029  }
1030  ContractionDimensions contractionDims = maybeContractionDims.value();
1031 
1032  if constexpr (reduceLeft) {
1033  auto m = contractionDims.m[0];
1035  op.mapIterationSpaceDimToAllOperandDims(m, mOperands);
1036  if (mOperands.size() != 2)
1037  return failure();
1038  if (llvm::all_of(mOperands, [](auto pair) {
1039  return cast<ShapedType>(std::get<0>(pair).getType())
1040  .getShape()[std::get<1>(pair)] == 1;
1041  })) {
1042  operandUnitDims = SmallVector<int64_t>{std::get<1>(mOperands[0]), -1,
1043  std::get<1>(mOperands[1])};
1044  return success();
1045  }
1046  } else {
1047  auto n = contractionDims.n[0];
1049  op.mapIterationSpaceDimToAllOperandDims(n, nOperands);
1050  if (nOperands.size() != 2)
1051  return failure();
1052  if (llvm::all_of(nOperands, [](auto pair) {
1053  return cast<ShapedType>(std::get<0>(pair).getType())
1054  .getShape()[std::get<1>(pair)] == 1;
1055  })) {
1056  operandUnitDims = SmallVector<int64_t>{-1, std::get<1>(nOperands[0]),
1057  std::get<1>(nOperands[1])};
1058  return success();
1059  }
1060  }
1061  LLVM_DEBUG(llvm::dbgs() << "specified unit dims not found");
1062  return failure();
1063  }
1064 };
1065 
1066 } // namespace
1067 
1070  MLIRContext *context = patterns.getContext();
1071  // Unbatching patterns for unit batch size
1072  patterns.add<RankReduceToUnBatched<BatchMatmulOp, MatmulOp>>(context);
1073  patterns
1074  .add<RankReduceToUnBatched<BatchMatmulTransposeAOp, MatmulTransposeAOp>>(
1075  context);
1076  patterns
1077  .add<RankReduceToUnBatched<BatchMatmulTransposeBOp, MatmulTransposeBOp>>(
1078  context);
1079  patterns.add<RankReduceToUnBatched<BatchMatvecOp, MatvecOp>>(context);
1080  patterns.add<RankReduceToUnBatched<BatchVecmatOp, VecmatOp>>(context);
1081 
1082  // Non-batch rank 1 reducing patterns
1083  patterns.add<RankReduceMatmul<MatmulOp, VecmatOp>>(context);
1084  patterns.add<RankReduceMatmul<MatmulOp, MatvecOp>>(context);
1085  patterns.add<RankReduceMatmul<MatmulTransposeAOp, VecmatOp>>(context);
1086  patterns.add<RankReduceMatmul<MatmulTransposeBOp, MatvecOp>>(context);
1087  // Batch rank 1 reducing patterns
1088  patterns.add<RankReduceMatmul<BatchMatmulOp, BatchVecmatOp>>(context);
1089  patterns.add<RankReduceMatmul<BatchMatmulOp, BatchMatvecOp>>(context);
1090  patterns.add<RankReduceMatmul<BatchMatmulTransposeAOp, BatchVecmatOp>>(
1091  context);
1092  patterns.add<RankReduceMatmul<BatchMatmulTransposeBOp, BatchMatvecOp>>(
1093  context);
1094 
1095  // Non-batch rank 0 reducing patterns
1096  patterns.add<RankReduceMatmul<MatvecOp, DotOp>>(context);
1097  patterns.add<RankReduceMatmul<VecmatOp, DotOp>>(context);
1098 }
static Value expandValue(RewriterBase &rewriter, Location loc, Value result, Value origDest, ArrayRef< ReassociationIndices > reassociation, ControlDropUnitDims::RankReductionStrategy rankReductionStrategy)
Expand the given value so that the type matches the type of origDest.
static void replaceUnitDimIndexOps(GenericOp genericOp, const llvm::SmallDenseSet< unsigned > &unitDims, RewriterBase &rewriter)
Implements a pass that canonicalizes the uses of unit-extent dimensions for broadcasting.
static UnitExtentReplacementInfo dropUnitExtentFromOperandMetadata(MLIRContext *context, GenericOp genericOp, OpOperand *opOperand, llvm::SmallDenseMap< unsigned, unsigned > &oldDimsToNewDimsMap, ArrayRef< AffineExpr > dimReplacements)
static void populateFoldUnitExtentDimsViaReshapesPatterns(RewritePatternSet &patterns, ControlDropUnitDims &options)
Patterns that are used to canonicalize the use of unit-extent dims for broadcasting.
static void populateFoldUnitExtentDimsViaSlicesPatterns(RewritePatternSet &patterns, ControlDropUnitDims &options)
static Value collapseValue(RewriterBase &rewriter, Location loc, Value operand, ArrayRef< int64_t > targetShape, ArrayRef< ReassociationIndices > reassociation, ControlDropUnitDims::RankReductionStrategy rankReductionStrategy)
Collapse the given value so that the type matches the type of origOutput.
static llvm::ManagedStatic< PassManagerOptions > options
A dimensional identifier appearing in an affine expression.
Definition: AffineExpr.h:236
Base type for affine expression.
Definition: AffineExpr.h:68
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:46
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
unsigned getNumSymbols() const
Definition: AffineMap.cpp:398
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:407
AffineMap replaceDimsAndSymbols(ArrayRef< AffineExpr > dimReplacements, ArrayRef< AffineExpr > symReplacements, unsigned numResultDims, unsigned numResultSyms) const
This method substitutes any uses of dimensions and symbols (e.g.
Definition: AffineMap.cpp:500
This class represents an argument of a Block.
Definition: Value.h:319
Block represents an ordered list of Operations.
Definition: Block.h:33
BlockArgument addArgument(Type type, Location loc)
Add one value to the argument list.
Definition: Block.cpp:155
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:148
MLIRContext * getContext() const
Definition: Builders.h:56
This is a utility class for mapping one set of IR entities to another.
Definition: IRMapping.h:26
void map(Value from, Value to)
Inserts a new mapping for 'from' to 'to'.
Definition: IRMapping.h:30
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:66
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:357
Operation * clone(Operation &op, IRMapping &mapper)
Creates a deep copy of the specified operation, remapping any operands that use values outside of the...
Definition: Builders.cpp:588
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:407
Block * createBlock(Region *parent, Region::iterator insertPt={}, TypeRange argTypes=std::nullopt, ArrayRef< Location > locs=std::nullopt)
Add new block with 'argTypes' arguments and set the insertion point to the end of it.
Definition: Builders.cpp:470
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:529
void setInsertionPointAfterValue(Value val)
Sets the insertion point to the node after the specified value.
Definition: Builders.h:430
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:497
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
This class represents an operand of an operation.
Definition: Value.h:267
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:216
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
Definition: Operation.h:582
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
Definition: PatternMatch.h:34
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:791
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Definition: Region.h:26
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:400
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
Definition: PatternMatch.h:724
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
void inlineRegionBefore(Region &region, Region &parent, Region::iterator before)
Move the blocks that belong to "region" before the given position in another region "parent".
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:542
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:381
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:129
Location getLoc() const
Return the location of this value.
Definition: Value.cpp:26
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:344
void populateMoveInitOperandsToInputPattern(RewritePatternSet &patterns)
A pattern that converts init operands to input operands.
void populateContractionOpRankReducingPatterns(RewritePatternSet &patterns)
Adds patterns that reduce the rank of named contraction ops that have unit dimensions in the operand(...
FailureOr< DropUnitDimsResult > dropUnitDims(RewriterBase &rewriter, GenericOp genericOp, const ControlDropUnitDims &options)
SmallVector< NamedAttribute > getPrunedAttributeList(OpTy op)
Returns an attribute list that excludes pre-defined attributes.
Definition: Utils.h:364
std::optional< SmallVector< ReassociationIndices > > getReassociationMapForFoldingUnitDims(ArrayRef< OpFoldResult > mixedSizes)
Get the reassociation maps to fold the result of a extract_slice (or source of a insert_slice) operat...
Definition: Utils.cpp:852
void populateFoldUnitExtentDimsPatterns(RewritePatternSet &patterns, ControlDropUnitDims &options)
Patterns to fold unit-extent dimensions in operands/results of linalg ops on tensors via reassociativ...
FailureOr< ContractionDimensions > inferContractionDims(LinalgOp linalgOp)
Find at least 2 parallel (m and n) and 1 reduction (k) dimension candidates that form a matmul subcom...
void populateResolveRankedShapedTypeResultDimsPatterns(RewritePatternSet &patterns)
Appends patterns that resolve memref.dim operations with values that are defined by operations that i...
void populateResolveShapedTypeResultDimsPatterns(RewritePatternSet &patterns)
Appends patterns that resolve memref.dim operations with values that are defined by operations that i...
void populateFoldTensorEmptyPatterns(RewritePatternSet &patterns, bool foldSingleUseOnly=false)
Populates patterns with patterns that fold tensor.empty with its consumers.
OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value, int64_t dim)
Return the dimension of the given tensor value.
Definition: TensorOps.cpp:56
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given tensor value.
Definition: TensorOps.cpp:66
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:836
std::optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
Type getType(OpFoldResult ofr)
Returns the int type of the integer in ofr.
Definition: Utils.cpp:305
LogicalResult applyPatternsGreedily(Region &region, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig(), bool *changed=nullptr)
Rewrite ops in the given region, which must be isolated from above, by repeatedly applying the highes...
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
Definition: AffineMap.cpp:791
const FrozenRewritePatternSet & patterns
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:641
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)
These free functions allow clients of the API to not use classes in detail.
Definition: AffineExpr.cpp:617
Compute the modified metadata for an operands of operation whose unit dims are being dropped.
SmallVector< ReassociationIndices > reassociation
SmallVector< int64_t > targetShape
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:358
Positions of a Linalg op loops that correspond to different kinds of a contraction dimension.
SmallVector< unsigned, 2 > batch
SmallVector< unsigned, 2 > m
SmallVector< unsigned, 2 > n
Transformation to drop unit-extent dimensions from linalg.generic operations.
Definition: Transforms.h:473