1 //===- LinalgTransforms.cpp - Linalg transformations as patterns ----------===//
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 logic and helpers to expose Linalg transforms as rewrite
10 // patterns.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
15 #include "mlir/Dialect/Affine/Utils.h"
16 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
17 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
18 #include "mlir/Dialect/Linalg/Utils/Utils.h"
19 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
20 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
21 #include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
22 #include "mlir/Dialect/Vector/VectorOps.h"
23 #include "mlir/IR/AffineExpr.h"
24 #include "mlir/IR/Matchers.h"
25 #include "mlir/Pass/Pass.h"
26 #include "mlir/Support/LLVM.h"
27 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
28 #include "llvm/ADT/ScopeExit.h"
29 #include "llvm/Support/Debug.h"
30 #include "llvm/Support/raw_ostream.h"
31 #include <type_traits>
32 
33 #define DEBUG_TYPE "linalg-transforms"
34 
35 using namespace mlir;
36 using namespace mlir::edsc;
37 using namespace mlir::edsc::intrinsics;
38 using namespace mlir::linalg;
39 
40 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
41 
42 //===----------------------------------------------------------------------===//
43 // Transformations exposed as rewrite patterns.
44 //===----------------------------------------------------------------------===//
45 // Marker used as attribute name in generated Linalg rewriting transformations.
46 const StringLiteral mlir::linalg::LinalgTransforms::kLinalgTransformMarker =
47     "__internal_linalg_transform__";
48 
49 mlir::linalg::LinalgTransformationFilter::LinalgTransformationFilter(
50     ArrayRef<Identifier> matchDisjunction, Optional<Identifier> replacement)
51     : matchDisjunction(matchDisjunction.begin(), matchDisjunction.end()),
52       replacement(replacement) {}
53 
54 mlir::linalg::LinalgTransformationFilter::LinalgTransformationFilter(
55     FilterFunction f, ArrayRef<Identifier> matchDisjunction,
56     Optional<Identifier> replacement)
57     : filters(),
58       matchDisjunction(matchDisjunction.begin(), matchDisjunction.end()),
59       replacement(replacement) {
60   if (f)
61     filters.push_back(f);
62 }
63 
64 LogicalResult mlir::linalg::LinalgTransformationFilter::checkAndNotify(
65     PatternRewriter &rewriter, Operation *op) const {
66   if (llvm::any_of(filters,
67                    [&](const FilterFunction &f) { return failed(f(op)); }))
68     return failure();
69 
70   auto attr = op->template getAttrOfType<StringAttr>(
71       LinalgTransforms::kLinalgTransformMarker);
72 
73   if (!attr) {
74     // 1. Has no filter case and matchDisjunction is empty.
75     if (matchDisjunction.empty())
76       return success();
77 
78     // 2. Has no filter but was expecting a filter.
79     return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
80       diag << " does not have any filter from list: ";
81       interleaveComma(matchDisjunction, diag);
82     });
83   }
84 
85   // 4. Match explicit filter.
86   for (auto filter : matchDisjunction)
87     if (attr.getValue() == filter)
88       return success();
89 
90   // 5. Fail to match.
91   return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
92     diag << " does not have any filter from list: ";
93     interleaveComma(matchDisjunction, diag);
94   });
95 }
96 
97 void mlir::linalg::LinalgTransformationFilter::
98     replaceLinalgTransformationFilter(PatternRewriter &rewriter,
99                                       Operation *op) const {
100   if (replacement.hasValue())
101     op->setAttr(LinalgTransforms::kLinalgTransformMarker,
102                 rewriter.getStringAttr(replacement.getValue()));
103   else
104     op->removeAttr(Identifier::get(LinalgTransforms::kLinalgTransformMarker,
105                                    rewriter.getContext()));
106 }
107 
108 LinalgTilingOptions &
109 mlir::linalg::LinalgTilingOptions::setTileSizes(ArrayRef<int64_t> ts) {
110   SmallVector<int64_t, 4> tileSizes(ts.begin(), ts.end());
111   tileSizeComputationFunction = [tileSizes](OpBuilder &b, Operation *op) {
112     OpBuilder::InsertionGuard guard(b);
113     b.setInsertionPointToStart(
114         &op->getParentOfType<FuncOp>().getBody().front());
115     return llvm::to_vector<4>(map_range(tileSizes, [&](int64_t s) {
116       Value v = b.create<ConstantIndexOp>(op->getLoc(), s);
117       return v;
118     }));
119   };
120   return *this;
121 }
122 
123 /// Try to compute a static bounding box for `operand`
124 /// Return success if either:
125 ///   1. The operand is already statically shaped, `result` is left unchanged.
126 ///   2. The operand is (partially) dynamic, `result` is the result of a freshly
127 ///      created PadTensorOp.
128 /// Return failure if the operand cannot be padded to a static shape.
129 static LogicalResult padOperandToSmallestStaticBoundingBox(
130     PatternRewriter &rewriter, linalg::LinalgOp opToPad, OpOperand &operand,
131     const LinalgTilingOptions &options, Value &result) {
132   auto tensorType = operand.get().getType().cast<RankedTensorType>();
133   // Already static shape, no need to pad.
134   if (tensorType.hasStaticShape())
135     return success();
136   auto subtensor = operand.get().getDefiningOp<SubTensorOp>();
137   // Not a subtensor, cannot construct a static bounding box.
138   if (!subtensor)
139     return failure();
140   SmallVector<int64_t> staticSizes;
141   staticSizes.reserve(tensorType.getRank());
142   auto shapedOp =
143       cast<OffsetSizeAndStrideOpInterface>(subtensor.getOperation());
144   for (auto size : shapedOp.getMixedSizes()) {
145     auto indexAttr = size.is<Attribute>()
146                          ? size.get<Attribute>().dyn_cast<IntegerAttr>()
147                          : linalg::getSmallestBoundingIndex(size.get<Value>());
148     // SmallestBoundingIndex must exist for all sizes.
149     // For now return an error if we can't find it.
150     if (!indexAttr)
151       return rewriter.notifyMatchFailure(
152           opToPad, "No constant bounding box can be found for padding");
153     staticSizes.push_back(indexAttr.getInt());
154   }
155   Value pad = options.paddingValueComputationFunction(rewriter, operand);
156   auto staticTensorType =
157       RankedTensorType::get(staticSizes, tensorType.getElementType());
158   result = linalg::PadTensorOp::createPadHighOp(
159       staticTensorType, operand.get(), pad, opToPad->getLoc(), rewriter);
160   return success();
161 }
162 
163 // Try to create a static bounding box around each operand of `res.op`.
164 // If successful, `res.op` is rewritten in static form with padded operands.
165 // `res.op` is updated to the cloned static form of the op on success.
166 static LogicalResult rewriteAsPaddedOp(PatternRewriter &rewriter,
167                                        TiledLinalgOp &res,
168                                        const LinalgTilingOptions &options) {
169   LinalgOp opToPad = res.op;
170   Location loc = opToPad->getLoc();
171 
172   // If the op is fully static, it does not need padding.
173   // TODO: there are cases where we may still want to pad to larger sizes.
174   if (llvm::all_of(opToPad.getShapedOperands(), [](Value v) {
175         return v.getType().cast<RankedTensorType>().hasStaticShape();
176       }))
177     return success();
178 
179   OpBuilder::InsertionGuard g(rewriter);
180   // Set IP after op because we also take the dims of the original output.
181   rewriter.setInsertionPointAfter(opToPad);
182   // Make a copy of the shaped operands and update it.
183   SmallVector<Value> newOperands;
184   newOperands.reserve(opToPad.getNumShapedOperands());
185   for (OpOperand &operand : opToPad.getShapedOpOperands()) {
186     Value paddedOperand;
187     // If padding was requested but the shape cannot be bounded statically then
188     // the pattern fails to apply.
189     if (failed(padOperandToSmallestStaticBoundingBox(rewriter, opToPad, operand,
190                                                      options, paddedOperand))) {
191       return failure();
192     }
193     newOperands.push_back(paddedOperand ? paddedOperand : operand.get());
194   }
195 
196   // Clone `opToPad` to operate on the statically padded shapes.
197   auto resultTensorTypes =
198       ValueRange(newOperands).take_back(opToPad.getNumOutputs()).getTypes();
199   ValueRange otherOperands = opToPad.getAssumedNonShapedOperands();
200   newOperands.append(otherOperands.begin(), otherOperands.end());
201   linalg::LinalgOp paddedOp =
202       opToPad.clone(rewriter, loc, resultTensorTypes, newOperands);
203 
204   // Recover the subtensor out of the new static results. This keeps the
205   // original linalg op around because it uses the dims of the original results.
206   // This later folds away.
207   SmallVector<Value> paddedSubviewResults;
208   paddedSubviewResults.reserve(opToPad->getNumResults());
209   SetVector<Operation *> newUsersOfOpToPad;
210   for (auto it : llvm::zip(opToPad->getResults(), paddedOp->getResults())) {
211     auto rank = std::get<0>(it).getType().cast<RankedTensorType>().getRank();
212     SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
213     auto sizes = llvm::to_vector<4>(llvm::map_range(
214         llvm::seq<unsigned>(0, rank), [&](unsigned d) -> OpFoldResult {
215           auto dimOp = rewriter.create<memref::DimOp>(loc, std::get<0>(it), d);
216           newUsersOfOpToPad.insert(dimOp);
217           return dimOp.getResult();
218         }));
219     SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
220     paddedSubviewResults.push_back(rewriter.create<SubTensorOp>(
221         loc, std::get<1>(it), offsets, sizes, strides));
222   }
223   // Replace the transient `opToPad` locally, except for uses that we just
224   // created for the purpose of extracting the dims.
225   rewriter.replaceOpWithIf(opToPad, paddedSubviewResults, [&](OpOperand &opOp) {
226     return !newUsersOfOpToPad.contains(opOp.getOwner());
227   });
228 
229   res = TiledLinalgOp{paddedOp, res.loops, res.tensorResults};
230   return success();
231 }
232 
233 /// Linalg base tiling pattern.
234 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
235     StringRef opName, MLIRContext *context, LinalgTilingOptions options,
236     LinalgTransformationFilter filter, PatternBenefit benefit)
237     : RewritePattern(opName, benefit, context), filter(filter),
238       options(options) {}
239 
240 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
241     MLIRContext *context, LinalgTilingOptions options,
242     LinalgTransformationFilter filter, PatternBenefit benefit)
243     : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter),
244       options(options) {}
245 
246 LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
247     Operation *op, PatternRewriter &rewriter, TiledLinalgOp &result) const {
248   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
249   if (!linalgOp)
250     return failure();
251   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
252     return failure();
253 
254   Optional<TiledLinalgOp> res = tileLinalgOp(rewriter, linalgOp, options);
255 
256   if (!res)
257     return failure();
258 
259   // Setup RAII guard to return properly.
260   LinalgOp tiledOp = res->op;
261   auto guard = llvm::make_scope_exit([&]() {
262     // Return relevant information to derived pattern.
263     result = *res;
264     // Replace filter on both tiledOp and tiledAndPaddedOp, if necessary.
265     filter.replaceLinalgTransformationFilter(rewriter, tiledOp);
266     if (tiledOp != res->op)
267       filter.replaceLinalgTransformationFilter(rewriter, res->op);
268   });
269 
270   // Consider padding on the fly only if the op has tensor semantics.
271   if (!options.paddingValueComputationFunction ||
272       !linalgOp.hasTensorSemantics())
273     return success();
274 
275   // Try to pad on the fly by rewriting res->op as a padded op.
276   if (failed(rewriteAsPaddedOp(rewriter, *res, options))) {
277     // Set so RAII guard does not propagate TiledLinalgOp to `result`.
278     return failure();
279   }
280 
281   // Do not perform replacement of `linalgOp`, let the derived patterns
282   // do this as they see fit, from the resulting TiledLinalgOp.
283   return success();
284 }
285 
286 static ValueRange getTiledOpResult(TiledLinalgOp tiledOp) {
287   if (tiledOp.loops.empty())
288     return tiledOp.op.getOperation()->getResults();
289   return tiledOp.loops.front()->getResults();
290 }
291 
292 static ValueRange
293 getTiledAndFusedOpResult(TiledAndFusedLinalgOps tiledAndFusedOp) {
294   if (tiledAndFusedOp.fusedLoops.empty())
295     return tiledAndFusedOp.op.getOperation()->getResults();
296   return tiledAndFusedOp.fusedLoops.front()->getResults();
297 }
298 
299 mlir::linalg::LinalgBaseTileAndFusePattern::LinalgBaseTileAndFusePattern(
300     StringRef opName, MLIRContext *context,
301     const LinalgDependenceGraph &dependenceGraph,
302     LinalgTilingOptions tilingOptions, LinalgFusionOptions fusionOptions,
303     LinalgTransformationFilter filter, LinalgTransformationFilter fusedOpMarker,
304     LinalgTransformationFilter originalOpMarker, PatternBenefit benefit)
305     : RewritePattern(opName, benefit, context, {}),
306       dependenceGraph(dependenceGraph), tilingOptions(tilingOptions),
307       fusionOptions(fusionOptions), filter(filter),
308       fusedOpMarker(fusedOpMarker), originalOpMarker(originalOpMarker) {}
309 
310 LogicalResult mlir::linalg::LinalgBaseTileAndFusePattern::matchAndRewrite(
311     Operation *op, PatternRewriter &rewriter) const {
312   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
313   // TODO: remove hasIndexSemantics check once index ops are supported.
314   if (!linalgOp || linalgOp.hasIndexSemantics())
315     return failure();
316   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
317     return failure();
318 
319   DenseSet<Operation *> producers;
320   producers.insert(linalgOp);
321   for (auto dependence : dependenceGraph.getDependentOperationsInto(linalgOp)) {
322     Optional<unsigned> operandNumber = dependence.getIndexingOpViewOperandNum();
323     // When looking at dependences into, indexingOp is always OpOperand. We
324     // could assert, but continue if this is not the case.
325     if (!operandNumber)
326       continue;
327     if (!fusionOptions.indicesToFuse.count(operandNumber.getValue()))
328       continue;
329     if (isa<LinalgOp>(dependence.getDependentOp()))
330       producers.insert(dependence.getDependentOp());
331   }
332 
333   SmallVector<LinalgOp, 1> fusionOps;
334   for (auto it = op->getBlock()->begin(), ie = Block::iterator(op); it != ie;
335        ++it) {
336     auto producerLinalgOp = dyn_cast<LinalgOp>(&(*it));
337     if (producerLinalgOp && producers.count(producerLinalgOp))
338       fusionOps.push_back(producerLinalgOp);
339   }
340   fusionOps.push_back(linalgOp);
341 
342   SmallVector<Value, 4> tileSizes =
343       tilingOptions.tileSizeComputationFunction(rewriter, op);
344   LinalgTilingOptions instanceTilingOptions = tilingOptions;
345   instanceTilingOptions.setTileSizes(tileSizes);
346   Optional<TiledAndFusedLinalgOps> tiledAndFusedOps = tileAndFuseLinalgOps(
347       rewriter, fusionOps, dependenceGraph, instanceTilingOptions);
348   if (!tiledAndFusedOps)
349     return failure();
350 
351   // Tile the unfused loops;
352   SmallVector<Value, 4> unfusedLoopTileSizes;
353   Value zero = rewriter.create<ConstantIndexOp>(op->getLoc(), 0);
354   for (auto tileSize : enumerate(tileSizes)) {
355     if (tiledAndFusedOps->fusedLoopDims.count(tileSize.index()))
356       unfusedLoopTileSizes.push_back(zero);
357     else
358       unfusedLoopTileSizes.push_back(tileSize.value());
359   }
360   // Tile the loop only if there is a non-zero tile size.
361   if (unfusedLoopTileSizes.size() > linalgOp.getNumLoops())
362     unfusedLoopTileSizes.resize(linalgOp.getNumLoops());
363   if (llvm::any_of(unfusedLoopTileSizes, [](Value val) {
364         if (auto cst = val.getDefiningOp<ConstantIndexOp>())
365           return cst.getValue() != 0;
366         return true;
367       })) {
368     LinalgTilingOptions unfusedTilingOptions = tilingOptions;
369     unfusedTilingOptions.setTileSizes(unfusedLoopTileSizes);
370     Optional<TiledLinalgOp> unfusedTiledOp =
371         tileLinalgOp(rewriter, tiledAndFusedOps->op, unfusedTilingOptions);
372     if (!unfusedTiledOp)
373       return failure();
374     rewriter.replaceOp(tiledAndFusedOps->op,
375                        getTiledOpResult(unfusedTiledOp.getValue()));
376     tiledAndFusedOps->op = unfusedTiledOp->op;
377   }
378   op->replaceAllUsesWith(getTiledAndFusedOpResult(tiledAndFusedOps.getValue()));
379 
380   filter.replaceLinalgTransformationFilter(rewriter,
381                                            tiledAndFusedOps->op.getOperation());
382   for (auto fusedOp : tiledAndFusedOps->fusedProducers) {
383     fusedOpMarker.replaceLinalgTransformationFilter(rewriter,
384                                                     fusedOp.getOperation());
385   }
386   for (auto origProducerOp : ArrayRef<LinalgOp>(fusionOps).drop_back()) {
387     originalOpMarker.replaceLinalgTransformationFilter(
388         rewriter, origProducerOp.getOperation());
389   }
390   rewriter.updateRootInPlace(op, [&]() {
391     originalOpMarker.replaceLinalgTransformationFilter(rewriter, op);
392   });
393   return success();
394 }
395 
396 /// Linalg base interchange pattern.
397 mlir::linalg::LinalgBaseInterchangePattern::LinalgBaseInterchangePattern(
398     StringRef opName, MLIRContext *context,
399     ArrayRef<unsigned> interchangeVector, LinalgTransformationFilter filter,
400     PatternBenefit benefit)
401     : RewritePattern(opName, benefit, context, {}), filter(filter),
402       interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
403 
404 LogicalResult mlir::linalg::LinalgBaseInterchangePattern::matchAndRewrite(
405     Operation *op, PatternRewriter &rewriter) const {
406   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
407   if (!linalgOp)
408     return failure();
409   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
410     return failure();
411   if (failed(interchangeGenericLinalgOpPrecondition(op, interchangeVector)))
412     return failure();
413 
414   // TODO: figure out how this interplays with named ops. In particular this
415   // should break the named op property.
416   rewriter.updateRootInPlace(op, [&]() {
417     interchange(rewriter, linalgOp, interchangeVector);
418     // New filter if specified.
419     filter.replaceLinalgTransformationFilter(rewriter, op);
420   });
421   return success();
422 }
423 
424 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
425     StringRef opName, MLIRContext *context, LinalgPromotionOptions options,
426     LinalgTransformationFilter filter, PatternBenefit benefit)
427     : RewritePattern(opName, benefit, context, {}), filter(filter),
428       options(options) {}
429 
430 LogicalResult mlir::linalg::LinalgBasePromotionPattern::matchAndRewrite(
431     Operation *op, PatternRewriter &rewriter) const {
432   if (failed(filter.checkAndNotify(rewriter, op)))
433     return failure();
434   if (failed(promoteSubviewsPrecondition(op, options)))
435     return failure();
436 
437   // TODO: We cannot use root update here. This pattern is creating other ops,
438   // so if the promotion fails, those need to be cleaned up, which doesnt seem
439   // to be happening here. So to fail properly, we should be cloning the op and
440   // deleting the previous op. This needs more investigation.
441   rewriter.startRootUpdate(op);
442   Optional<LinalgOp> promotedOp = promoteSubViews(rewriter, op, options);
443   if (!promotedOp) {
444     rewriter.cancelRootUpdate(op);
445     return op->emitError("subview promotion failed");
446   }
447   rewriter.finalizeRootUpdate(op);
448   filter.replaceLinalgTransformationFilter(rewriter, op);
449   return success();
450 }
451 
452 mlir::linalg::LinalgBaseVectorizationPattern::LinalgBaseVectorizationPattern(
453     MLIRContext *context, LinalgTransformationFilter filter,
454     PatternBenefit benefit)
455     : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter) {}
456 
457 mlir::linalg::LinalgBaseVectorizationPattern::LinalgBaseVectorizationPattern(
458     StringRef opName, MLIRContext *context, LinalgTransformationFilter filter,
459     PatternBenefit benefit)
460     : RewritePattern(opName, benefit, context, {}), filter(filter) {}
461 
462 LogicalResult mlir::linalg::LinalgBaseVectorizationPattern::matchAndRewrite(
463     Operation *op, PatternRewriter &rewriter) const {
464   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
465   if (!linalgOp)
466     return failure();
467   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
468     return failure();
469   SmallVector<Value> newResults;
470   if (failed(vectorizeLinalgOp(rewriter, op, newResults)))
471     return failure();
472   if (!newResults.empty())
473     rewriter.replaceOp(op, newResults);
474   else
475     rewriter.eraseOp(op);
476   return success();
477 }
478 
479 LogicalResult mlir::linalg::applyStagedPatterns(
480     Operation *op, ArrayRef<FrozenRewritePatternSet> stage1Patterns,
481     const FrozenRewritePatternSet &stage2Patterns,
482     function_ref<LogicalResult(Operation *)> stage3Lambda) {
483   unsigned iteration = 0;
484   (void)iteration;
485   for (const auto &patterns : stage1Patterns) {
486     LLVM_DEBUG(DBGS() << "Before 1st stage, iter: " << ++iteration << "\n"
487                       << *op);
488     if (failed(applyPatternsAndFoldGreedily(op, patterns))) {
489       LLVM_DEBUG(DBGS() << "Underlying first stage rewrite did not converge");
490       return failure();
491     }
492     LLVM_DEBUG(DBGS() << "After 1st stage, iter: " << ++iteration << "\n"
493                       << *op);
494     if (failed(applyPatternsAndFoldGreedily(op, stage2Patterns))) {
495       LLVM_DEBUG(DBGS() << "Underlying 2nd stage rewrite did not converge");
496       return failure();
497     }
498     LLVM_DEBUG(DBGS() << "After 2nd stage, iter : " << iteration << "\n"
499                       << *op);
500     if (stage3Lambda) {
501       if (failed(stage3Lambda(op)))
502         return failure();
503       LLVM_DEBUG(DBGS() << "After 3rd stage, iter : " << iteration << "\n"
504                         << *op);
505     }
506   }
507   return success();
508 }
509 
510 /// Traverse the `dims` and substitute known min or max expressions returned by
511 /// the lambda |getMinMaxExpr|.
512 static AffineMap substitute(AffineMap map, SmallVectorImpl<Value> &dims,
513                             SmallVectorImpl<Value> &symbols,
514                             GetMinMaxExprFn getMinMaxExpr) {
515   auto exprs = llvm::to_vector<4>(map.getResults());
516   for (AffineExpr &expr : exprs) {
517     bool substituted = true;
518     while (substituted) {
519       substituted = false;
520       for (unsigned dimIdx = 0; dimIdx < dims.size(); ++dimIdx) {
521         Value dim = dims[dimIdx];
522         auto minMax = getMinMaxExpr(dim, dims, symbols);
523         if (!minMax)
524           continue;
525         AffineExpr dimExpr = getAffineDimExpr(dimIdx, expr.getContext());
526         LLVM_DEBUG(DBGS() << "Subst: " << dim << " @ " << dimExpr << "\n");
527         LLVM_DEBUG(DBGS() << "Before: " << expr << "\n");
528         // Substitute occurrences of `dimExpr` by either the min expression or
529         // the max expression depending on whether the value is used with a
530         // positive or negative  coefficient.
531         AffineExpr substitutedExpr =
532             substWithMin(expr, dimExpr, minMax->first, minMax->second);
533         LLVM_DEBUG(DBGS() << "After: " << substitutedExpr << "\n");
534         substituted = (substitutedExpr != expr);
535         expr = substitutedExpr;
536       }
537     }
538 
539     // Cleanup and simplify the results.
540     // This needs to happen outside of the loop iterating on dims.size() since
541     // it modifies dims.
542     SmallVector<Value, 4> operands(dims.begin(), dims.end());
543     operands.append(symbols.begin(), symbols.end());
544     auto map = AffineMap::get(dims.size(), symbols.size(), exprs,
545                               exprs.front().getContext());
546 
547     LLVM_DEBUG({
548       DBGS() << "Map to simplify: " << map << "\n";
549       DBGS() << "Operands:\n";
550       for (Value v : operands)
551         DBGS() << v << "\n";
552     });
553 
554     // Pull in affine.apply operations and compose them fully into the
555     // result.
556     fullyComposeAffineMapAndOperands(&map, &operands);
557     canonicalizeMapAndOperands(&map, &operands);
558     map = simplifyAffineMap(map);
559     // Assign the results.
560     exprs.assign(map.getResults().begin(), map.getResults().end());
561     dims.assign(operands.begin(), operands.begin() + map.getNumDims());
562     symbols.assign(operands.begin() + map.getNumDims(), operands.end());
563 
564     LLVM_DEBUG(DBGS() << "Map simplified: " << map << "\n");
565   }
566 
567   assert(!exprs.empty() && "Unexpected empty exprs");
568   return AffineMap::get(dims.size(), symbols.size(), exprs, map.getContext());
569 }
570 
571 /// Traverse the dims of the AffineMap of `affineMinOp` and substitute
572 /// dimensions with known range by new expressions involving the min or max
573 /// expression:
574 ///   - If the AffineDimExpr mapped to a known value has a positive sign, it
575 ///     is replaced by the min expression.
576 ///   - If the AffineDimExpr mapped to a known value has a negative sign, it is
577 ///     replaced by the max expression.
578 /// All known values are iteratively replaced.
579 /// This is used as an intermediate step in computing bounding boxes and
580 /// canonicalize AffineMinOps. All dim and symbol operands are assumed to have
581 /// positive values (positive orthant assumptions).
582 /// Return a new AffineMap, dims and symbols that have been canonicalized and
583 /// simplified.
584 AffineMapAndOperands
585 mlir::linalg::substituteMin(AffineMinOp affineMinOp,
586                             GetMinMaxExprFn getMinMaxExpr) {
587   AffineMapAndOperands res{affineMinOp.getAffineMap(),
588                            SmallVector<Value>(affineMinOp.getDimOperands()),
589                            SmallVector<Value>(affineMinOp.getSymbolOperands())};
590   res.map = substitute(affineMinOp.getAffineMap(), res.dims, res.symbols,
591                        getMinMaxExpr);
592   return res;
593 }
594 
595 LogicalResult AffineMinRangeCanonicalizationPattern::matchAndRewrite(
596     AffineMinOp minOp, PatternRewriter &rewriter) const {
597   LLVM_DEBUG(DBGS() << "Canonicalize AffineMinSCF: " << *minOp.getOperation()
598                     << "\n");
599 
600   auto affineMapAndOperands = substituteMin(minOp, getMinMaxFn);
601   AffineMap map = affineMapAndOperands.map;
602 
603   LLVM_DEBUG(DBGS() << "Resulting map: " << map << "\n");
604 
605   // Check whether any of the expressions, when subtracted from all other
606   // expressions, produces only >= 0 constants. If so, it is the min.
607   for (auto e : minOp.getAffineMap().getResults()) {
608     LLVM_DEBUG(DBGS() << "Candidate min: " << e << "\n");
609     if (!e.isSymbolicOrConstant())
610       continue;
611 
612     auto isNonPositive = [](AffineExpr e) {
613       if (auto cst = e.dyn_cast<AffineConstantExpr>())
614         return cst.getValue() < 0;
615       return true;
616     };
617 
618     // Build the subMap and check everything is statically known to be
619     // positive.
620     SmallVector<AffineExpr, 4> subExprs;
621     subExprs.reserve(map.getNumResults());
622     for (auto ee : map.getResults())
623       subExprs.push_back(ee - e);
624     MLIRContext *ctx = minOp.getContext();
625     AffineMap subMap = simplifyAffineMap(
626         AffineMap::get(map.getNumDims(), map.getNumSymbols(), subExprs, ctx));
627     LLVM_DEBUG(DBGS() << "simplified subMap: " << subMap << "\n");
628     if (llvm::any_of(subMap.getResults(), isNonPositive))
629       continue;
630 
631     // Static min found.
632     if (auto cst = e.dyn_cast<AffineConstantExpr>()) {
633       rewriter.replaceOpWithNewOp<ConstantIndexOp>(minOp, cst.getValue());
634     } else {
635       auto resultMap = AffineMap::get(0, map.getNumSymbols(), {e}, ctx);
636       SmallVector<Value> resultOperands = affineMapAndOperands.dims;
637       llvm::append_range(resultOperands, affineMapAndOperands.symbols);
638       canonicalizeMapAndOperands(&resultMap, &resultOperands);
639       resultMap = simplifyAffineMap(resultMap);
640       rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, resultMap,
641                                                  resultOperands);
642     }
643     return success();
644   }
645 
646   return failure();
647 }
648