MLIR  19.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 /// args_in = 2,
182 /// args_out = 1,
183 /// indexing_maps = #accesses,
184 /// iterator_types = ["parallel", "parallel"],
185 /// library_call = "some_external_fn"
186 /// }
187 ///
188 /// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
189 /// tensor<5x5xf32>
190 /// {
191 /// %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1) -> (d0, d1)>] :
192 /// tensor<5xf32> into tensor<1x5xf32>
193 /// %1 = linalg.tensor_reshape %arg1 [affine_map<(d0, d1) -> (d0, d1)>] :
194 /// tensor<5xf32> into tensor<5x1xf32>
195 /// %2 = linalg.generic #trait %0, %1 {
196 /// ^bb0(%arg2: f32, %arg3: f32):
197 /// %3 = arith.addf %arg2, %arg3 : f32
198 /// linalg.yield %3 : f32
199 /// } : tensor<1x5xf32>, tensor<5x1xf32> -> tensor<5x5xf32>
200 /// return %2 : tensor<5x5xf32>
201 /// }
202 ///
203 /// would canonicalize to
204 ///
205 /// ```mlir
206 /// #accesses = [
207 /// affine_map<(d0, d1) -> (d1)>,
208 /// affine_map<(d0, d1) -> (d0)>,
209 /// affine_map<(d0, d1) -> (d0, d1)>
210 /// ]
211 ///
212 /// #trait = {
213 /// args_in = 2,
214 /// args_out = 1,
215 /// indexing_maps = #accesses,
216 /// iterator_types = ["parallel", "parallel"],
217 /// library_call = "some_external_fn"
218 /// }
219 ///
220 /// func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
221 /// tensor<5x5xf32>
222 /// {
223 /// %0 = linalg.generic #trait %arg0, %arg1 {
224 /// ^bb0(%arg2: f32, %arg3: f32):
225 /// %3 = arith.addf %arg2, %arg3 : f32
226 /// linalg.yield %3 : f32
227 /// } : tensor<5xf32>, tensor<5xf32> -> tensor<5x5xf32>
228 /// return %0 : tensor<5x5xf32>
229 /// }
230 
231 /// Update the index accesses of linalg operations having index semantics.
232 static void
233 replaceUnitDimIndexOps(GenericOp genericOp,
234  const llvm::SmallDenseSet<unsigned> &unitDims,
235  RewriterBase &rewriter) {
236  for (IndexOp indexOp :
237  llvm::make_early_inc_range(genericOp.getBody()->getOps<IndexOp>())) {
238  OpBuilder::InsertionGuard guard(rewriter);
239  rewriter.setInsertionPoint(indexOp);
240  if (unitDims.count(indexOp.getDim()) != 0) {
241  rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(indexOp, 0);
242  } else {
243  // Update the dimension of the index operation if needed.
244  unsigned droppedDims = llvm::count_if(
245  unitDims, [&](unsigned dim) { return dim < indexOp.getDim(); });
246  if (droppedDims != 0)
247  rewriter.replaceOpWithNewOp<IndexOp>(indexOp,
248  indexOp.getDim() - droppedDims);
249  }
250  }
251 }
252 
253 /// Expand the given `value` so that the type matches the type of `origDest`.
254 /// The `reassociation` is used when `rankReductionStrategy` is set to
255 /// `RankReductionStrategy::ReassociativeReshape`.
256 static Value
257 expandValue(RewriterBase &rewriter, Location loc, Value result, Value origDest,
258  ArrayRef<ReassociationIndices> reassociation,
259  ControlDropUnitDims::RankReductionStrategy rankReductionStrategy) {
260  // There are no results for memref outputs.
261  auto origResultType = cast<RankedTensorType>(origDest.getType());
262  if (rankReductionStrategy ==
264  unsigned rank = origResultType.getRank();
265  SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
267  tensor::getMixedSizes(rewriter, loc, origDest);
268  SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
269  return rewriter.createOrFold<tensor::InsertSliceOp>(
270  loc, result, origDest, offsets, sizes, strides);
271  }
272 
273  assert(rankReductionStrategy ==
275  "unknown rank reduction strategy");
276  return rewriter
277  .create<tensor::ExpandShapeOp>(loc, origResultType, result, reassociation)
278  .getResult();
279 }
280 
281 /// Collapse the given `value` so that the type matches the type of
282 /// `origOutput`. The `reassociation` is used when `rankReductionStrategy` is
283 /// set to `RankReductionStrategy::ReassociativeReshape`.
285  RewriterBase &rewriter, Location loc, Value operand,
286  ArrayRef<int64_t> targetShape, ArrayRef<ReassociationIndices> reassociation,
287  ControlDropUnitDims::RankReductionStrategy rankReductionStrategy) {
288  if (auto memrefType = dyn_cast<MemRefType>(operand.getType())) {
289  if (rankReductionStrategy ==
291  FailureOr<Value> rankReducingExtract =
292  memref::SubViewOp::rankReduceIfNeeded(rewriter, loc, operand,
293  targetShape);
294  assert(succeeded(rankReducingExtract) && "not a unit-extent collapse");
295  return *rankReducingExtract;
296  }
297 
298  assert(
299  rankReductionStrategy ==
301  "unknown rank reduction strategy");
302  MemRefLayoutAttrInterface layout;
303  auto targetType = MemRefType::get(targetShape, memrefType.getElementType(),
304  layout, memrefType.getMemorySpace());
305  return rewriter.create<memref::CollapseShapeOp>(loc, targetType, operand,
306  reassociation);
307  }
308  if (auto tensorType = dyn_cast<RankedTensorType>(operand.getType())) {
309  if (rankReductionStrategy ==
311  FailureOr<Value> rankReducingExtract =
312  tensor::ExtractSliceOp::rankReduceIfNeeded(rewriter, loc, operand,
313  targetShape);
314  assert(succeeded(rankReducingExtract) && "not a unit-extent collapse");
315  return *rankReducingExtract;
316  }
317 
318  assert(
319  rankReductionStrategy ==
321  "unknown rank reduction strategy");
322  auto targetType =
323  RankedTensorType::get(targetShape, tensorType.getElementType());
324  return rewriter.create<tensor::CollapseShapeOp>(loc, targetType, operand,
325  reassociation);
326  }
327  llvm_unreachable("unsupported operand type");
328 }
329 
330 /// Compute the modified metadata for an operands of operation
331 /// whose unit dims are being dropped. Return the new indexing map
332 /// to use, the shape of the operand in the replacement op
333 /// and the `reassocation` to use to go from original operand shape
334 /// to modified operand shape.
339 };
341  MLIRContext *context, GenericOp genericOp, OpOperand *opOperand,
342  llvm::SmallDenseMap<unsigned, unsigned> &oldDimsToNewDimsMap,
343  ArrayRef<AffineExpr> dimReplacements) {
345  ReassociationIndices reassociationGroup;
346  SmallVector<AffineExpr> newIndexExprs;
347  AffineMap indexingMap = genericOp.getMatchingIndexingMap(opOperand);
348  ArrayRef<int64_t> operandShape = genericOp.getShape(opOperand);
349  ArrayRef<AffineExpr> exprs = indexingMap.getResults();
350 
351  auto isUnitDim = [&](unsigned dim) {
352  if (auto dimExpr = dyn_cast<AffineDimExpr>(exprs[dim])) {
353  unsigned oldPosition = dimExpr.getPosition();
354  return !oldDimsToNewDimsMap.count(oldPosition);
355  }
356  // Handle the other case where the shape is 1, and is accessed using a
357  // constant 0.
358  if (operandShape[dim] == 1) {
359  auto constAffineExpr = dyn_cast<AffineConstantExpr>(exprs[dim]);
360  return constAffineExpr && constAffineExpr.getValue() == 0;
361  }
362  return false;
363  };
364 
365  unsigned dim = 0;
366  while (dim < operandShape.size() && isUnitDim(dim))
367  reassociationGroup.push_back(dim++);
368  while (dim < operandShape.size()) {
369  assert(!isUnitDim(dim) && "expected non unit-extent");
370  reassociationGroup.push_back(dim);
371  AffineExpr newExpr = exprs[dim].replaceDims(dimReplacements);
372  newIndexExprs.push_back(newExpr);
373  info.targetShape.push_back(operandShape[dim]);
374  ++dim;
375  // Fold all following dimensions that are unit-extent.
376  while (dim < operandShape.size() && isUnitDim(dim)) {
377  reassociationGroup.push_back(dim++);
378  }
379  info.reassociation.push_back(reassociationGroup);
380  reassociationGroup.clear();
381  }
382  info.indexMap =
383  AffineMap::get(oldDimsToNewDimsMap.size(), indexingMap.getNumSymbols(),
384  newIndexExprs, context);
385  return info;
386 }
387 
388 LogicalResult linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
389  const ControlDropUnitDims &options) {
390  SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
391  if (indexingMaps.empty())
392  return failure();
393 
394  // 1. Check if any of the iteration dimensions are unit-trip count. They will
395  // end up being unit-trip count if they are used to index into a unit-dim
396  // tensor/memref.
397  AffineMap invertedMap = inversePermutation(concatAffineMaps(indexingMaps));
398  if (!invertedMap) {
399  return rewriter.notifyMatchFailure(genericOp,
400  "invalid indexing maps for operation");
401  }
402  SmallVector<int64_t> dims = genericOp.getStaticShape();
403 
404  // 1a. Get the allowed list of dimensions to drop from the `options`.
405  SmallVector<unsigned> allowedUnitDims = options.controlFn(genericOp);
406  if (allowedUnitDims.empty()) {
407  return rewriter.notifyMatchFailure(
408  genericOp, "control function returns no allowed unit dims to prune");
409  }
410  llvm::SmallDenseSet<unsigned> unitDimsFilter(allowedUnitDims.begin(),
411  allowedUnitDims.end());
412  llvm::SmallDenseSet<unsigned> unitDims;
413  for (const auto &expr : enumerate(invertedMap.getResults())) {
414  if (AffineDimExpr dimExpr = dyn_cast<AffineDimExpr>(expr.value())) {
415  if (dims[dimExpr.getPosition()] == 1 &&
416  unitDimsFilter.count(expr.index()))
417  unitDims.insert(expr.index());
418  }
419  }
420 
421  // 2. Compute the iterator types of the modified op by dropping the one-trip
422  // count loops.
423  SmallVector<utils::IteratorType> newIteratorTypes;
424  llvm::SmallDenseMap<unsigned, unsigned> oldDimToNewDimMap;
425  SmallVector<AffineExpr> dimReplacements;
426  unsigned newDims = 0;
427  for (auto [index, attr] :
428  llvm::enumerate(genericOp.getIteratorTypesArray())) {
429  if (unitDims.count(index)) {
430  dimReplacements.push_back(
431  getAffineConstantExpr(0, rewriter.getContext()));
432  } else {
433  newIteratorTypes.push_back(attr);
434  oldDimToNewDimMap[index] = newDims;
435  dimReplacements.push_back(
436  getAffineDimExpr(newDims, rewriter.getContext()));
437  newDims++;
438  }
439  }
440 
441  // 3. For each of the operands, find the
442  // - modified affine map to use.
443  // - shape of the operands after the unit-dims are dropped.
444  // - the reassociation indices used to convert from the original
445  // operand type to modified operand (needed only when using reshapes
446  // for rank reduction strategy)
447  // Note that the indexing maps might need changing even if there are no
448  // unit dimensions that are dropped to handle cases where `0` is used to
449  // access a unit-extent tensor. Consider moving this out of this specific
450  // transformation as a stand-alone transformation. Kept here right now due
451  // to legacy.
452  SmallVector<AffineMap> newIndexingMaps;
454  SmallVector<SmallVector<int64_t>> targetShapes;
455  SmallVector<bool> collapsed;
456  auto hasCollapsibleType = [](OpOperand &operand) {
457  Type operandType = operand.get().getType();
458  if (auto memrefOperandType = dyn_cast_or_null<MemRefType>(operandType)) {
459  return memrefOperandType.getLayout().isIdentity();
460  }
461  if (auto tensorOperandType = dyn_cast<RankedTensorType>(operandType)) {
462  return tensorOperandType.getEncoding() == nullptr;
463  }
464  return false;
465  };
466  for (OpOperand &opOperand : genericOp->getOpOperands()) {
467  auto indexingMap = genericOp.getMatchingIndexingMap(&opOperand);
468  ArrayRef<int64_t> shape = genericOp.getShape(&opOperand);
469  if (!hasCollapsibleType(opOperand)) {
470  AffineMap newIndexingMap = indexingMap.replaceDimsAndSymbols(
471  dimReplacements, ArrayRef<AffineExpr>{}, oldDimToNewDimMap.size(), 0);
472  newIndexingMaps.push_back(newIndexingMap);
473  targetShapes.push_back(llvm::to_vector(shape));
474  collapsed.push_back(false);
475  reassociations.push_back({});
476  continue;
477  }
478  auto replacementInfo = dropUnitExtentFromOperandMetadata(
479  rewriter.getContext(), genericOp, &opOperand, oldDimToNewDimMap,
480  dimReplacements);
481  reassociations.push_back(replacementInfo.reassociation);
482  newIndexingMaps.push_back(replacementInfo.indexMap);
483  targetShapes.push_back(replacementInfo.targetShape);
484  collapsed.push_back(!(replacementInfo.indexMap.getNumResults() ==
485  indexingMap.getNumResults()));
486  }
487 
488  // Abort if the indexing maps of the result operation are not invertible
489  // (i.e. not legal) or if no dimension was reduced.
490  if (newIndexingMaps == indexingMaps ||
491  !inversePermutation(concatAffineMaps(newIndexingMaps)))
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
532  // type to the new type.
533  SmallVector<Value> resultReplacements;
534  for (auto [index, result] : llvm::enumerate(replacementOp.getResults())) {
535  unsigned opOperandIndex = index + replacementOp.getNumDpsInputs();
536  Value origDest = genericOp.getDpsInitOperand(index)->get();
537  if (!collapsed[opOperandIndex]) {
538  resultReplacements.push_back(result);
539  continue;
540  }
541  Value expandedValue = expandValue(rewriter, loc, result, origDest,
542  reassociations[opOperandIndex],
543  options.rankReductionStrategy);
544  resultReplacements.push_back(expandedValue);
545  }
546 
547  rewriter.replaceOp(genericOp, resultReplacements);
548  return success();
549 }
550 
551 namespace {
552 struct DropUnitDims : public OpRewritePattern<GenericOp> {
553  DropUnitDims(MLIRContext *context, ControlDropUnitDims options = {},
554  PatternBenefit benefit = 1)
555  : OpRewritePattern(context, benefit), options(std::move(options)) {}
556 
557  LogicalResult matchAndRewrite(GenericOp genericOp,
558  PatternRewriter &rewriter) const override {
559  return dropUnitDims(rewriter, genericOp, options);
560  }
561 
562 private:
564 };
565 } // namespace
566 
567 //===---------------------------------------------------------------------===//
568 // Drop dimensions that are unit-extents within tensor operations.
569 //===---------------------------------------------------------------------===//
570 
571 namespace {
572 struct DropPadUnitDims : public OpRewritePattern<tensor::PadOp> {
573  DropPadUnitDims(MLIRContext *context, ControlDropUnitDims options = {},
574  PatternBenefit benefit = 1)
575  : OpRewritePattern(context, benefit), options(std::move(options)) {}
576 
577  LogicalResult matchAndRewrite(tensor::PadOp padOp,
578  PatternRewriter &rewriter) const override {
579  // 1a. Get the allowed list of dimensions to drop from the `options`.
580  SmallVector<unsigned> allowedUnitDims = options.controlFn(padOp);
581  if (allowedUnitDims.empty()) {
582  return rewriter.notifyMatchFailure(
583  padOp, "control function returns no allowed unit dims to prune");
584  }
585 
586  if (padOp.getSourceType().getEncoding()) {
587  return rewriter.notifyMatchFailure(
588  padOp, "cannot collapse dims of tensor with encoding");
589  }
590 
591  // Fail for non-constant padding values. The body of the pad could
592  // depend on the padding indices and/or properties of the padded
593  // tensor so for now we fail.
594  // TODO: Support non-constant padding values.
595  Value paddingVal = padOp.getConstantPaddingValue();
596  if (!paddingVal) {
597  return rewriter.notifyMatchFailure(
598  padOp, "unimplemented: non-constant padding value");
599  }
600 
601  ArrayRef<int64_t> sourceShape = padOp.getSourceType().getShape();
602  int64_t padRank = sourceShape.size();
603 
604  auto isStaticZero = [](OpFoldResult f) {
605  std::optional<int64_t> maybeInt = getConstantIntValue(f);
606  return maybeInt && *maybeInt == 0;
607  };
608 
609  llvm::SmallDenseSet<unsigned> unitDimsFilter(allowedUnitDims.begin(),
610  allowedUnitDims.end());
611  llvm::SmallDenseSet<unsigned> unitDims;
612  SmallVector<int64_t> newShape;
613  SmallVector<OpFoldResult> newLowPad;
614  SmallVector<OpFoldResult> newHighPad;
615  for (const auto [dim, size, low, high] :
616  zip_equal(llvm::seq(static_cast<int64_t>(0), padRank), sourceShape,
617  padOp.getMixedLowPad(), padOp.getMixedHighPad())) {
618  if (unitDimsFilter.contains(dim) && size == 1 && isStaticZero(low) &&
619  isStaticZero(high)) {
620  unitDims.insert(dim);
621  } else {
622  newShape.push_back(size);
623  newLowPad.push_back(low);
624  newHighPad.push_back(high);
625  }
626  }
627 
628  if (unitDims.empty()) {
629  return rewriter.notifyMatchFailure(padOp, "no unit dims to collapse");
630  }
631 
632  ReassociationIndices reassociationGroup;
633  SmallVector<ReassociationIndices> reassociationMap;
634  int64_t dim = 0;
635  while (dim < padRank && unitDims.contains(dim))
636  reassociationGroup.push_back(dim++);
637  while (dim < padRank) {
638  assert(!unitDims.contains(dim) && "expected non unit-extent");
639  reassociationGroup.push_back(dim);
640  dim++;
641  // Fold all following dimensions that are unit-extent.
642  while (dim < padRank && unitDims.contains(dim))
643  reassociationGroup.push_back(dim++);
644  reassociationMap.push_back(reassociationGroup);
645  reassociationGroup.clear();
646  }
647 
648  Value collapsedSource =
649  collapseValue(rewriter, padOp.getLoc(), padOp.getSource(), newShape,
650  reassociationMap, options.rankReductionStrategy);
651 
652  auto newPadOp = rewriter.create<tensor::PadOp>(
653  padOp.getLoc(), /*result=*/Type(), collapsedSource, newLowPad,
654  newHighPad, paddingVal, padOp.getNofold());
655 
656  Value dest = padOp.getResult();
657  if (options.rankReductionStrategy ==
659  SmallVector<OpFoldResult> expandedSizes;
660  int64_t numUnitDims = 0;
661  for (auto dim : llvm::seq(static_cast<int64_t>(0), padRank)) {
662  if (unitDims.contains(dim)) {
663  expandedSizes.push_back(rewriter.getIndexAttr(1));
664  numUnitDims++;
665  continue;
666  }
667  expandedSizes.push_back(tensor::getMixedSize(
668  rewriter, padOp.getLoc(), newPadOp, dim - numUnitDims));
669  }
670  dest = rewriter.create<tensor::EmptyOp>(
671  padOp.getLoc(), expandedSizes,
672  padOp.getResultType().getElementType());
673  }
674 
675  Value expandedValue =
676  expandValue(rewriter, padOp.getLoc(), newPadOp.getResult(), dest,
677  reassociationMap, options.rankReductionStrategy);
678  rewriter.replaceOp(padOp, expandedValue);
679  return success();
680  }
681 
682 private:
684 };
685 } // namespace
686 
687 namespace {
688 /// Convert `extract_slice` operations to rank-reduced versions.
689 struct RankReducedExtractSliceOp
690  : public OpRewritePattern<tensor::ExtractSliceOp> {
692 
693  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
694  PatternRewriter &rewriter) const override {
695  RankedTensorType resultType = sliceOp.getType();
696  SmallVector<OpFoldResult> targetShape;
697  for (auto size : resultType.getShape())
698  targetShape.push_back(rewriter.getIndexAttr(size));
699  auto reassociation = getReassociationMapForFoldingUnitDims(targetShape);
700  if (!reassociation ||
701  reassociation->size() == static_cast<size_t>(resultType.getRank()))
702  return failure();
703 
704  SmallVector<OpFoldResult> offsets = sliceOp.getMixedOffsets();
705  SmallVector<OpFoldResult> strides = sliceOp.getMixedStrides();
706  SmallVector<OpFoldResult> sizes = sliceOp.getMixedSizes();
707  auto rankReducedType = cast<RankedTensorType>(
708  tensor::ExtractSliceOp::inferCanonicalRankReducedResultType(
709  reassociation->size(), sliceOp.getSourceType(), offsets, sizes,
710  strides));
711 
712  Location loc = sliceOp.getLoc();
713  Value newSlice = rewriter.create<tensor::ExtractSliceOp>(
714  loc, rankReducedType, sliceOp.getSource(), offsets, sizes, strides);
715  rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
716  sliceOp, resultType, newSlice, *reassociation);
717  return success();
718  }
719 };
720 
721 /// Convert `insert_slice` operations to rank-reduced versions.
722 /// This patterns works with both InsertSliceOp and ParallelInsertSliceOp.
723 template <typename InsertOpTy>
724 struct RankReducedInsertSliceOp : public OpRewritePattern<InsertOpTy> {
726 
727  LogicalResult matchAndRewrite(InsertOpTy insertSliceOp,
728  PatternRewriter &rewriter) const override {
729  RankedTensorType sourceType = insertSliceOp.getSourceType();
730  SmallVector<OpFoldResult> targetShape;
731  for (auto size : sourceType.getShape())
732  targetShape.push_back(rewriter.getIndexAttr(size));
733  auto reassociation = getReassociationMapForFoldingUnitDims(targetShape);
734  if (!reassociation ||
735  reassociation->size() == static_cast<size_t>(sourceType.getRank()))
736  return failure();
737 
738  Location loc = insertSliceOp.getLoc();
739  tensor::CollapseShapeOp reshapedSource;
740  {
741  OpBuilder::InsertionGuard g(rewriter);
742  // The only difference between InsertSliceOp and ParallelInsertSliceOp
743  // is the insertion point is just before the ParallelCombiningOp in the
744  // parallel case.
745  if (std::is_same<InsertOpTy, tensor::ParallelInsertSliceOp>::value)
746  rewriter.setInsertionPoint(insertSliceOp->getParentOp());
747  reshapedSource = rewriter.create<tensor::CollapseShapeOp>(
748  loc, insertSliceOp.getSource(), *reassociation);
749  }
750  rewriter.replaceOpWithNewOp<InsertOpTy>(
751  insertSliceOp, reshapedSource, insertSliceOp.getDest(),
752  insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
753  insertSliceOp.getMixedStrides());
754  return success();
755  }
756 };
757 } // namespace
758 
759 /// Patterns that are used to canonicalize the use of unit-extent dims for
760 /// broadcasting.
761 static void
764  auto *context = patterns.getContext();
765  patterns.add<DropUnitDims>(context, options);
766  patterns.add<DropPadUnitDims>(context, options);
767  // TODO: Patterns unrelated to unit dim folding should be factored out.
768  patterns.add<RankReducedExtractSliceOp,
769  RankReducedInsertSliceOp<tensor::InsertSliceOp>,
770  RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>(
771  context);
772  linalg::FillOp::getCanonicalizationPatterns(patterns, context);
773  tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
774  tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
775  tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
779 }
780 
781 static void
784  auto *context = patterns.getContext();
785  options.rankReductionStrategy =
787  patterns.add<DropUnitDims>(context, options);
788  patterns.add<DropPadUnitDims>(context, options);
789  // TODO: Patterns unrelated to unit dim folding should be factored out.
790  linalg::FillOp::getCanonicalizationPatterns(patterns, context);
791  tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
795 }
796 
799  if (options.rankReductionStrategy ==
802  } else if (options.rankReductionStrategy ==
804  ReassociativeReshape) {
806  }
807 }
808 
810  RewritePatternSet &patterns) {
811  patterns.add<MoveInitOperandsToInput>(patterns.getContext());
812 }
813 
814 namespace {
815 /// Pass that removes unit-extent dims within generic ops.
816 struct LinalgFoldUnitExtentDimsPass
817  : public impl::LinalgFoldUnitExtentDimsPassBase<
818  LinalgFoldUnitExtentDimsPass> {
819  using impl::LinalgFoldUnitExtentDimsPassBase<
820  LinalgFoldUnitExtentDimsPass>::LinalgFoldUnitExtentDimsPassBase;
821  void runOnOperation() override {
822  Operation *op = getOperation();
823  MLIRContext *context = op->getContext();
824  RewritePatternSet patterns(context);
826  if (useRankReducingSlices) {
827  options.rankReductionStrategy = linalg::ControlDropUnitDims::
829  }
832  (void)applyPatternsAndFoldGreedily(op, std::move(patterns));
833  }
834 };
835 } // namespace
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:237
Base type for affine expression.
Definition: AffineExpr.h:69
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:47
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
unsigned getNumSymbols() const
Definition: AffineMap.cpp:382
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:391
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:484
This class represents an argument of a Block.
Definition: Value.h:319
Block represents an ordered list of Operations.
Definition: Block.h:30
BlockArgument addArgument(Type type, Location loc)
Add one value to the argument list.
Definition: Block.cpp:152
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:124
MLIRContext * getContext() const
Definition: Builders.h:55
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
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:63
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:350
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:555
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:400
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:437
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:522
void setInsertionPointAfterValue(Value val)
Sets the insertion point to the node after the specified value.
Definition: Builders.h:423
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:464
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
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:785
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Definition: Region.h:26
MLIRContext * getContext() const
Definition: PatternMatch.h:822
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
Definition: PatternMatch.h:846
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:718
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:536
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:129
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
void populateMoveInitOperandsToInputPattern(RewritePatternSet &patterns)
A pattern that converts init operands to input operands.
SmallVector< NamedAttribute > getPrunedAttributeList(OpTy op)
Returns an attribute list that excludes pre-defined attributes.
Definition: Utils.h:371
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:886
void populateFoldUnitExtentDimsPatterns(RewritePatternSet &patterns, ControlDropUnitDims &options)
Patterns to fold unit-extent dimensions in operands/results of linalg ops on tensors via reassociativ...
LogicalResult dropUnitDims(RewriterBase &rewriter, GenericOp genericOp, const ControlDropUnitDims &options)
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 tensor.
OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value, int64_t dim)
Return the dimension of the given tensor value.
Definition: TensorOps.cpp:51
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given tensor value.
Definition: TensorOps.cpp:61
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
std::optional< int64_t > getConstantIntValue(OpFoldResult ofr)
If ofr is a constant integer or an IntegerAttr, return the integer.
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
Definition: AffineMap.cpp:753
bool succeeded(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a success value.
Definition: LogicalResult.h:68
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
AffineMap concatAffineMaps(ArrayRef< AffineMap > maps)
Concatenates a list of maps into a single AffineMap, stepping over potentially empty maps.
Definition: AffineMap.cpp:798
LogicalResult applyPatternsAndFoldGreedily(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...
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)
Definition: AffineExpr.cpp:623
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:599
Compute the modified metadata for an operands of operation whose unit dims are being dropped.
SmallVector< ReassociationIndices > reassociation
SmallVector< int64_t > targetShape
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:358
Transformation to drop unit-extent dimensions from linalg.generic operations.
Definition: Transforms.h:474