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