1 //===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//
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 utilities for the Linalg dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/Linalg/Utils/Utils.h"
14 
15 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
16 #include "mlir/Dialect/Affine/IR/AffineOps.h"
17 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
18 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
19 #include "mlir/Dialect/SCF/EDSC/Builders.h"
20 #include "mlir/Dialect/SCF/SCF.h"
21 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
22 #include "mlir/Dialect/StandardOps/IR/Ops.h"
23 #include "mlir/IR/AffineExpr.h"
24 #include "mlir/IR/AffineExprVisitor.h"
25 #include "mlir/IR/AffineMap.h"
26 #include "mlir/IR/Matchers.h"
27 #include "mlir/IR/OpImplementation.h"
28 #include "mlir/Pass/Pass.h"
29 #include "mlir/Transforms/LoopUtils.h"
30 #include "llvm/Support/Debug.h"
31 
32 #define DEBUG_TYPE "linalg-utils"
33 
34 using namespace mlir;
35 using namespace mlir::edsc;
36 using namespace mlir::edsc::intrinsics;
37 using namespace mlir::linalg;
38 using namespace mlir::scf;
39 
40 static bool isZero(Value v) {
41   if (auto cst = v.getDefiningOp<ConstantIndexOp>())
42     return cst.getValue() == 0;
43   return false;
44 }
45 
46 namespace {
47 
48 // Helper visitor to determine whether an AffineExpr is tiled.
49 // This is achieved by traversing every AffineDimExpr with position `pos` and
50 // checking whether the corresponding `tileSizes[pos]` is non-zero.
51 // This also enforces only positive coefficients occur in multiplications.
52 //
53 // Example:
54 //   `d0 + 2 * d1 + d3` is tiled by [0, 0, 0, 2] but not by [0, 0, 2, 0]
55 //
56 struct TileCheck : public AffineExprVisitor<TileCheck> {
57   TileCheck(ValueRange tileSizes) : isTiled(false), tileSizes(tileSizes) {}
58 
59   void visitDimExpr(AffineDimExpr expr) {
60     isTiled |= !isZero(tileSizes[expr.getPosition()]);
61   }
62   void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) {
63     visit(expr.getLHS());
64     visit(expr.getRHS());
65     if (expr.getKind() == mlir::AffineExprKind::Mul)
66       assert(expr.getRHS().cast<AffineConstantExpr>().getValue() > 0 &&
67              "nonpositive multiplying coefficient");
68   }
69   bool isTiled;
70   ValueRange tileSizes;
71 };
72 
73 } // namespace
74 
75 static bool isTiled(AffineExpr expr, ValueRange tileSizes) {
76   if (!expr)
77     return false;
78   TileCheck t(tileSizes);
79   t.visit(expr);
80   return t.isTiled;
81 }
82 
83 // Checks whether the `map  varies with respect to a non-zero `tileSize`.
84 static bool isTiled(AffineMap map, ValueRange tileSizes) {
85   if (!map)
86     return false;
87   for (unsigned r = 0; r < map.getNumResults(); ++r)
88     if (isTiled(map.getResult(r), tileSizes))
89       return true;
90   return false;
91 }
92 
93 Optional<RegionMatcher::BinaryOpKind>
94 RegionMatcher::matchAsScalarBinaryOp(GenericOp op) {
95   auto &region = op.region();
96   if (!llvm::hasSingleElement(region))
97     return llvm::None;
98 
99   Block &block = region.front();
100   if (block.getNumArguments() != 2 ||
101       !block.getArgument(0).getType().isSignlessIntOrFloat() ||
102       !block.getArgument(1).getType().isSignlessIntOrFloat())
103     return llvm::None;
104 
105   auto &ops = block.getOperations();
106   if (!llvm::hasSingleElement(block.without_terminator()))
107     return llvm::None;
108 
109   using mlir::matchers::m_Val;
110   auto a = m_Val(block.getArgument(0));
111   auto b = m_Val(block.getArgument(1));
112 
113   auto addPattern = m_Op<linalg::YieldOp>(m_Op<AddIOp>(a, b));
114   if (addPattern.match(&ops.back()))
115     return BinaryOpKind::IAdd;
116 
117   return llvm::None;
118 }
119 
120 bool mlir::linalg::isParallelIteratorType(Attribute attr) {
121   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
122     return strAttr.getValue() == getParallelIteratorTypeName();
123   }
124   return false;
125 }
126 
127 bool mlir::linalg::isReductionIteratorType(Attribute attr) {
128   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
129     return strAttr.getValue() == getReductionIteratorTypeName();
130   }
131   return false;
132 }
133 
134 bool mlir::linalg::isWindowIteratorType(Attribute attr) {
135   if (auto strAttr = attr.dyn_cast<StringAttr>()) {
136     return strAttr.getValue() == getWindowIteratorTypeName();
137   }
138   return false;
139 }
140 
141 /// Explicit instantiation of loop nest generator for different loop types.
142 template struct mlir::linalg::GenerateLoopNest<scf::ForOp>;
143 template struct mlir::linalg::GenerateLoopNest<scf::ParallelOp>;
144 template struct mlir::linalg::GenerateLoopNest<AffineForOp>;
145 template struct mlir::linalg::GenerateLoopNest<TiledLoopOp>;
146 
147 /// Given a list of subview ranges, extract individual values for lower, upper
148 /// bounds and steps and put them into the corresponding vectors.
149 static void unpackRanges(ArrayRef<Range> ranges, SmallVectorImpl<Value> &lbs,
150                          SmallVectorImpl<Value> &ubs,
151                          SmallVectorImpl<Value> &steps) {
152   for (Range range : ranges) {
153     lbs.emplace_back(range.offset);
154     ubs.emplace_back(range.size);
155     steps.emplace_back(range.stride);
156   }
157 }
158 
159 namespace mlir {
160 namespace linalg {
161 
162 /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp
163 /// is a constant then return a new value set to the smallest such constant.
164 /// Otherwise returngetSmallestBoundingIndex nullptr.
165 IntegerAttr getSmallestBoundingIndex(Value size) {
166   Optional<int64_t> boundingConst = {};
167   if (auto affineMinOp = size.getDefiningOp<AffineMinOp>()) {
168     for (auto e : affineMinOp.getAffineMap().getResults())
169       if (auto cst = e.dyn_cast<AffineConstantExpr>())
170         boundingConst = boundingConst
171                             ? std::min(boundingConst.getValue(), cst.getValue())
172                             : cst.getValue();
173   } else if (auto constIndexOp = size.getDefiningOp<ConstantOp>()) {
174     if (constIndexOp.getType().isa<IndexType>())
175       boundingConst = constIndexOp.value().cast<IntegerAttr>().getInt();
176   } else if (auto affineApplyOp = size.getDefiningOp<AffineApplyOp>()) {
177     if (auto cExpr = affineApplyOp.getAffineMap()
178                          .getResult(0)
179                          .dyn_cast<AffineConstantExpr>())
180       boundingConst = cExpr.getValue();
181   }
182   if (boundingConst && *boundingConst >= 0)
183     return Builder(size.getContext()).getIndexAttr(*boundingConst);
184   return nullptr;
185 }
186 
187 /// Specialization to build an scf "for" nest.
188 template <>
189 void GenerateLoopNest<scf::ForOp>::doit(
190     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
191     ArrayRef<Attribute> iteratorTypes,
192     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
193     Optional<LinalgLoopDistributionOptions> distributionOptions) {
194   auto iterArgInitValues = linalgOp.getOutputTensors();
195   // Create procInfo so it dominates loops, if appropriate.
196   OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
197   Location loc = edsc::ScopedContext::getLocation();
198   SmallVector<ProcInfo, 2> procInfo;
199   if (distributionOptions.hasValue())
200     procInfo = distributionOptions->procInfo(builder, loc, loopRanges);
201 
202   SmallVector<Value, 4> lbs, ubs, steps;
203   unpackRanges(loopRanges, lbs, ubs, steps);
204   LoopNest loopNest =
205       edsc::loopNestBuilder(lbs, ubs, steps, iterArgInitValues, bodyBuilderFn);
206 
207   if (!distributionOptions.hasValue() || loopNest.loops.empty())
208     return;
209 
210   // Only supports cyclic distribution for now.
211   for (auto it : llvm::zip(loopNest.loops, procInfo,
212                            distributionOptions->distributionMethod))
213     if (std::get<2>(it) == DistributionMethod::Cyclic)
214       mapLoopToProcessorIds(std::get<0>(it), std::get<1>(it).procId,
215                             std::get<1>(it).nprocs);
216 }
217 
218 /// Specialization to build affine "for" nest.
219 template <>
220 void GenerateLoopNest<AffineForOp>::doit(
221     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
222     ArrayRef<Attribute> iteratorTypes,
223     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
224     Optional<LinalgLoopDistributionOptions>) {
225   auto iterArgInitValues = linalgOp.getOutputTensors();
226   assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
227   SmallVector<Value, 4> lbs, ubs, steps;
228   unpackRanges(loopRanges, lbs, ubs, steps);
229 
230   // Affine loops require constant steps.
231   SmallVector<int64_t, 4> constantSteps;
232   constantSteps.reserve(steps.size());
233   for (Value v : steps) {
234     auto op = v.getDefiningOp<ConstantIndexOp>();
235     assert(op && "Affine loops require constant steps");
236     constantSteps.push_back(op.getValue());
237   }
238 
239   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
240     bodyBuilderFn(ivs, {});
241   };
242   edsc::affineLoopNestBuilder(lbs, ubs, constantSteps,
243                               bodyBuilderWithoutIterArgsFn);
244 }
245 
246 /// Specialization to build an linalg.tiled_loop
247 template <>
248 void GenerateLoopNest<TiledLoopOp>::doit(
249     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
250     ArrayRef<Attribute> iteratorTypes,
251     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
252     Optional<LinalgLoopDistributionOptions>) {
253   OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
254   Location loc = edsc::ScopedContext::getLocation();
255   SmallVector<ProcInfo, 2> procInfo;
256 
257   SmallVector<Value, 4> lbs, ubs, steps;
258   unpackRanges(loopRanges, lbs, ubs, steps);
259 
260   auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
261                               ValueRange ivs, ValueRange inputs,
262                               ValueRange outputs) {
263     ScopedContext context(nestedBuilder, nestedLoc);
264     scf::ValueVector results = bodyBuilderFn(ivs, linalgOp.getOutputTensors());
265     nestedBuilder.create<linalg::YieldOp>(nestedLoc, results);
266   };
267 
268   auto tiledLoop = builder.create<TiledLoopOp>(
269       loc, lbs, ubs, steps, linalgOp.getInputs(), linalgOp.getOutputs(),
270       builder.getArrayAttr(iteratorTypes), wrappedBuilderFn);
271 
272   // Replace inputs/outputs with the corresponding region args.
273   auto isInsideTiledLoop = [&](OpOperand &operand) {
274     return operand.getOwner()->getBlock() == tiledLoop.getBody();
275   };
276   for (auto it :
277        llvm::zip(linalgOp.getInputs(), tiledLoop.getRegionInputArgs()))
278     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop);
279   for (auto it :
280        llvm::zip(linalgOp.getOutputs(), tiledLoop.getRegionOutputArgs()))
281     std::get<0>(it).replaceUsesWithIf(std::get<1>(it), isInsideTiledLoop);
282 }
283 
284 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
285 void updateBoundsForCyclicDistribution(OpBuilder &builder, Location loc,
286                                        Value procId, Value nprocs, Value &lb,
287                                        Value &ub, Value &step) {
288   using edsc::op::operator+;
289   using edsc::op::operator*;
290   lb = lb + (procId * step);
291   step = nprocs * step;
292 }
293 
294 /// Generates a loop nest consisting of scf.parallel and scf.for, depending
295 /// on the `iteratorTypes.` Consecutive parallel loops create a single
296 /// scf.parallel operation; each sequential loop creates a new scf.for
297 /// operation. The body of the innermost loop is populated by
298 /// `bodyBuilderFn` that accepts a range of induction variables for all
299 /// loops. `ivStorage` is used to store the partial list of induction
300 /// variables.
301 // TODO: this function can be made iterative instead. However, it
302 // will have at most as many recursive calls as nested loops, which rarely
303 // exceeds 10.
304 static void
305 generateParallelLoopNest(ValueRange lbs, ValueRange ubs, ValueRange steps,
306                          ArrayRef<Attribute> iteratorTypes,
307                          function_ref<void(ValueRange)> bodyBuilderFn,
308                          SmallVectorImpl<Value> &ivStorage,
309                          ArrayRef<DistributionMethod> distributionMethod = {}) {
310   assert(lbs.size() == ubs.size());
311   assert(lbs.size() == steps.size());
312   assert(lbs.size() == iteratorTypes.size());
313 
314   // If there are no (more) loops to be generated, generate the body and be
315   // done with it.
316   if (iteratorTypes.empty())
317     return bodyBuilderFn(ivStorage);
318 
319   // Find the outermost parallel loops and drop their types from the list.
320   unsigned nLoops = iteratorTypes.size();
321   unsigned nOuterPar =
322       nLoops - iteratorTypes.drop_while(isParallelIteratorType).size();
323 
324   // If there are no outer parallel loops, generate one sequential loop and
325   // recurse. Note that we wouldn't have dropped anything from `iteratorTypes`
326   // in this case.
327   if (nOuterPar == 0) {
328     edsc::loopNestBuilder(lbs[0], ubs[0], steps[0], [&](Value iv) {
329       ivStorage.push_back(iv);
330       generateParallelLoopNest(lbs.drop_front(), ubs.drop_front(),
331                                steps.drop_front(), iteratorTypes.drop_front(),
332                                bodyBuilderFn, ivStorage, distributionMethod);
333     });
334     return;
335   }
336   if (distributionMethod.empty()) {
337     // Generate a single parallel loop-nest operation for all outermost
338     // parallel loops and recurse.
339     edsc::OperationBuilder<scf::ParallelOp>(
340         lbs.take_front(nOuterPar), ubs.take_front(nOuterPar),
341         steps.take_front(nOuterPar),
342         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
343           edsc::ScopedContext context(nestedBuilder, nestedLoc);
344           ivStorage.append(localIvs.begin(), localIvs.end());
345           generateParallelLoopNest(
346               lbs.drop_front(nOuterPar), ubs.drop_front(nOuterPar),
347               steps.drop_front(nOuterPar), iteratorTypes.drop_front(nOuterPar),
348               bodyBuilderFn, ivStorage,
349               (distributionMethod.size() < nOuterPar)
350                   ? ArrayRef<DistributionMethod>()
351                   : distributionMethod.drop_front(nOuterPar));
352         });
353     return;
354   }
355 
356   // Process all consecutive similarly distributed loops simultaneously.
357   DistributionMethod methodToUse = distributionMethod[0];
358   unsigned numProcessed = 1;
359   for (unsigned i = 1; i < nOuterPar && i < distributionMethod.size(); ++i) {
360     if (distributionMethod[i] != methodToUse)
361       break;
362     numProcessed++;
363   }
364 
365   switch (methodToUse) {
366   case DistributionMethod::Cyclic: {
367     // Generate a single parallel loop-nest operation for all outermost
368     // parallel loops and recurse.
369     edsc::OperationBuilder<scf::ParallelOp>(
370         lbs.take_front(numProcessed), ubs.take_front(numProcessed),
371         steps.take_front(numProcessed),
372         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
373           edsc::ScopedContext context(nestedBuilder, nestedLoc);
374           ivStorage.append(localIvs.begin(), localIvs.end());
375           generateParallelLoopNest(
376               lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
377               steps.drop_front(numProcessed),
378               iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
379               (distributionMethod.size() < numProcessed)
380                   ? ArrayRef<DistributionMethod>()
381                   : distributionMethod.drop_front(numProcessed));
382         });
383     return;
384   }
385   case DistributionMethod::CyclicNumProcsGeNumIters: {
386     // Check (for the processed loops) that the iteration is in-bounds.
387     using edsc::op::slt;
388     using edsc::op::operator&&;
389     Value cond = slt(lbs[0], ubs[0]);
390     for (unsigned i = 1; i < numProcessed; ++i)
391       cond = cond && slt(lbs[i], ubs[i]);
392     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
393     edsc::conditionBuilder(cond, [&]() {
394       generateParallelLoopNest(
395           lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
396           steps.drop_front(numProcessed),
397           iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
398           distributionMethod.drop_front(numProcessed));
399     });
400     return;
401   }
402   case DistributionMethod::CyclicNumProcsEqNumIters:
403     // No check/loops needed here. Set the `%iv` to be the `%lb` and proceed
404     // with inner loop generation.
405     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
406     generateParallelLoopNest(
407         lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
408         steps.drop_front(numProcessed), iteratorTypes.drop_front(numProcessed),
409         bodyBuilderFn, ivStorage, distributionMethod.drop_front(numProcessed));
410     return;
411   }
412 }
413 
414 /// Specialization for generating a mix of parallel and sequential scf loops.
415 template <>
416 void GenerateLoopNest<scf::ParallelOp>::doit(
417     ArrayRef<Range> loopRanges, LinalgOp linalgOp,
418     ArrayRef<Attribute> iteratorTypes,
419     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
420     Optional<LinalgLoopDistributionOptions> distributionOptions) {
421   auto iterArgInitValues = linalgOp.getOutputTensors();
422   assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");
423   // This function may be passed more iterator types than ranges.
424   assert(iteratorTypes.size() >= loopRanges.size() &&
425          "expected iterator type for all ranges");
426   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
427   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
428   unsigned numLoops = iteratorTypes.size();
429   ivs.reserve(numLoops);
430   lbsStorage.reserve(numLoops);
431   ubsStorage.reserve(numLoops);
432   stepsStorage.reserve(numLoops);
433 
434   // Get the loop lb, ub, and step.
435   unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
436 
437   // Modify the lb, ub, and step based on the distribution options.
438   SmallVector<DistributionMethod, 0> distributionMethod;
439   if (distributionOptions) {
440     auto &options = distributionOptions.getValue();
441     OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
442     Location loc = edsc::ScopedContext::getLocation();
443     distributionMethod.assign(distributionOptions->distributionMethod.begin(),
444                               distributionOptions->distributionMethod.end());
445     SmallVector<Range, 2> parallelLoopRanges;
446     for (auto iteratorType : enumerate(iteratorTypes)) {
447       if (isParallelIteratorType(iteratorType.value()))
448         parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
449     }
450     if (distributionMethod.size() < parallelLoopRanges.size())
451       parallelLoopRanges.resize(distributionMethod.size());
452     SmallVector<ProcInfo, 2> procInfo =
453         options.procInfo(builder, loc, parallelLoopRanges);
454     unsigned index = 0;
455     for (auto iteratorType : enumerate(iteratorTypes)) {
456       if (index >= procInfo.size())
457         break;
458       if (isParallelIteratorType(iteratorType.value())) {
459         unsigned i = iteratorType.index();
460         updateBoundsForCyclicDistribution(builder, loc, procInfo[index].procId,
461                                           procInfo[index].nprocs, lbsStorage[i],
462                                           ubsStorage[i], stepsStorage[i]);
463         index++;
464       }
465     }
466   }
467   ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
468   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
469     bodyBuilderFn(ivs, {});
470   };
471   generateParallelLoopNest(lbs, ubs, steps, iteratorTypes,
472                            bodyBuilderWithoutIterArgsFn, ivs,
473                            distributionMethod);
474 
475   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
476 }
477 
478 SmallVector<Value, 4> makeTiledShapes(OpBuilder &builder, Location loc,
479                                       LinalgOp linalgOp,
480                                       ArrayRef<Value> tiledOperands,
481                                       ValueRange ivs, ValueRange tileSizes,
482                                       ArrayRef<Value> sizeBounds) {
483   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
484                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
485                            [](Value v) { return !isZero(v); })) &&
486          "expected as many ivs as non-zero sizes");
487 
488   using namespace edsc::op;
489 
490   // Construct (potentially temporary) mins and maxes on which to apply maps
491   // that define tile subshapes.
492   SmallVector<Value, 8> lbs, subShapeSizes;
493   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
494     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n");
495     bool isTiled = !isZero(tileSizes[idx]);
496     lbs.push_back(isTiled ? ivs[idxIvs++] : (Value)std_constant_index(0));
497     // Before composing, we need to make range a closed interval.
498     Value size = isTiled ? tileSizes[idx] : sizeBounds[idx];
499     subShapeSizes.push_back(size - std_constant_index(1));
500     LLVM_DEBUG(llvm::dbgs() << "lb: " << lbs.back() << "\n");
501     LLVM_DEBUG(llvm::dbgs() << "size: " << subShapeSizes.back() << "\n");
502   }
503 
504   MLIRContext *context = builder.getContext();
505   SmallVector<Value, 4> tiledShapes;
506   tiledShapes.reserve(tiledOperands.size());
507   for (auto en : llvm::enumerate(tiledOperands)) {
508     Value shapedOp = en.value();
509     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
510     ShapedType shapedType = shapedOp.getType().cast<ShapedType>();
511     unsigned rank = shapedType.getRank();
512     AffineMap map = linalgOp.getIndexingMap(en.index());
513     // If the shape is not tiled, we can use it as is.
514     if (!isTiled(map, tileSizes)) {
515       tiledShapes.push_back(shapedOp);
516       LLVM_DEBUG(llvm::dbgs()
517                  << ": not tiled: use shape: " << shapedType << "\n");
518       continue;
519     }
520     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subshape...\n");
521 
522     // Construct a new subview / subtensor for the tile.
523     SmallVector<OpFoldResult, 4> offsets, sizes, strides;
524     offsets.reserve(rank);
525     sizes.reserve(rank);
526     strides.reserve(rank);
527     for (unsigned r = 0; r < rank; ++r) {
528       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for dim#" << r);
529       if (!isTiled(map.getSubMap({r}), tileSizes)) {
530         offsets.push_back(builder.getIndexAttr(0));
531         Value dim = memref_dim(shapedOp, r).value;
532         sizes.push_back(dim);
533         strides.push_back(builder.getIndexAttr(1));
534         LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n");
535         continue;
536       }
537       LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");
538 
539       // Tiling creates a new slice at the proper index, the slice step is 1
540       // (i.e. the op does not subsample, stepping occurs in the loop).
541       auto m = map.getSubMap({r});
542       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: submap: " << map << "\n");
543       auto offset = applyMapToValues(builder, loc, m, lbs).front();
544       offsets.push_back(offset);
545       auto closedIntSize =
546           applyMapToValues(builder, loc, m, subShapeSizes).front();
547       // Resulting size needs to be made half open interval again.
548       auto size = closedIntSize + std_constant_index(1);
549       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: raw size: " << size << "\n");
550 
551       // The size of the subview / subtensor should be trimmed to avoid
552       // out-of-bounds accesses, unless we statically know the subshape size
553       // divides the shape size evenly.
554       int64_t shapeSize = shapedType.getDimSize(r);
555       auto sizeCst = size.getDefiningOp<ConstantIndexOp>();
556       if (ShapedType::isDynamic(shapeSize) || !sizeCst ||
557           (shapeSize % sizeCst.getValue()) != 0) {
558         LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: shapeSize=" << shapeSize
559                                 << ", size: " << size
560                                 << ": make sure in bound with affine.min\n");
561         AffineExpr dim0, dim1, dim2;
562         bindDims(context, dim0, dim1, dim2);
563         // Compute min(size, dim - offset) to avoid out-of-bounds accesses.
564         auto minMap = AffineMap::get(
565             /*dimCount=*/3, /*symbolCount=*/0, {dim0, dim1 - dim2}, context);
566         Value d = memref_dim(shapedOp, r);
567         SmallVector<Value, 4> operands{size, d, offset};
568         fullyComposeAffineMapAndOperands(&minMap, &operands);
569         size = affine_min(builder.getIndexType(), minMap, operands);
570       }
571 
572       sizes.push_back(size);
573       LLVM_DEBUG(llvm::dbgs()
574                  << "makeTiledShapes: new offset: " << offset << "\n");
575       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: new size: " << size << "\n");
576       strides.push_back(builder.getIndexAttr(1));
577     }
578 
579     if (shapedType.isa<MemRefType>())
580       tiledShapes.push_back(builder.create<memref::SubViewOp>(
581           loc, shapedOp, offsets, sizes, strides));
582     else
583       tiledShapes.push_back(
584           builder.create<SubTensorOp>(loc, shapedOp, offsets, sizes, strides));
585   }
586 
587   return tiledShapes;
588 }
589 
590 } // namespace linalg
591 } // namespace mlir
592