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 
146 /// Given a list of subview ranges, extract individual values for lower, upper
147 /// bounds and steps and put them into the corresponding vectors.
148 static void unpackRanges(ArrayRef<Range> ranges, SmallVectorImpl<Value> &lbs,
149                          SmallVectorImpl<Value> &ubs,
150                          SmallVectorImpl<Value> &steps) {
151   for (Range range : ranges) {
152     lbs.emplace_back(range.offset);
153     ubs.emplace_back(range.size);
154     steps.emplace_back(range.stride);
155   }
156 }
157 
158 namespace mlir {
159 namespace linalg {
160 
161 /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp
162 /// is a constant then return a new value set to the smallest such constant.
163 /// Otherwise returngetSmallestBoundingIndex nullptr.
164 IntegerAttr getSmallestBoundingIndex(Value size) {
165   Optional<int64_t> boundingConst = {};
166   if (auto affineMinOp = size.getDefiningOp<AffineMinOp>()) {
167     for (auto e : affineMinOp.getAffineMap().getResults())
168       if (auto cst = e.dyn_cast<AffineConstantExpr>())
169         boundingConst = boundingConst
170                             ? std::min(boundingConst.getValue(), cst.getValue())
171                             : cst.getValue();
172   } else if (auto constIndexOp = size.getDefiningOp<ConstantOp>()) {
173     if (constIndexOp.getType().isa<IndexType>())
174       boundingConst = constIndexOp.value().cast<IntegerAttr>().getInt();
175   } else if (auto affineApplyOp = size.getDefiningOp<AffineApplyOp>()) {
176     if (auto cExpr = affineApplyOp.getAffineMap()
177                          .getResult(0)
178                          .dyn_cast<AffineConstantExpr>())
179       boundingConst = cExpr.getValue();
180   }
181   if (boundingConst && *boundingConst >= 0)
182     return Builder(size.getContext()).getIndexAttr(*boundingConst);
183   return nullptr;
184 }
185 
186 /// Specialization to build an scf "for" nest.
187 template <>
188 void GenerateLoopNest<scf::ForOp>::doit(
189     ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
190     ArrayRef<Attribute> iteratorTypes,
191     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
192     Optional<LinalgLoopDistributionOptions> distributionOptions) {
193   // Create procInfo so it dominates loops, if appropriate.
194   OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
195   Location loc = edsc::ScopedContext::getLocation();
196   SmallVector<ProcInfo, 2> procInfo;
197   if (distributionOptions.hasValue())
198     procInfo = distributionOptions->procInfo(builder, loc, loopRanges);
199 
200   SmallVector<Value, 4> lbs, ubs, steps;
201   unpackRanges(loopRanges, lbs, ubs, steps);
202   LoopNest loopNest =
203       edsc::loopNestBuilder(lbs, ubs, steps, iterArgInitValues, bodyBuilderFn);
204 
205   if (!distributionOptions.hasValue() || loopNest.loops.empty())
206     return;
207 
208   // Only supports cyclic distribution for now.
209   for (auto it : llvm::zip(loopNest.loops, procInfo,
210                            distributionOptions->distributionMethod))
211     if (std::get<2>(it) == DistributionMethod::Cyclic)
212       mapLoopToProcessorIds(std::get<0>(it), std::get<1>(it).procId,
213                             std::get<1>(it).nprocs);
214 }
215 
216 /// Specialization to build affine "for" nest.
217 template <>
218 void GenerateLoopNest<AffineForOp>::doit(
219     ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
220     ArrayRef<Attribute> iteratorTypes,
221     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
222     Optional<LinalgLoopDistributionOptions>) {
223   assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
224   SmallVector<Value, 4> lbs, ubs, steps;
225   unpackRanges(loopRanges, lbs, ubs, steps);
226 
227   // Affine loops require constant steps.
228   SmallVector<int64_t, 4> constantSteps;
229   constantSteps.reserve(steps.size());
230   for (Value v : steps) {
231     auto op = v.getDefiningOp<ConstantIndexOp>();
232     assert(op && "Affine loops require constant steps");
233     constantSteps.push_back(op.getValue());
234   }
235 
236   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
237     bodyBuilderFn(ivs, {});
238   };
239   edsc::affineLoopNestBuilder(lbs, ubs, constantSteps,
240                               bodyBuilderWithoutIterArgsFn);
241 }
242 
243 /// Update the `lb`, `ub` and `step` to get per processor `lb`, `ub` and `step`.
244 static void updateBoundsForCyclicDistribution(OpBuilder &builder, Location loc,
245                                               Value procId, Value nprocs,
246                                               Value &lb, Value &ub,
247                                               Value &step) {
248   using edsc::op::operator+;
249   using edsc::op::operator*;
250   lb = lb + (procId * step);
251   step = nprocs * step;
252 }
253 
254 /// Generates a loop nest consisting of scf.parallel and scf.for, depending
255 /// on the `iteratorTypes.` Consecutive parallel loops create a single
256 /// scf.parallel operation; each sequential loop creates a new scf.for
257 /// operation. The body of the innermost loop is populated by
258 /// `bodyBuilderFn` that accepts a range of induction variables for all
259 /// loops. `ivStorage` is used to store the partial list of induction
260 /// variables.
261 // TODO: this function can be made iterative instead. However, it
262 // will have at most as many recursive calls as nested loops, which rarely
263 // exceeds 10.
264 static void
265 generateParallelLoopNest(ValueRange lbs, ValueRange ubs, ValueRange steps,
266                          ArrayRef<Attribute> iteratorTypes,
267                          function_ref<void(ValueRange)> bodyBuilderFn,
268                          SmallVectorImpl<Value> &ivStorage,
269                          ArrayRef<DistributionMethod> distributionMethod = {}) {
270   assert(lbs.size() == ubs.size());
271   assert(lbs.size() == steps.size());
272   assert(lbs.size() == iteratorTypes.size());
273 
274   // If there are no (more) loops to be generated, generate the body and be
275   // done with it.
276   if (iteratorTypes.empty())
277     return bodyBuilderFn(ivStorage);
278 
279   // Find the outermost parallel loops and drop their types from the list.
280   unsigned nLoops = iteratorTypes.size();
281   unsigned nOuterPar =
282       nLoops - iteratorTypes.drop_while(isParallelIteratorType).size();
283 
284   // If there are no outer parallel loops, generate one sequential loop and
285   // recurse. Note that we wouldn't have dropped anything from `iteratorTypes`
286   // in this case.
287   if (nOuterPar == 0) {
288     edsc::loopNestBuilder(lbs[0], ubs[0], steps[0], [&](Value iv) {
289       ivStorage.push_back(iv);
290       generateParallelLoopNest(lbs.drop_front(), ubs.drop_front(),
291                                steps.drop_front(), iteratorTypes.drop_front(),
292                                bodyBuilderFn, ivStorage, distributionMethod);
293     });
294     return;
295   }
296   if (distributionMethod.empty()) {
297     // Generate a single parallel loop-nest operation for all outermost
298     // parallel loops and recurse.
299     edsc::OperationBuilder<scf::ParallelOp>(
300         lbs.take_front(nOuterPar), ubs.take_front(nOuterPar),
301         steps.take_front(nOuterPar),
302         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
303           edsc::ScopedContext context(nestedBuilder, nestedLoc);
304           ivStorage.append(localIvs.begin(), localIvs.end());
305           generateParallelLoopNest(
306               lbs.drop_front(nOuterPar), ubs.drop_front(nOuterPar),
307               steps.drop_front(nOuterPar), iteratorTypes.drop_front(nOuterPar),
308               bodyBuilderFn, ivStorage,
309               (distributionMethod.size() < nOuterPar)
310                   ? ArrayRef<DistributionMethod>()
311                   : distributionMethod.drop_front(nOuterPar));
312         });
313     return;
314   }
315 
316   // Process all consecutive similarly distributed loops simultaneously.
317   DistributionMethod methodToUse = distributionMethod[0];
318   unsigned numProcessed = 1;
319   for (unsigned i = 1; i < nOuterPar && i < distributionMethod.size(); ++i) {
320     if (distributionMethod[i] != methodToUse)
321       break;
322     numProcessed++;
323   }
324 
325   switch (methodToUse) {
326   case DistributionMethod::Cyclic: {
327     // Generate a single parallel loop-nest operation for all outermost
328     // parallel loops and recurse.
329     edsc::OperationBuilder<scf::ParallelOp>(
330         lbs.take_front(numProcessed), ubs.take_front(numProcessed),
331         steps.take_front(numProcessed),
332         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange localIvs) {
333           edsc::ScopedContext context(nestedBuilder, nestedLoc);
334           ivStorage.append(localIvs.begin(), localIvs.end());
335           generateParallelLoopNest(
336               lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
337               steps.drop_front(numProcessed),
338               iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
339               (distributionMethod.size() < numProcessed)
340                   ? ArrayRef<DistributionMethod>()
341                   : distributionMethod.drop_front(numProcessed));
342         });
343     return;
344   }
345   case DistributionMethod::CyclicNumProcsGeNumIters: {
346     // Check (for the processed loops) that the iteration is in-bounds.
347     using edsc::op::slt;
348     using edsc::op::operator&&;
349     Value cond = slt(lbs[0], ubs[0]);
350     for (unsigned i = 1; i < numProcessed; ++i)
351       cond = cond && slt(lbs[i], ubs[i]);
352     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
353     edsc::conditionBuilder(cond, [&]() {
354       generateParallelLoopNest(
355           lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
356           steps.drop_front(numProcessed),
357           iteratorTypes.drop_front(numProcessed), bodyBuilderFn, ivStorage,
358           distributionMethod.drop_front(numProcessed));
359     });
360     return;
361   }
362   case DistributionMethod::CyclicNumProcsEqNumIters:
363     // No check/loops needed here. Set the `%iv` to be the `%lb` and proceed
364     // with inner loop generation.
365     ivStorage.append(lbs.begin(), std::next(lbs.begin(), numProcessed));
366     generateParallelLoopNest(
367         lbs.drop_front(numProcessed), ubs.drop_front(numProcessed),
368         steps.drop_front(numProcessed), iteratorTypes.drop_front(numProcessed),
369         bodyBuilderFn, ivStorage, distributionMethod.drop_front(numProcessed));
370     return;
371   }
372 }
373 
374 /// Specialization for generating a mix of parallel and sequential scf loops.
375 template <>
376 void GenerateLoopNest<scf::ParallelOp>::doit(
377     ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
378     ArrayRef<Attribute> iteratorTypes,
379     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
380     Optional<LinalgLoopDistributionOptions> distributionOptions) {
381   assert(iterArgInitValues.empty() && "unexpected ParallelOp init values");
382   // This function may be passed more iterator types than ranges.
383   assert(iteratorTypes.size() >= loopRanges.size() &&
384          "expected iterator type for all ranges");
385   iteratorTypes = iteratorTypes.take_front(loopRanges.size());
386   SmallVector<Value, 8> lbsStorage, ubsStorage, stepsStorage, ivs;
387   unsigned numLoops = iteratorTypes.size();
388   ivs.reserve(numLoops);
389   lbsStorage.reserve(numLoops);
390   ubsStorage.reserve(numLoops);
391   stepsStorage.reserve(numLoops);
392 
393   // Get the loop lb, ub, and step.
394   unpackRanges(loopRanges, lbsStorage, ubsStorage, stepsStorage);
395 
396   // Modify the lb, ub, and step based on the distribution options.
397   SmallVector<DistributionMethod, 0> distributionMethod;
398   if (distributionOptions) {
399     auto &options = distributionOptions.getValue();
400     OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
401     Location loc = edsc::ScopedContext::getLocation();
402     distributionMethod.assign(distributionOptions->distributionMethod.begin(),
403                               distributionOptions->distributionMethod.end());
404     SmallVector<Range, 2> parallelLoopRanges;
405     for (auto iteratorType : enumerate(iteratorTypes)) {
406       if (isParallelIteratorType(iteratorType.value()))
407         parallelLoopRanges.push_back(loopRanges[iteratorType.index()]);
408     }
409     if (distributionMethod.size() < parallelLoopRanges.size())
410       parallelLoopRanges.resize(distributionMethod.size());
411     SmallVector<ProcInfo, 2> procInfo =
412         options.procInfo(builder, loc, parallelLoopRanges);
413     unsigned index = 0;
414     for (auto iteratorType : enumerate(iteratorTypes)) {
415       if (index >= procInfo.size())
416         break;
417       if (isParallelIteratorType(iteratorType.value())) {
418         unsigned i = iteratorType.index();
419         updateBoundsForCyclicDistribution(builder, loc, procInfo[index].procId,
420                                           procInfo[index].nprocs, lbsStorage[i],
421                                           ubsStorage[i], stepsStorage[i]);
422         index++;
423       }
424     }
425   }
426   ValueRange lbs(lbsStorage), ubs(ubsStorage), steps(stepsStorage);
427   auto bodyBuilderWithoutIterArgsFn = [&](ValueRange ivs) {
428     bodyBuilderFn(ivs, {});
429   };
430   generateParallelLoopNest(lbs, ubs, steps, iteratorTypes,
431                            bodyBuilderWithoutIterArgsFn, ivs,
432                            distributionMethod);
433 
434   assert(ivs.size() == iteratorTypes.size() && "did not generate enough loops");
435 }
436 
437 SmallVector<Value, 4> makeTiledShapes(OpBuilder &builder, Location loc,
438                                       LinalgOp linalgOp,
439                                       ArrayRef<Value> tiledOperands,
440                                       ValueRange ivs, ValueRange tileSizes,
441                                       ArrayRef<Value> sizeBounds) {
442   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
443                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
444                            [](Value v) { return !isZero(v); })) &&
445          "expected as many ivs as non-zero sizes");
446 
447   using namespace edsc::op;
448 
449   // Construct (potentially temporary) mins and maxes on which to apply maps
450   // that define tile subshapes.
451   SmallVector<Value, 8> lbs, subShapeSizes;
452   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
453     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for loop#" << idx << "\n");
454     bool isTiled = !isZero(tileSizes[idx]);
455     lbs.push_back(isTiled ? ivs[idxIvs++] : (Value)std_constant_index(0));
456     // Before composing, we need to make range a closed interval.
457     Value size = isTiled ? tileSizes[idx] : sizeBounds[idx];
458     subShapeSizes.push_back(size - std_constant_index(1));
459     LLVM_DEBUG(llvm::dbgs() << "lb: " << lbs.back() << "\n");
460     LLVM_DEBUG(llvm::dbgs() << "size: " << subShapeSizes.back() << "\n");
461   }
462 
463   MLIRContext *context = builder.getContext();
464   SmallVector<Value, 4> tiledShapes;
465   tiledShapes.reserve(tiledOperands.size());
466   for (auto en : llvm::enumerate(tiledOperands)) {
467     Value shapedOp = en.value();
468     LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for operand " << shapedOp);
469     ShapedType shapedType = shapedOp.getType().cast<ShapedType>();
470     unsigned rank = shapedType.getRank();
471     AffineMap map = linalgOp.getIndexingMap(en.index());
472     // If the shape is not tiled, we can use it as is.
473     if (!isTiled(map, tileSizes)) {
474       tiledShapes.push_back(shapedOp);
475       LLVM_DEBUG(llvm::dbgs()
476                  << ": not tiled: use shape: " << shapedType << "\n");
477       continue;
478     }
479     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subshape...\n");
480 
481     // Construct a new subview / subtensor for the tile.
482     SmallVector<OpFoldResult, 4> offsets, sizes, strides;
483     offsets.reserve(rank);
484     sizes.reserve(rank);
485     strides.reserve(rank);
486     for (unsigned r = 0; r < rank; ++r) {
487       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: for dim#" << r);
488       if (!isTiled(map.getSubMap({r}), tileSizes)) {
489         offsets.push_back(builder.getIndexAttr(0));
490         Value dim = memref_dim(shapedOp, r).value;
491         sizes.push_back(dim);
492         strides.push_back(builder.getIndexAttr(1));
493         LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n");
494         continue;
495       }
496       LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");
497 
498       // Tiling creates a new slice at the proper index, the slice step is 1
499       // (i.e. the op does not subsample, stepping occurs in the loop).
500       auto m = map.getSubMap({r});
501       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: submap: " << map << "\n");
502       auto offset = applyMapToValues(builder, loc, m, lbs).front();
503       offsets.push_back(offset);
504       auto closedIntSize =
505           applyMapToValues(builder, loc, m, subShapeSizes).front();
506       // Resulting size needs to be made half open interval again.
507       auto size = closedIntSize + std_constant_index(1);
508       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: raw size: " << size << "\n");
509 
510       // The size of the subview / subtensor should be trimmed to avoid
511       // out-of-bounds accesses, unless we statically know the subshape size
512       // divides the shape size evenly.
513       int64_t shapeSize = shapedType.getDimSize(r);
514       auto sizeCst = size.getDefiningOp<ConstantIndexOp>();
515       if (ShapedType::isDynamic(shapeSize) || !sizeCst ||
516           (shapeSize % sizeCst.getValue()) != 0) {
517         LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: shapeSize=" << shapeSize
518                                 << ", size: " << size
519                                 << ": make sure in bound with affine.min\n");
520         AffineExpr dim0, dim1, dim2;
521         bindDims(context, dim0, dim1, dim2);
522         // Compute min(size, dim - offset) to avoid out-of-bounds accesses.
523         auto minMap = AffineMap::get(
524             /*dimCount=*/3, /*symbolCount=*/0, {dim0, dim1 - dim2}, context);
525         Value d = memref_dim(shapedOp, r);
526         SmallVector<Value, 4> operands{size, d, offset};
527         fullyComposeAffineMapAndOperands(&minMap, &operands);
528         size = affine_min(builder.getIndexType(), minMap, operands);
529       }
530 
531       sizes.push_back(size);
532       LLVM_DEBUG(llvm::dbgs()
533                  << "makeTiledShapes: new offset: " << offset << "\n");
534       LLVM_DEBUG(llvm::dbgs() << "makeTiledShapes: new size: " << size << "\n");
535       strides.push_back(builder.getIndexAttr(1));
536     }
537 
538     if (shapedType.isa<MemRefType>())
539       tiledShapes.push_back(builder.create<memref::SubViewOp>(
540           loc, shapedOp, offsets, sizes, strides));
541     else
542       tiledShapes.push_back(
543           builder.create<SubTensorOp>(loc, shapedOp, offsets, sizes, strides));
544   }
545 
546   return tiledShapes;
547 }
548 
549 } // namespace linalg
550 } // namespace mlir
551